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 3, 14:04 UTC
[ROCm] Restore fast top_k_per_row kernels for sparse MLA when topk_tokens=2048 (#42072)
Signed-off-by: Frida Andersson <fanderss@amd.com>
frida-andersson committed
May 15, 2026 at 21:02 UTC
4d67d3bde25f94b6199ce16c7ef239ae4412bb8f
1 file changed
+77
-4
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+77
-4
@@ -562,6 +562,64 @@ def _topk_indices_torch(logits: torch.Tensor, topk_tokens: int) -> torch.Tensor:
562
return padded
563
564
565
+# topk_tokens values with dedicated fused C++ kernel support.
566
+_TOPK_FAST_PATH_VALUES = frozenset({2048})
567
+
568
+
569
+def _topk_indices_prefill(
570
+ logits: torch.Tensor,
571
+ topk_tokens: int,
572
+ topk_out: torch.Tensor,
573
+ cu_seqlen_ks: torch.Tensor,
574
+ cu_seqlen_ke: torch.Tensor,
575
+) -> None:
576
+ """Top-k indices for the prefill path.
577
+
578
+ Writes ``logits.shape[0]`` rows into ``topk_out``; caller must size the
579
+ view accordingly.
580
+ """
581
+ if topk_tokens in _TOPK_FAST_PATH_VALUES:
582
+ torch.ops._C.top_k_per_row_prefill(
583
+ logits,
584
+ cu_seqlen_ks,
585
+ cu_seqlen_ke,
586
+ topk_out,
587
+ logits.shape[0],
588
+ logits.stride(0),
589
+ logits.stride(1),
590
+ topk_tokens,
591
+ )
592
+ else:
593
+ topk_out.copy_(_topk_indices_torch(logits, topk_tokens))
594
+
595
+
596
+def _topk_indices_decode(
597
+ logits: torch.Tensor,
598
+ topk_tokens: int,
599
+ topk_out: torch.Tensor,
600
+ seq_lens: torch.Tensor,
601
+ next_n: int,
602
+) -> None:
603
+ """Top-k indices for the decode path.
604
+
605
+ Writes ``logits.shape[0] == batch_size * next_n`` rows into ``topk_out``;
606
+ caller must size the view to ``num_padded_tokens``.
607
+ """
608
+ if topk_tokens in _TOPK_FAST_PATH_VALUES:
609
+ torch.ops._C.top_k_per_row_decode(
610
+ logits,
611
+ next_n,
612
+ seq_lens,
613
+ topk_out,
614
+ logits.shape[0],
615
+ logits.stride(0),
616
+ logits.stride(1),
617
+ topk_tokens,
618
+ )
619
+ else:
620
+ topk_out.copy_(_topk_indices_torch(logits, topk_tokens))
621
+
622
+
623
def rocm_aiter_sparse_attn_indexer_fake(
624
hidden_states: torch.Tensor,
625
k_cache_prefix: LayerNameType,
@@ -709,7 +767,13 @@ def rocm_aiter_sparse_attn_indexer_native(
767
topk_indices = topk_indices_buffer[
768
chunk.token_start : chunk.token_end, :topk_tokens
769
]
712
- topk_indices.copy_(_topk_indices_torch(logits, topk_tokens))
770
+ _topk_indices_prefill(
771
+ logits,
772
+ topk_tokens,
773
+ topk_indices,
774
+ chunk.cu_seqlen_ks,
775
+ chunk.cu_seqlen_ke,
776
+ )
777
778
if has_decode:
779
decode_metadata = layer_attn_metadata.decode
@@ -746,14 +810,23 @@ def rocm_aiter_sparse_attn_indexer_native(
810
max_model_len=max_model_len,
811
)
812
749
- topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
750
- topk_indices.copy_(_topk_indices_torch(logits, topk_tokens)[:num_decode_tokens])
813
+ # Size the view to num_padded_tokens: top_k_per_row_decode writes
814
+ # logits.shape[0] == num_padded_tokens rows, and the unpack below
815
+ # reshapes to (batch_size, next_n, ...).
816
+ topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
817
+ _topk_indices_decode(
818
+ logits,
819
+ topk_tokens,
820
+ topk_indices,
821
+ decode_metadata.seq_lens,
822
+ next_n,
823
+ )
824
825
if decode_metadata.requires_padding:
826
# if padded, we need to unpack
827
# the topk indices removing padded tokens
828
topk_indices = unpack_seq_triton(
756
- topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
829
+ topk_indices.reshape(batch_size, next_n, topk_indices.shape[-1]),
830
decode_lens,
831
)
832
topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (