[Spec Decode] Remove eager metadata rebuild during MTP fused multi-step decode (#58463)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

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]]

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
 

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]):

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,

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,

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",

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,

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)