A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Performance] Remove redundant clone() calls in cutlass_mla (#24891)
Alexander Matveev committed
Sep 15, 2025 at 16:21 UTC
aae725af7cee79628e576e7b394811050f651e49
1 file changed
+8
-8
vllm/v1/attention/backends/mla/cutlass_mla.py
+8
-8
@@ -210,9 +210,14 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
210
sm_scale,
211
num_kv_splits,
212
)
213
- returned_lse = lse[:, :H].contiguous(
214
- ) if self.need_to_return_lse_for_decode else lse
215
- return out[:, :H].contiguous(), returned_lse
213
+
214
+ if H < MAX_HEADS:
215
+ # Extract the subsets of the outputs
216
+ returned_lse = lse[:, :H].contiguous(
217
+ ) if self.need_to_return_lse_for_decode else lse
218
+ out = out[:, :H]
219
+
220
+ return out, returned_lse
221
222
def _sm100_forward_decode(
223
self,
@@ -228,11 +233,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
233
self._workspace.ensure_size(attn_metadata, self._num_kv_splits)
234
235
# Run MLA
231
- # Clone q_nope and q_pe to make sure strides computation is correct.
232
- # TODO: Check if we really need it
233
- q_nope = q_nope.clone()
234
- q_pe = q_pe.clone()
235
-
236
o, lse = self._sm100_cutlass_mla_decode(
237
q_nope,
238
q_pe,