[ROCm][DSv4.1][Perf] Reuse the decode topk ragged metadata across layers (#57434)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
Fangzhou Ai committed
Sep 19, 2026 at 23:46 UTC
1596fa5f8825d1b949463d51f3ed03650b0f539a
1 file changed
+66
-13
vllm/models/deepseek_v41/amd/rocm.py
+66
-13
@@ -13,7 +13,10 @@ from vllm.distributed import (
13
)
14
from vllm.forward_context import get_forward_context
15
from vllm.logger import init_logger
16
-from vllm.models.deepseek_v41.attention import DeepseekV4Attention
16
+from vllm.models.deepseek_v41.attention import (
17
+ DeepseekV4Attention,
18
+ _replace_layer_index,
19
+)
20
from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
21
from vllm.models.deepseek_v41.sparse_mla import (
22
DeepseekV4FlashMLAMetadata,
@@ -391,6 +394,10 @@ def _copy_ragged_to_graph_buffers(
394
return ragged_out, indptr_out
395
396
397
+# (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode.
398
+_TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
399
+
400
+
401
@dataclass
402
class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
403
decode_swa_ragged_indices: torch.Tensor | None = None
@@ -500,6 +507,23 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
507
self._fused_compressor_weight: torch.Tensor | None
508
self.register_buffer("_fused_compressor_weight", None, persistent=False)
509
self._fused_compressor_split_sizes: tuple[int, int] | None = None
510
+ # Decode ragged topk metadata is a pure function of the indices its
511
+ # index source published, so every consumer of a source rebuilds the
512
+ # same thing. The source owns one cache per compress ratio, refreshed
513
+ # when it runs; consumers below it read through.
514
+ self._topk_ragged_cache: dict[int, _TopkRagged] = {}
515
+ self._index_source_prefix: str | None = None
516
+ if self.compress_ratio > 0:
517
+ assert self.index_source_layer_id is not None
518
+ self._index_source_prefix = _replace_layer_index(
519
+ self.prefix, self.index_source_layer_id
520
+ )
521
+ if self._index_source_prefix not in self._static_forward_context:
522
+ raise NotImplementedError(
523
+ f"Index source {self._index_source_prefix} not found on "
524
+ "this rank; PP splits inside a v4.1 index-sharing group "
525
+ "are not supported."
526
+ )
527
528
@classmethod
529
def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -733,6 +757,42 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
757
output=output[:num_decode_tokens],
758
)
759
760
+ def _decode_topk_ragged(
761
+ self,
762
+ swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
763
+ attn_metadata: DeepseekV4FlashMLAMetadata | None,
764
+ num_decodes: int,
765
+ num_decode_tokens: int,
766
+ ) -> _TopkRagged:
767
+ """Ragged form of the topk indices this layer's index source published.
768
+
769
+ The packing depends only on the shared ``topk_indices_buffer`` and
770
+ per-step metadata, apart from the layer's own compress ratio, so the
771
+ source memoizes one result per ratio for all the layers below it.
772
+ """
773
+ assert attn_metadata is not None
774
+ assert swa_metadata.is_valid_token is not None
775
+ assert self.topk_indices_buffer is not None
776
+ assert self._index_source_prefix is not None
777
+
778
+ source = self._static_forward_context[self._index_source_prefix]
779
+ if source is self:
780
+ # Fresh indices as of this layer; drop what the last step cached.
781
+ source._topk_ragged_cache = {}
782
+ cached = source._topk_ragged_cache.get(self.compress_ratio)
783
+ if cached is not None:
784
+ return cached
785
+
786
+ built = compute_global_topk_ragged_indices_and_indptr(
787
+ self.topk_indices_buffer[:num_decode_tokens],
788
+ swa_metadata.token_to_req_indices,
789
+ attn_metadata.block_table[:num_decodes],
790
+ attn_metadata.block_size // self.compress_ratio,
791
+ swa_metadata.is_valid_token[:num_decode_tokens],
792
+ )
793
+ source._topk_ragged_cache[self.compress_ratio] = built
794
+ return built
795
+
796
def _forward_decode(
797
self,
798
q: torch.Tensor,
@@ -749,22 +809,15 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
809
topk_ragged_indices = None
810
topk_ragged_indptr = None
811
if not swa_only:
752
- # Local indices filled by the index-source layer's indexer.
753
- assert attn_metadata is not None
754
- assert swa_metadata.is_valid_token is not None
755
- assert self.topk_indices_buffer is not None
756
- block_size = attn_metadata.block_size // self.compress_ratio
757
- is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
812
(
813
topk_ragged_indices,
814
topk_ragged_indptr,
815
topk_lens,
762
- ) = compute_global_topk_ragged_indices_and_indptr(
763
- self.topk_indices_buffer[:num_decode_tokens],
764
- swa_metadata.token_to_req_indices,
765
- attn_metadata.block_table[:num_decodes],
766
- block_size,
767
- is_valid,
816
+ ) = self._decode_topk_ragged(
817
+ swa_metadata=swa_metadata,
818
+ attn_metadata=attn_metadata,
819
+ num_decodes=num_decodes,
820
+ num_decode_tokens=num_decode_tokens,
821
)
822
823
rocm_sparse_attn_decode(