From fcf53d240c73d44bd9e3fe62996f8a2f1169fee9 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 6 Oct 2026 09:17:30 -0700 Subject: [PATCH] [Spec Decode] Remove eager metadata rebuild during MTP fused multi-step decode (#58463) Signed-off-by: Giancarlo Delfin --- .../test_indexer_deepseek_v4_slot_mapping.py | 61 +++++++++++++++++++ vllm/v1/attention/backend.py | 29 ++++----- .../backends/mla/flashattn_mla_sparse.py | 15 ++++- vllm/v1/attention/backends/mla/indexer.py | 22 +++++++ vllm/v1/attention/backends/mla/sparse_swa.py | 15 ++++- vllm/v1/attention/ops/metadata.py | 24 ++++++++ .../spec_decode/autoregressive/speculator.py | 50 ++++++--------- vllm/v1/worker/gpu/spec_decode/speculator.py | 16 +++++ 8 files changed, 182 insertions(+), 50 deletions(-) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 052088a621..17dea5d03c 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -122,6 +122,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = 16 + builder.kernel_block_size = None builder.num_sms = num_sms builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) builder.kv_cache_spec = SimpleNamespace(num_states=64) @@ -140,6 +141,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): requires_padding=False, schedule_metadata=schedule_metadata, ), + block_table=block_table, ) builder.update_draft_decode_metadata(metadata) @@ -167,6 +169,65 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): torch.testing.assert_close(metadata.slot_mapping, expected_slot_mapping) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + "block_size,kernel_block_size,compress_ratio", [(64, 64, 1), (256, 64, 4)] +) +def test_indexer_draft_decode_refresh_follows_new_batch( + monkeypatch, block_size, kernel_block_size, compress_ratio +): + """Without a prior build(), the refresh must not keep the last batch's state.""" + monkeypatch.setattr(indexer_module, "has_deep_gemm", lambda: False) + device = torch.device("cuda") + num_reqs, width = 3, 8 + factor = block_size // kernel_block_size + + builder = object.__new__(DeepseekV32IndexerMetadataBuilder) + builder.dcp_world_size = 1 + builder.compress_ratio = compress_ratio + builder.kernel_block_size = kernel_block_size + builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) + builder.kv_cache_spec = SimpleNamespace(block_size=block_size, num_states=64) + + def stale(*shape: int) -> torch.Tensor: + return torch.full(shape, -7, dtype=torch.int32, device=device) + + decode = DeepSeekV32IndexerDecodeMetadata( + block_table=stale(num_reqs, width // factor), + seq_lens=stale(num_reqs, 1), + decode_lens=stale(num_reqs), + requires_padding=False, + schedule_metadata=stale(1), + indices=stale(num_reqs), + per_req_decode_lens=stale(num_reqs), + ) + request_block_table = torch.randint( + 0, 1000, (num_reqs, width), dtype=torch.int32, device=device + ) + seq_lens = torch.tensor([17, 33, 65], dtype=torch.int32, device=device) + metadata = DeepseekV32IndexerMetadata( + seq_lens=seq_lens, + max_seq_len=128, + slot_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), + num_decodes=num_reqs, + num_decode_tokens=num_reqs, + num_prefills=0, + num_prefill_tokens=0, + decode=decode, + block_table=request_block_table, + ) + + builder.update_draft_decode_metadata(metadata) + + expected_block_table = request_block_table[:, ::factor] // factor + torch.testing.assert_close(decode.block_table, expected_block_table) + torch.testing.assert_close(decode.indices, builder.arange_buffer[:num_reqs]) + assert torch.all(decode.per_req_decode_lens == 1) + assert torch.all(decode.decode_lens == 1) + expected_seq_lens = torch.div(seq_lens, compress_ratio, rounding_mode="floor") + torch.testing.assert_close(decode.seq_lens.flatten(), expected_seq_lens) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize( "query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0], [0, 0, 0]] diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ffb68582ea..3387a31794 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -506,19 +506,13 @@ class CommonAttentionMetadata: # but not the right per-request boundaries. Padding requests have a query # length of zero and are skipped by the device boundary search. num_mapped_tokens = int(self.query_start_loc_cpu[-1]) - from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel + from vllm.v1.attention.ops.metadata import compute_token_to_req_indices num_output_tokens = max(num_mapped_tokens, num_tokens) assert buffer.shape[0] >= num_output_tokens - _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)]( - self.query_start_loc, - buffer, - self.query_start_loc.shape[0] - 1, - num_mapped_tokens, - num_output_tokens, - num_warps=4, + self._token_to_req_indices_cache = compute_token_to_req_indices( + self.query_start_loc, buffer, num_mapped_tokens, num_output_tokens ) - self._token_to_req_indices_cache = buffer[:num_output_tokens] return self._token_to_req_indices_cache[:num_tokens] # TODO(lucas): remove once we have FULL-CG spec-decode support @@ -609,8 +603,11 @@ class AttentionMetadataBuilder(ABC, Generic[M]): supports_update_block_table: bool = False # Whether the builder constructor requires the block-table width. requires_block_table_width: ClassVar[bool] = False - # Whether all step-dependent draft decode metadata can be updated in place, - # allowing one metadata build to be reused across autoregressive draft steps. + # Whether update_draft_decode_metadata() can regenerate all decode metadata + # from persistent device buffers with capture-safe ops only. Lets a + # speculator record the refresh inside its draft decode CUDA graph, so no + # eager build() is needed between steps or between the batches that replay + # it, whether the graph holds one draft forward or a multi-step loop. supports_draft_decode_metadata_update: bool = False @abstractmethod @@ -758,12 +755,12 @@ class AttentionMetadataBuilder(ABC, Generic[M]): ) def update_draft_decode_metadata(self, metadata: M) -> None: - """Update step-dependent draft decode metadata in place. + """Update draft decode metadata in place. - The fused draft loop may call this method during full CUDA graph - capture. CUDA graph replay does not run this Python method, so - implementations must emit capture-safe operations and keep replayed - tensor state in persistent storage. + Speculators can record this call inside their draft decode CUDA graphs, + so it runs once at capture and never again in Python. Implementations + must emit only capture-safe operations and write results into persistent + tensors that ``metadata`` already references. """ raise NotImplementedError diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 3ca88992aa..aabd682240 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -27,6 +27,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, ) +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func @@ -132,8 +133,18 @@ class FlashAttnMLASparseMetadataBuilder( self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True) self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 - def update_draft_decode_metadata(self, _metadata: SparseMLACommonMetadata) -> None: - pass + def update_draft_decode_metadata(self, metadata: SparseMLACommonMetadata) -> None: + num_tokens = metadata.num_decode_tokens + if num_tokens == 0: + return + # Everything else is a view of runner buffers; the per-token request + # map is a builder buffer that build() filled for the captured batch. + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.req_id_per_token, + num_tokens, + num_tokens, + ) class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 9d3823b880..8b186502f2 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -678,6 +678,7 @@ class DeepseekV32IndexerMetadata: decode: DeepSeekV32IndexerDecodeMetadata | None = None prefill: DeepseekV32IndexerPrefillMetadata | None = None positions: torch.Tensor | None = None + block_table: torch.Tensor | None = None @triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) @@ -1713,6 +1714,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): num_prefill_tokens=num_prefill_tokens, prefill=prefill_metadata, decode=decode_metadata, + block_table=common_attn_metadata.block_table_tensor, ) return attn_metadata @@ -1728,8 +1730,28 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): assert metadata.num_prefills == 0 assert metadata.num_decodes == metadata.num_decode_tokens assert decode.seq_lens.numel() == metadata.num_decode_tokens + assert metadata.block_table is not None assert self.dcp_world_size == 1 + # One query token per request, so the per-token decode block table is + # the request block table (at the indexer kernel's page size). build() + # may copy it into a builder buffer, so re-derive it for this batch. + block_table = metadata.block_table[: metadata.num_decode_tokens] + kernel_block_size = self.kernel_block_size + if ( + self.compress_ratio > 1 + and kernel_block_size is not None + and self.kv_cache_spec.block_size != kernel_block_size + and self.kv_cache_spec.block_size % kernel_block_size == 0 + ): + factor = self.kv_cache_spec.block_size // kernel_block_size + block_table = block_table[:, ::factor] // factor + decode.block_table.copy_(block_table) + if decode.indices is not None: + decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens]) + if decode.per_req_decode_lens is not None: + decode.per_req_decode_lens.fill_(1) + if self.compress_ratio > 1: get_compressed_slot_mapping( metadata.num_decode_tokens, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 0f5bff4144..c5cbb4ef28 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -33,6 +33,7 @@ from vllm.v1.attention.backends.mla.compressor_utils import ( ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, @@ -829,7 +830,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): self, metadata: DeepseekSparseSWAMetadata, ) -> None: - if metadata.num_decode_tokens == 0: + num_tokens = metadata.num_decode_tokens + if num_tokens == 0: return assert metadata.query_start_loc is not None assert metadata.seq_lens is not None @@ -839,6 +841,17 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): assert metadata.decode_swa_lens is not None assert metadata.replay_start is not None + # Recompute the validity mask and token->request map from the device + # buffers. Their padding differs from the dummy batch the graph was + # captured with. + torch.ge(metadata.slot_mapping, 0, out=metadata.is_valid_token) + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.token_to_req_indices, + num_tokens, + num_tokens, + ) + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( metadata.decode_swa_indices, metadata.decode_swa_lens, diff --git a/vllm/v1/attention/ops/metadata.py b/vllm/v1/attention/ops/metadata.py index f5baa51a92..1852e77ee2 100644 --- a/vllm/v1/attention/ops/metadata.py +++ b/vllm/v1/attention/ops/metadata.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Device-side request mapping and sparse indexer metadata.""" +import torch + from vllm.triton_utils import tl, triton @@ -27,6 +29,28 @@ def _token_request_mapping_kernel( tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens) +def compute_token_to_req_indices( + query_start_loc: torch.Tensor, + out: torch.Tensor, + num_mapped_tokens: int, + num_tokens: int, +) -> torch.Tensor: + """Map each of the first ``num_tokens`` tokens to its request index. + + Reads only the device ``query_start_loc``, so it is safe to record in a + CUDA graph. Tokens at or past ``num_mapped_tokens`` are mapped to 0. + """ + _token_request_mapping_kernel[((num_tokens + 255) // 256,)]( + query_start_loc, + out, + query_start_loc.shape[0] - 1, + num_mapped_tokens, + num_tokens, + num_warps=4, + ) + return out[:num_tokens] + + @triton.jit( do_not_specialize=[ "num_reqs", diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 7242a840d2..46ff826815 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -99,10 +99,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.use_fused_multi_step_decode = False return - if not self.advance_draft_positions: - self.use_fused_multi_step_decode = True - return - unsupported_backends = sorted( { attn_group.backend.get_name() @@ -574,6 +570,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] idx_mapping = self.idx_mapping[:num_reqs] + if batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.run_fullgraph(batch_desc) + return + attn_metadata = None slot_mappings_by_layer = None if not skip_attn: @@ -583,10 +584,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): positions, batch_desc.num_tokens, ) - if batch_desc.cg_mode != CUDAGraphMode.FULL: - slot_mappings_by_layer = build_slot_mappings_by_layer( - slot_mappings, self.kv_cache_config - ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, self.kv_cache_config + ) attn_metadata = self._build_uniform_attn_metadata( num_reqs=num_reqs, batch_desc=batch_desc, @@ -595,11 +595,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): step=1, ) - if batch_desc.cg_mode == CUDAGraphMode.FULL: - assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.run_fullgraph(batch_desc) - return - self._generate_fused_drafts( num_reqs, batch_desc.num_tokens, @@ -623,13 +618,19 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): idx_mapping = self.idx_mapping[:num_reqs] positions = self.input_buffers.positions[:num_reqs] query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] - attn_groups = ( - [group for groups in self.attn_groups for group in groups] - if attn_metadata is not None - else [] - ) for step in range(1, num_speculative_steps): + if attn_metadata is not None and ( + self.advance_draft_positions or step == 1 + ): + self.block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded, + ) + self._update_draft_decode_metadata(attn_metadata, num_reqs) + self.current_draft_step.fill_(step) self._generate_draft( num_reqs, @@ -640,19 +641,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): cudagraph_runtime_mode, num_speculative_steps, ) - if ( - step < num_speculative_steps - 1 - and attn_metadata is not None - and self.advance_draft_positions - ): - self.block_tables.compute_slot_mappings( - idx_mapping, - query_start_loc, - positions, - num_tokens_padded, - ) - for attn_group in attn_groups: - attn_group.update_draft_decode_metadata(attn_metadata) def _generate_draft( self, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 1110663288..d07c043de9 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -562,3 +562,19 @@ class DraftModelSpeculator(BaseSpeculator): causal=causal, dcp_local_seq_lens=dcp_local_seq_lens, ) + + def _update_draft_decode_metadata( + self, attn_metadata: dict[str, Any], num_reqs: int + ) -> None: + if self.block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + for groups in self.attn_groups: + for group in groups: + group.update_draft_decode_metadata(attn_metadata)