[Perf] Cache token_to_req_indices for dsv4, 5x~6x kernel performance improvement (#47474)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/models/deepseek_v4/compressor.py

@@ -104,12 +104,9 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder):
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> CompressorMetadata:
-        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
-        num_reqs = common_attn_metadata.num_reqs
-        query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
-        x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
-        token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
-        token_to_req_indices.copy_(x, non_blocking=True)
+        token_to_req_indices = common_attn_metadata.token_to_req_indices(
+            self.token_to_req_indices
+        )
         return CompressorMetadata(
             block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
             slot_mapping=common_attn_metadata.slot_mapping,

vllm/models/deepseek_v4/sparse_mla.py

@@ -5,7 +5,6 @@
 from dataclasses import dataclass
 from typing import Any, ClassVar
 
-import numpy as np
 import torch
 
 from vllm.config import VllmConfig
@@ -13,7 +12,6 @@ from vllm.config.cache import CacheDType
 from vllm.platforms.interface import DeviceCapability
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -203,18 +201,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
         fast_build: bool = False,
     ) -> DeepseekV4FlashMLAMetadata:
         cm = common_attn_metadata
-        num_tokens = cm.num_actual_tokens
-        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
-        seg_lengths = np.diff(starts)
-        req_id_per_token = np.repeat(
-            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
-        )
-        # Zero-fill for cudagraphs
-        self.req_id_per_token_buffer.fill_(0)
-        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            np_to_pinned_tensor(req_id_per_token), non_blocking=True
-        )
-        req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
+        req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer)
 
         slot_mapping = cm.slot_mapping
         if self.compress_ratio > 1:

vllm/v1/attention/backend.py

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
     kNvfp4Dynamic,
 )
+from vllm.utils.torch_utils import np_to_pinned_tensor
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -467,6 +468,7 @@ class CommonAttentionMetadata:
     _num_computed_tokens_cpu: torch.Tensor | None = None
 
     _num_computed_tokens_cache: torch.Tensor | None = None
+    _token_to_req_indices_cache: torch.Tensor | None = None
 
     def batch_size(self) -> int:
         return self.seq_lens.shape[0]
@@ -515,6 +517,31 @@ class CommonAttentionMetadata:
             self._num_computed_tokens_cache = self.seq_lens - query_lens
         return self._num_computed_tokens_cache
 
+    def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
+        """Build or reuse the per-token request index mapping."""
+        num_tokens = self.num_actual_tokens
+        if self._token_to_req_indices_cache is not None:
+            assert self._token_to_req_indices_cache.device == buffer.device
+            assert self._token_to_req_indices_cache.dtype == torch.int32
+            assert self._token_to_req_indices_cache.shape[0] >= num_tokens
+            return self._token_to_req_indices_cache[:num_tokens]
+
+        starts = np.asarray(self.query_start_loc_cpu, dtype=np.int32)
+        query_lens = np.diff(starts)
+        token_to_req_indices = np.repeat(
+            np.arange(query_lens.shape[0], dtype=np.int32), query_lens
+        )
+        num_mapped_tokens = token_to_req_indices.shape[0]
+        assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
+        # copy from CPU to GPU
+        buffer[:num_mapped_tokens].copy_(
+            np_to_pinned_tensor(token_to_req_indices), non_blocking=True
+        )
+        if num_mapped_tokens < num_tokens:
+            buffer[num_mapped_tokens:num_tokens].zero_()
+        self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
+        return self._token_to_req_indices_cache[:num_tokens]
+
     # TODO(lucas): remove once we have FULL-CG spec-decode support
     def unpadded(
         self, num_actual_tokens: int, num_actual_reqs: int

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -399,7 +399,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
 
         For prefill, we use chunked prefill to align with the indexer's chunking.
         """
-        num_reqs = common_attn_metadata.num_reqs
         seq_lens = common_attn_metadata.seq_lens
         seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
         query_start_loc = common_attn_metadata.query_start_loc
@@ -416,10 +415,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
 
         # NOTE: Ensure all metadata tensors maintain fixed memory addresses
         # for CUDA graph compatibility.
-        query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
-        x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
-        token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
-        token_to_req_indices.copy_(x, non_blocking=True)
+        token_to_req_indices = common_attn_metadata.token_to_req_indices(
+            self.token_to_req_indices
+        )
 
         is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
         is_valid_token.copy_(slot_mapping >= 0)