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 8, 22:45 UTC
[Bugfix] Fix sparse MLA metadata building (#33579)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Feb 3, 2026 at 18:29 UTC
bd8da29a66ea8c0e0f208cab7ba0b6be640f3faa
1 file changed
+22
-31
vllm/model_executor/layers/attention/mla_attention.py
+22
-31
@@ -522,22 +522,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
522
k_c_normed = k_c_normed[:num_actual_toks, ...]
523
k_pe = k_pe[:num_actual_toks, ...]
524
525
- assert (
526
- attn_metadata.num_decodes is not None
527
- and attn_metadata.num_prefills is not None
528
- and attn_metadata.num_decode_tokens is not None
529
- )
530
-
531
- has_decode = attn_metadata.num_decodes > 0
532
- has_prefill = attn_metadata.num_prefills > 0
533
- num_decode_tokens = attn_metadata.num_decode_tokens
534
-
535
- decode_q = q[:num_decode_tokens]
536
-
537
- prefill_q = q[num_decode_tokens:]
538
- prefill_k_pe = k_pe[num_decode_tokens:]
539
- prefill_k_c_normed = k_c_normed[num_decode_tokens:]
540
-
525
# write the latent and rope to kv cache
526
if kv_cache.numel() > 0:
527
ops.concat_and_cache_mla(
@@ -555,27 +539,32 @@ class MLAAttention(nn.Module, AttentionLayerBase):
539
# Sparse MLA impls only support forward_mqa (decode-style attention)
540
is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl)
541
558
- if has_prefill and not is_sparse_impl:
542
+ if is_sparse_impl:
543
+ num_mqa_tokens = q.size(0)
544
+ num_mha_tokens = 0
545
+ else:
546
+ assert (
547
+ attn_metadata.num_decodes is not None
548
+ and attn_metadata.num_prefills is not None
549
+ and attn_metadata.num_decode_tokens is not None
550
+ )
551
+ num_mqa_tokens = attn_metadata.num_decode_tokens
552
+ num_mha_tokens = q.size(0) - num_mqa_tokens
553
+
554
+ if num_mha_tokens > 0:
555
self.impl.forward_mha(
560
- prefill_q,
561
- prefill_k_c_normed,
562
- prefill_k_pe,
556
+ q[num_mqa_tokens:],
557
+ k_c_normed[num_mqa_tokens:],
558
+ k_pe[num_mqa_tokens:],
559
kv_cache,
560
attn_metadata,
561
self._k_scale,
566
- output=output[num_decode_tokens:],
562
+ output=output[num_mqa_tokens:],
563
)
564
569
- if has_decode or (has_prefill and is_sparse_impl):
570
- # For sparse impl, we always use forward_mqa for all tokens
571
- # For non-sparse impl, we only use forward_mqa for decode tokens
572
- if is_sparse_impl:
573
- mqa_q = q
574
- mqa_output_slice = output
575
- else:
576
- assert attn_metadata.decode is not None
577
- mqa_q = decode_q
578
- mqa_output_slice = output[:num_decode_tokens]
565
+ if num_mqa_tokens > 0:
566
+ mqa_q = q[:num_mqa_tokens]
567
+ mqa_output_slice = output[:num_mqa_tokens]
568
569
mqa_q_nope, mqa_q_pe = mqa_q.split(
570
[self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
@@ -644,6 +633,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
633
mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
634
635
# call decode attn
636
+ if not is_sparse_impl:
637
+ assert attn_metadata.decode is not None
638
attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
639
640
# correct dcp attn_out with lse.