diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 1efa987fe7..48cd234038 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/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, diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 136a96a45d..1aaf3f1a14 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/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: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 22c6a38228..bfecb3c952 100644 --- a/vllm/v1/attention/backend.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index ac722dca9f..46b55cc7b2 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/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)