A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix FlashMLA sparse accuracy with topk_length and zero-init padding (#36616)
Signed-off-by: AjAnubolu <anuboluajay@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Ajay Anubolu committed
Jun 16, 2026 at 07:09 UTC
bf5149b51606dc567503684905705cc26427f099
1 file changed
+13
-3
vllm/v1/attention/backends/mla/flashmla_sparse.py
+13
-3
@@ -612,18 +612,20 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
612
) -> torch.Tensor:
613
# Convert per-request indices to global slots (decode) or workspace
614
# offsets (prefill).
615
- topk_indices = triton_convert_req_index_to_global_index(
615
+ topk_indices, topk_length = triton_convert_req_index_to_global_index(
616
attn_metadata.req_id_per_token,
617
attn_metadata.block_table,
618
topk_indices,
619
BLOCK_SIZE=attn_metadata.block_size,
620
NUM_TOPK_TOKENS=topk_indices.shape[1],
621
+ return_valid_counts=True,
622
)
623
624
return self._bf16_flash_mla_kernel(
625
q,
626
kv_c_and_k_pe_cache,
627
topk_indices,
628
+ topk_length,
629
)
630
631
def _forward_fp8_kv_separate_prefill_decode(
@@ -651,7 +653,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
653
# For BF16 cache: always use global cache slots (no workspace)
654
# prefill_workspace_starts has been adjusted in-place per chunk so
655
# prefill indices automatically come out chunk-local
654
- topk_indices = triton_convert_req_index_to_global_index(
656
+ topk_indices, topk_length = triton_convert_req_index_to_global_index(
657
attn_metadata.req_id_per_token,
658
attn_metadata.block_table,
659
topk_indices,
@@ -660,6 +662,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
662
HAS_PREFILL_WORKSPACE=has_prefill_workspace,
663
prefill_workspace_request_ids=prefill_request_ids,
664
prefill_workspace_starts=prefill_workspace_starts,
665
+ return_valid_counts=True,
666
)
667
668
fp8_metadata = attn_metadata.fp8_extra_metadata
@@ -722,11 +725,13 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
725
726
chunk_q = q[chunk.tokens_slice]
727
chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
728
+ chunk_topk_length = topk_length[chunk.tokens_slice]
729
730
attn_out[chunk.tokens_slice] = self._bf16_flash_mla_kernel(
731
chunk_q,
732
chunk_workspace,
733
chunk_topk_indices_workspace,
734
+ chunk_topk_length,
735
)
736
737
return attn_out
@@ -814,6 +819,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
819
q: torch.Tensor,
820
kv_c_and_k_pe_cache: torch.Tensor,
821
topk_indices: torch.Tensor,
822
+ topk_length: torch.Tensor | None = None,
823
) -> torch.Tensor:
824
num_tokens = q.shape[0]
825
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -834,7 +840,11 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
840
841
topk_indices = topk_indices.view(num_tokens, 1, -1)
842
output = flash_mla_sparse_fwd(
837
- q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale
843
+ q,
844
+ kv_c_and_k_pe_cache,
845
+ topk_indices,
846
+ self.softmax_scale,
847
+ topk_length=topk_length,
848
)[0]
849
850
output = output[:, : self.num_heads, :]