[Perf] Cache `token_to_req_indices` for dsv4, 5x~6x kernel performance improvement (#47474)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 6, 2026 at 11:17 UTC
f70caef48b9263e8f86fdf20a04a4933a5f75083
4 files changed
+34
-25
vllm/models/deepseek_v4/compressor.py
+3
-6
@@ -104,12 +104,9 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder):
104
common_attn_metadata: CommonAttentionMetadata,
105
fast_build: bool = False,
106
) -> CompressorMetadata:
107
- query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
108
- num_reqs = common_attn_metadata.num_reqs
109
- query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
110
- x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
111
- token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
112
- token_to_req_indices.copy_(x, non_blocking=True)
107
+ token_to_req_indices = common_attn_metadata.token_to_req_indices(
108
+ self.token_to_req_indices
109
+ )
110
return CompressorMetadata(
111
block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
112
slot_mapping=common_attn_metadata.slot_mapping,
vllm/models/deepseek_v4/sparse_mla.py
+1
-14
@@ -5,7 +5,6 @@
5
from dataclasses import dataclass
6
from typing import Any, ClassVar
7
8
-import numpy as np
8
import torch
9
10
from vllm.config import VllmConfig
@@ -13,7 +12,6 @@ from vllm.config.cache import CacheDType
12
from vllm.platforms.interface import DeviceCapability
13
from vllm.triton_utils import tl, triton
14
from vllm.utils.math_utils import cdiv
16
-from vllm.utils.torch_utils import np_to_pinned_tensor
15
from vllm.v1.attention.backend import (
16
AttentionBackend,
17
AttentionCGSupport,
@@ -203,18 +201,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
201
fast_build: bool = False,
202
) -> DeepseekV4FlashMLAMetadata:
203
cm = common_attn_metadata
206
- num_tokens = cm.num_actual_tokens
207
- starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
208
- seg_lengths = np.diff(starts)
209
- req_id_per_token = np.repeat(
210
- np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
211
- )
212
- # Zero-fill for cudagraphs
213
- self.req_id_per_token_buffer.fill_(0)
214
- self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
215
- np_to_pinned_tensor(req_id_per_token), non_blocking=True
216
- )
217
- req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
204
+ req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer)
205
206
slot_mapping = cm.slot_mapping
207
if self.compress_ratio > 1:
vllm/v1/attention/backend.py
+27
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
kFp8StaticTensorSym,
17
kNvfp4Dynamic,
18
)
19
+from vllm.utils.torch_utils import np_to_pinned_tensor
20
21
if TYPE_CHECKING:
22
from vllm.config import VllmConfig
@@ -467,6 +468,7 @@ class CommonAttentionMetadata:
468
_num_computed_tokens_cpu: torch.Tensor | None = None
469
470
_num_computed_tokens_cache: torch.Tensor | None = None
471
+ _token_to_req_indices_cache: torch.Tensor | None = None
472
473
def batch_size(self) -> int:
474
return self.seq_lens.shape[0]
@@ -515,6 +517,31 @@ class CommonAttentionMetadata:
517
self._num_computed_tokens_cache = self.seq_lens - query_lens
518
return self._num_computed_tokens_cache
519
520
+ def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
521
+ """Build or reuse the per-token request index mapping."""
522
+ num_tokens = self.num_actual_tokens
523
+ if self._token_to_req_indices_cache is not None:
524
+ assert self._token_to_req_indices_cache.device == buffer.device
525
+ assert self._token_to_req_indices_cache.dtype == torch.int32
526
+ assert self._token_to_req_indices_cache.shape[0] >= num_tokens
527
+ return self._token_to_req_indices_cache[:num_tokens]
528
+
529
+ starts = np.asarray(self.query_start_loc_cpu, dtype=np.int32)
530
+ query_lens = np.diff(starts)
531
+ token_to_req_indices = np.repeat(
532
+ np.arange(query_lens.shape[0], dtype=np.int32), query_lens
533
+ )
534
+ num_mapped_tokens = token_to_req_indices.shape[0]
535
+ assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
536
+ # copy from CPU to GPU
537
+ buffer[:num_mapped_tokens].copy_(
538
+ np_to_pinned_tensor(token_to_req_indices), non_blocking=True
539
+ )
540
+ if num_mapped_tokens < num_tokens:
541
+ buffer[num_mapped_tokens:num_tokens].zero_()
542
+ self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
543
+ return self._token_to_req_indices_cache[:num_tokens]
544
+
545
# TODO(lucas): remove once we have FULL-CG spec-decode support
546
def unpadded(
547
self, num_actual_tokens: int, num_actual_reqs: int
vllm/v1/attention/backends/mla/sparse_swa.py
+3
-5
@@ -399,7 +399,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
399
400
For prefill, we use chunked prefill to align with the indexer's chunking.
401
"""
402
- num_reqs = common_attn_metadata.num_reqs
402
seq_lens = common_attn_metadata.seq_lens
403
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
404
query_start_loc = common_attn_metadata.query_start_loc
@@ -416,10 +415,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
415
416
# NOTE: Ensure all metadata tensors maintain fixed memory addresses
417
# for CUDA graph compatibility.
419
- query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
420
- x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
421
- token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
422
- token_to_req_indices.copy_(x, non_blocking=True)
418
+ token_to_req_indices = common_attn_metadata.token_to_req_indices(
419
+ self.token_to_req_indices
420
+ )
421
422
is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
423
is_valid_token.copy_(slot_mapping >= 0)