[MRV2][Spec] Fuse AR speculator multi-step decodes back into one CUDA graph (#46849)

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

docs/design/model_runner_v2.md

@@ -191,6 +191,14 @@ V1's CUDA graph handling is implicit and hard to reason about. MRV2 uses a `CUDA
 
 This makes graph lifecycle and execution mode decisions more understandable and easier to extend. Example: MRV2 can capture multiple draft-model forward passes into one CUDA graph.
 
+### Fused Multi-Step Draft Decoding
+
+Autoregressive speculative decoding executes several dependent draft steps per scheduler step. In the fused path, MRV2 captures all post-prefill draft steps in one full CUDA graph instead of replaying a separate graph for each draft token. Attention metadata is built once before the loop, while common step-dependent tensors keep stable addresses and are updated in place between draft steps.
+
+Some attention backends also materialize derived state, such as scheduler metadata or sparse indices. Before opting into the fused path, these backends must implement `AttentionMetadataBuilder.update_draft_decode_metadata()` to update or invalidate that state after the draft inputs advance. The hook runs during CUDA graph capture, so only the GPU operations it issues are recorded and executed during replay; its Python body is not run again. Implementations must therefore use capture-safe operations and keep all replayed tensor state in persistent storage.
+
+For draft models that advance positions, the fused path is enabled only when every draft attention group declares `supports_draft_decode_metadata_update`. Otherwise, MRV2 falls back to rebuilding attention metadata between draft steps. Draft models that keep positions fixed do not require this update. Before enabling a backend, developers must audit all derived metadata, including state inherited from parent builders or owned by auxiliary attention backends.
+
 ## Development Philosophy
 
 MRV2 changes should meet a higher code quality bar. As feature gaps with V1 are filled, features should be reconsidered from first principles in the MRV2 design context instead of quickly porting V1 behavior.

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -3,6 +3,7 @@
 
 from contextlib import nullcontext
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -16,6 +17,9 @@ from vllm.model_executor.models.mistral_eagle import EagleMistralForCausalLM
 from vllm.model_executor.models.mistral_large_3_eagle import (
     EagleMistralLarge3ForCausalLM,
 )
+from vllm.v1.attention.backends import flash_attn as flash_attn_module
+from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
 from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
@@ -278,3 +282,165 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
 
     assert actual_logits_hidden is hidden
     assert actual_feedback_hidden is hidden
+
+
+@pytest.mark.parametrize(
+    (
+        "method_name",
+        "cg_mode",
+        "expected_eager_calls",
+        "expected_graph_replays",
+    ),
+    [
+        ("_multi_step_decode", CUDAGraphMode.NONE, 3, 0),
+        ("_multi_step_decode", CUDAGraphMode.FULL, 0, 3),
+        ("_fused_multi_step_decode", CUDAGraphMode.NONE, 3, 0),
+        ("_fused_multi_step_decode", CUDAGraphMode.FULL, 0, 1),
+    ],
+)
+def test_multi_step_decode_replays_captured_graph_as_expected(
+    method_name,
+    cg_mode,
+    expected_eager_calls,
+    expected_graph_replays,
+):
+    speculator = object.__new__(_TestSpeculator)
+    speculator.num_speculative_steps = 4
+    speculator.current_draft_step = torch.tensor(0)
+    speculator.input_buffers = SimpleNamespace(
+        positions=torch.arange(2),
+        query_start_loc=torch.arange(3),
+    )
+    speculator.idx_mapping = torch.arange(2)
+    generate_draft = Mock()
+    speculator._generate_draft = generate_draft
+    run_fullgraph = Mock()
+    speculator.decode_cudagraph_manager = SimpleNamespace(run_fullgraph=run_fullgraph)
+    batch_desc = BatchExecutionDescriptor(
+        cg_mode=cg_mode,
+        num_tokens=2,
+        num_reqs=2,
+    )
+
+    getattr(speculator, method_name)(
+        num_reqs=2,
+        skip_attn=True,
+        batch_desc=batch_desc,
+        seq_lens_cpu_upper_bound=None,
+        num_tokens_across_dp=None,
+    )
+
+    assert generate_draft.call_count == expected_eager_calls
+    assert run_fullgraph.call_count == expected_graph_replays
+
+
+def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata(
+    monkeypatch,
+):
+    builder = object.__new__(flash_attn_module.FlashAttentionMetadataBuilder)
+    builder.aot_schedule = True
+    builder.use_full_cuda_graph = True
+    builder.scheduler_metadata = torch.zeros(8, dtype=torch.int32)
+    builder.cache_config = SimpleNamespace(cache_dtype="bfloat16")
+    builder.kv_cache_dtype = torch.bfloat16
+    builder.num_heads_q = 2
+    builder.num_heads_kv = 1
+    builder.headdim = 128
+    builder.block_size = 16
+    builder.dcp_world_size = 1
+    builder.dcp_rank = 0
+    builder.cp_kv_cache_interleave_size = 1
+    builder.aot_sliding_window = None
+
+    expected = torch.tensor([7, 8, 9], dtype=torch.int32)
+
+    def fake_get_scheduler_metadata(**kwargs):
+        return expected
+
+    monkeypatch.setattr(builder, "_get_scheduler_metadata", fake_get_scheduler_metadata)
+
+    metadata = FlashAttentionMetadata(
+        num_actual_tokens=3,
+        max_query_len=2,
+        query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),
+        max_seq_len=8,
+        seq_lens=torch.tensor([5, 6], dtype=torch.int32),
+        block_table=torch.zeros((2, 1), dtype=torch.int32),
+        slot_mapping=torch.zeros(3, dtype=torch.int32),
+        use_cascade=False,
+        common_prefix_len=0,
+        cu_prefix_query_lens=None,
+        prefix_kv_lens=None,
+        suffix_kv_lens=None,
+        max_dcp_context_kv_len=None,
+        dcp_context_kv_lens=None,
+        num_decode_reqs=2,
+        num_prefill_reqs=0,
+        num_decode_tokens=3,
+        num_prefill_tokens=0,
+        scheduler_metadata=torch.tensor([-1, -1, -1], dtype=torch.int32),
+        prefix_scheduler_metadata=None,
+        max_num_splits=4,
+        causal=True,
+        sliding_window=None,
+        mm_prefix_query_range_tensor=None,
+        rswa_prefix_lens=None,
+        rswa_window=None,
+        rswa_window_tensor=None,
+    )
+
+    builder.update_draft_decode_metadata(metadata)
+
+    assert torch.equal(metadata.scheduler_metadata, expected)
+    assert torch.equal(builder.scheduler_metadata[:3], expected)
+
+
+def test_update_draft_decode_metadata_skips_without_scheduler_metadata(monkeypatch):
+    builder = object.__new__(flash_attn_module.FlashAttentionMetadataBuilder)
+    builder.aot_schedule = True
+    builder.use_full_cuda_graph = True
+    builder.scheduler_metadata = torch.zeros(4, dtype=torch.int32)
+
+    called = False
+
+    def fake_get_scheduler_metadata(**kwargs):
+        nonlocal called
+        called = True
+        return torch.tensor([1], dtype=torch.int32)
+
+    monkeypatch.setattr(builder, "_get_scheduler_metadata", fake_get_scheduler_metadata)
+
+    metadata = FlashAttentionMetadata(
+        num_actual_tokens=1,
+        max_query_len=1,
+        query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
+        max_seq_len=1,
+        seq_lens=torch.tensor([1], dtype=torch.int32),
+        block_table=torch.zeros((1, 1), dtype=torch.int32),
+        slot_mapping=torch.zeros(1, dtype=torch.int32),
+        use_cascade=False,
+        common_prefix_len=0,
+        cu_prefix_query_lens=None,
+        prefix_kv_lens=None,
+        suffix_kv_lens=None,
+        max_dcp_context_kv_len=None,
+        dcp_context_kv_lens=None,
+        num_decode_reqs=1,
+        num_prefill_reqs=0,
+        num_decode_tokens=1,
+        num_prefill_tokens=0,
+        scheduler_metadata=None,
+        prefix_scheduler_metadata=None,
+        max_num_splits=1,
+        causal=True,
+        sliding_window=None,
+        mm_prefix_query_range_tensor=None,
+        rswa_prefix_lens=None,
+        rswa_window=None,
+        rswa_window_tensor=None,
+    )
+
+    builder.update_draft_decode_metadata(metadata)
+
+    assert not called
+    assert metadata.scheduler_metadata is None

vllm/models/deepseek_v4/amd/rocm.py

@@ -372,6 +372,10 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4FlashMLAMetadataBuil
 
 
 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
+    # Keep fused multi-step decode disabled until update_draft_decode_metadata()
+    # also refreshes the ROCm-specific ragged SWA indices and indptrs.
+    supports_draft_decode_metadata_update = False
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens

vllm/v1/attention/backend.py

@@ -633,6 +633,9 @@ 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.
+    supports_draft_decode_metadata_update: bool = False
 
     @abstractmethod
     def __init__(
@@ -757,6 +760,16 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
             fast_build=True,
         )
 
+    def update_draft_decode_metadata(self, metadata: M) -> None:
+        """Update step-dependent 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.
+        """
+        raise NotImplementedError
+
     def use_cascade_attention(
         self,
         common_prefix_len: int,

vllm/v1/attention/backends/flash_attn.py

@@ -370,6 +370,57 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
     ) -> AttentionCGSupport:
         return cls._cudagraph_support
 
+    def _get_scheduler_metadata(
+        self,
+        *,
+        aot_schedule: bool,
+        batch_size: int,
+        cu_query_lens: torch.Tensor,
+        max_query_len: int,
+        seqlens: torch.Tensor,
+        max_seq_len: int,
+        causal: bool | torch.Tensor,
+        max_num_splits: int,
+    ) -> torch.Tensor | None:
+        if not aot_schedule:
+            return None
+
+        cache_dtype = self.cache_config.cache_dtype
+        if is_quantized_kv_cache(cache_dtype):
+            qkv_dtype = current_platform.fp8_dtype()
+        else:
+            qkv_dtype = self.kv_cache_dtype
+        return get_scheduler_metadata(
+            batch_size=batch_size,
+            max_seqlen_q=max_query_len,
+            max_seqlen_k=max_seq_len,
+            num_heads_q=self.num_heads_q * self.dcp_world_size,
+            num_heads_kv=self.num_heads_kv,
+            headdim=self.headdim,
+            cache_seqlens=seqlens,
+            qkv_dtype=qkv_dtype,
+            cu_seqlens_q=cu_query_lens,
+            page_size=self.block_size,
+            causal=causal,
+            window_size=_maybe_symmetrize_window(self.aot_sliding_window, causal),
+            num_splits=max_num_splits,
+        )
+
+    def _store_scheduler_metadata(
+        self, scheduler_metadata: torch.Tensor | None
+    ) -> torch.Tensor | None:
+        if self.use_full_cuda_graph and scheduler_metadata is not None:
+            n = scheduler_metadata.shape[0]
+            assert self.scheduler_metadata is not None
+            self.scheduler_metadata[:n] = scheduler_metadata
+            # NOTE(woosuk): We should zero out the rest of the scheduler
+            # metadata to guarantee the correctness. Otherwise, some thread
+            # blocks may use the invalid scheduler metadata and overwrite the
+            # output buffer.
+            self.scheduler_metadata[n:] = 0
+            return self.scheduler_metadata[:n]
+        return scheduler_metadata
+
     def __init__(
         self,
         kv_cache_spec: AttentionSpec,
@@ -405,6 +456,14 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
             self.dcp_world_size = 1
             self.dcp_rank = 0
 
+        # Fused draft decode reuses the captured metadata object across draft
+        # steps. For DCP, build-time host-side decisions such as
+        # skip_dcp_context_attention() can change the metadata shape/control
+        # path (for example max_dcp_context_kv_len), and those Python-side
+        # fields are not refreshed in-place between graph replays. Keep the
+        # fused path disabled until DCP gets a full replay-safe refresh model.
+        self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
+
         self.cp_kv_cache_interleave_size = (
             self.parallel_config.cp_kv_cache_interleave_size
         )
@@ -533,34 +592,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         if envs.VLLM_BATCH_INVARIANT:
             max_num_splits = 1
 
-        def schedule(
-            batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
-        ):
-            cache_dtype = self.cache_config.cache_dtype
-            if is_quantized_kv_cache(cache_dtype):
-                qkv_dtype = current_platform.fp8_dtype()
-            else:
-                qkv_dtype = self.kv_cache_dtype
-            if aot_schedule:
-                return get_scheduler_metadata(
-                    batch_size=batch_size,
-                    max_seqlen_q=max_query_len,
-                    max_seqlen_k=max_seq_len,
-                    num_heads_q=self.num_heads_q * self.dcp_world_size,
-                    num_heads_kv=self.num_heads_kv,
-                    headdim=self.headdim,
-                    cache_seqlens=seqlens,
-                    qkv_dtype=qkv_dtype,
-                    cu_seqlens_q=cu_query_lens,
-                    page_size=self.block_size,
-                    causal=causal,
-                    window_size=_maybe_symmetrize_window(
-                        self.aot_sliding_window, causal
-                    ),
-                    num_splits=max_num_splits,
-                )
-            return None
-
         use_cascade = common_prefix_len > 0
         max_dcp_context_kv_len = 0
         dcp_context_kv_lens = None
@@ -627,13 +658,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                     (max_seq_len + num_partitions - 1) // num_partitions
                 ) * self.cp_kv_cache_interleave_size
 
-                scheduler_metadata = schedule(
+                scheduler_metadata = self._get_scheduler_metadata(
+                    aot_schedule=aot_schedule,
                     batch_size=num_reqs,
                     cu_query_lens=query_start_loc,
                     max_query_len=max_query_len,
                     seqlens=dcp_context_kv_lens,
                     max_seq_len=max_dcp_context_kv_len,
                     causal=False,
+                    max_num_splits=max_num_splits,
                 )
         elif use_cascade:
             cu_prefix_query_lens = torch.tensor(
@@ -644,41 +677,38 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
             )
             # Use GPU tensor directly - no CPU sync needed
             suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len
-            prefix_scheduler_metadata = schedule(
+            prefix_scheduler_metadata = self._get_scheduler_metadata(
+                aot_schedule=aot_schedule,
                 batch_size=1,
                 cu_query_lens=cu_prefix_query_lens,
                 max_query_len=num_actual_tokens,
                 seqlens=prefix_kv_lens,
                 max_seq_len=common_prefix_len,
                 causal=False,
+                max_num_splits=max_num_splits,
             )
-            scheduler_metadata = schedule(
+            scheduler_metadata = self._get_scheduler_metadata(
+                aot_schedule=aot_schedule,
                 batch_size=num_reqs,
                 cu_query_lens=query_start_loc,
                 max_query_len=max_query_len,
                 seqlens=suffix_kv_lens,
                 max_seq_len=max_seq_len - common_prefix_len,
                 causal=True,
+                max_num_splits=max_num_splits,
             )
         else:
-            scheduler_metadata = schedule(
+            scheduler_metadata = self._get_scheduler_metadata(
+                aot_schedule=aot_schedule,
                 batch_size=num_reqs,
                 cu_query_lens=query_start_loc,
                 max_query_len=max_query_len,
                 seqlens=seq_lens,
                 max_seq_len=max_seq_len,
                 causal=causal,
+                max_num_splits=max_num_splits,
             )
-        # For FA3 + full cudagraph
-        if self.use_full_cuda_graph and scheduler_metadata is not None:
-            n = scheduler_metadata.shape[0]
-            self.scheduler_metadata[:n] = scheduler_metadata
-            # NOTE(woosuk): We should zero out the rest of the scheduler
-            # metadata to guarantee the correctness. Otherwise, some thread
-            # blocks may use the invalid scheduler metadata and overwrite the
-            # output buffer.
-            self.scheduler_metadata[n:] = 0
-            scheduler_metadata = self.scheduler_metadata[:n]
+        scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
 
         if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
             causal = causal.to(torch.int32)
@@ -772,6 +802,28 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         new_metadata.slot_mapping = slot_mapping
         return new_metadata
 
+    def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None:
+        if metadata.scheduler_metadata is None:
+            return
+
+        num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0]
+
+        assert self.dcp_world_size == 1
+        assert not metadata.use_cascade
+
+        scheduler_metadata = self._get_scheduler_metadata(
+            aot_schedule=True,
+            batch_size=num_reqs,
+            cu_query_lens=metadata.query_start_loc,
+            max_query_len=metadata.max_query_len,
+            seqlens=metadata.seq_lens,
+            max_seq_len=metadata.max_seq_len,
+            causal=metadata.causal,
+            max_num_splits=metadata.max_num_splits,
+        )
+
+        metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
+
     def use_cascade_attention(self, *args, **kwargs) -> bool:
         return use_cascade_attention(*args, **kwargs)
 

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

@@ -400,6 +400,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
 
     reorder_batch_threshold: int | None = None
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    supports_draft_decode_metadata_update = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -649,6 +650,40 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             **deepseek_v4_fields,  # type: ignore[arg-type]
         )
 
+    def update_draft_decode_metadata(
+        self,
+        metadata: DeepseekSparseSWAMetadata,
+    ) -> None:
+        if metadata.num_decode_tokens == 0:
+            return
+        assert metadata.query_start_loc is not None
+        assert metadata.seq_lens is not None
+        assert metadata.token_to_req_indices is not None
+        assert metadata.is_valid_token is not None
+        assert metadata.decode_swa_indices is not None
+        assert metadata.decode_swa_lens is not None
+
+        _compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](
+            metadata.decode_swa_indices,
+            metadata.decode_swa_indices.stride(0),
+            metadata.decode_swa_lens,
+            metadata.decode_swa_indices.shape[-1],
+            metadata.query_start_loc,
+            metadata.seq_lens,
+            metadata.token_to_req_indices,
+            metadata.is_valid_token,
+            metadata.block_table,
+            metadata.block_table.stride(0),
+            self.block_size,
+            token_offset=0,
+            TRITON_BLOCK_SIZE=1024,
+        )
+        tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
+        metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
+        metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
+        metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
+        metadata.flashinfer_sparse_index_cache.clear()
+
     def build_tile_scheduler(
         self, num_decode_tokens: int
     ) -> dict[str, FlashMLASchedMeta | None]:

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

@@ -57,6 +57,8 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
 
     def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        # DCP local sequence lengths are not advanced between draft steps.
+        self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
         # Only the non-causal DSpark draft group serves multi-token blocks via
         # the decode path; raise its reorder threshold to the spec block length
         # so full-cudagraph capture admits it. Causal usage stays single-token.
@@ -64,6 +66,9 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
             self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
         self._reserve_attn_logits_workspace()
 
+    def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None:
+        pass
+
     def _reserve_attn_logits_workspace(self) -> None:
         """Pre-size the shared workspace for the decode split-KV attn logits.
 

vllm/v1/attention/backends/triton_attn.py

@@ -100,6 +100,8 @@ class TritonAttentionMetadata:
 
 class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+    # Step-dependent fields reference persistent input buffers directly.
+    supports_draft_decode_metadata_update = True
 
     def __init__(
         self,
@@ -267,6 +269,9 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
 
         return attn_metadata
 
+    def update_draft_decode_metadata(self, _metadata: TritonAttentionMetadata) -> None:
+        pass
+
 
 class TritonAttentionBackend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -10,17 +10,21 @@ from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
+from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     get_uniform_token_count,
 )
 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
     SpeculatorCudaGraphManager,
 )
 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.utils import AttentionGroup
 
 logger = init_logger(__name__)
 
@@ -41,6 +45,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
 
         self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None
         self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None
+        self.use_fused_multi_step_decode = False
 
     def load_model(self, target_model: nn.Module) -> None:
         super().load_model(target_model)
@@ -76,6 +81,49 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         """
         return True
 
+    def set_attn(
+        self,
+        model_state: ModelState,
+        kv_cache_config: KVCacheConfig,
+        block_tables: BlockTables,
+        target_input_buffers: InputBuffers,
+        target_attn_groups: list[list[AttentionGroup]],
+    ) -> None:
+        super().set_attn(
+            model_state,
+            kv_cache_config,
+            block_tables,
+            target_input_buffers,
+            target_attn_groups,
+        )
+        self._configure_fused_multi_step_decode()
+
+    def _configure_fused_multi_step_decode(self) -> None:
+        if self.num_speculative_steps == 1:
+            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()
+                for attn_groups in self.attn_groups
+                for attn_group in attn_groups
+                if not attn_group.supports_draft_decode_metadata_update
+            }
+        )
+        self.use_fused_multi_step_decode = not unsupported_backends
+        if unsupported_backends:
+            logger.info_once(
+                "Fused multi-step draft decode is not supported by attention "
+                "backend(s) %s; falling back to rebuilding attention metadata "
+                "between draft steps.",
+                ", ".join(unsupported_backends),
+            )
+
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
         # Initialize cudagraph manager for draft prefill (draft position 0).
         self.prefill_cudagraph_manager = SpeculatorCudaGraphManager(
@@ -104,6 +152,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         # Reset indices to zeros to prevent stale values from prior
         # dummy runs to cause out-of-bounds indexing during capture.
         self.last_token_indices.zero_()
+        self.idx_mapping.zero_()
 
         # Capture the prefill routine (model forward + compute_logits +
         # sample).
@@ -133,12 +182,15 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             return
 
         self.on_multi_step_decode_begin(self.max_num_reqs)
-        # Capture the decode draft generation routine (model forward +
-        # sample + update_draft_inputs) for a single
-        # step.
+        # Capture either the fused decode loop or one decode step per graph.
         assert self.decode_cudagraph_manager is not None
+        decode_fn = (
+            self._generate_fused_drafts
+            if self.use_fused_multi_step_decode
+            else self._generate_draft
+        )
         self.decode_cudagraph_manager.capture(
-            self._generate_draft,
+            decode_fn,
             self.model_state,
             self.input_buffers,
             self.block_tables,
@@ -290,7 +342,12 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
 
         self.on_multi_step_decode_begin(num_reqs)
         # Generate the remaining num_speculative_steps - 1 draft tokens.
-        self._multi_step_decode(
+        decode_fn = (
+            self._fused_multi_step_decode
+            if self.use_fused_multi_step_decode
+            else self._multi_step_decode
+        )
+        decode_fn(
             num_reqs,
             dummy_run and skip_attn_for_dummy_run,
             decode_batch_desc,
@@ -435,10 +492,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
                     step=step,
                 )
 
-            # Update the current draft step.
             self.current_draft_step.fill_(step)
 
-            # Generate draft tokens for the current step.
             if batch_desc.cg_mode == CUDAGraphMode.FULL:
                 assert self.decode_cudagraph_manager is not None
                 self.decode_cudagraph_manager.run_fullgraph(batch_desc)
@@ -452,6 +507,95 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
                     cudagraph_runtime_mode=batch_desc.cg_mode,
                 )
 
+    def _fused_multi_step_decode(
+        self,
+        num_reqs: int,
+        skip_attn: bool,
+        batch_desc: BatchExecutionDescriptor,
+        num_tokens_across_dp: torch.Tensor | None,
+        seq_lens_cpu_upper_bound: torch.Tensor,
+    ) -> None:
+        positions = self.input_buffers.positions[:num_reqs]
+        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+        idx_mapping = self.idx_mapping[:num_reqs]
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
+        if not skip_attn:
+            slot_mappings = self.block_tables.compute_slot_mappings(
+                idx_mapping,
+                query_start_loc,
+                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
+                )
+            attn_metadata = self._build_draft_attn_metadata(
+                num_reqs=num_reqs,
+                num_reqs_padded=batch_desc.num_reqs or num_reqs,
+                num_tokens_padded=batch_desc.num_tokens,
+                seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+                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,
+            attn_metadata,
+            slot_mappings_by_layer,
+            num_tokens_across_dp,
+            batch_desc.cg_mode,
+        )
+
+    def _generate_fused_drafts(
+        self,
+        num_reqs: int,
+        num_tokens_padded: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> None:
+        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, self.num_speculative_steps):
+            self.current_draft_step.fill_(step)
+            self._generate_draft(
+                num_reqs,
+                num_tokens_padded,
+                attn_metadata,
+                slot_mappings,
+                num_tokens_across_dp,
+                cudagraph_runtime_mode,
+            )
+            if (
+                step < self.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,
         num_reqs: int,

vllm/v1/worker/utils.py

@@ -299,6 +299,17 @@ class AttentionGroup:
         assert len(self.metadata_builders) > ubatch_id
         return self.metadata_builders[ubatch_id]
 
+    @property
+    def supports_draft_decode_metadata_update(self) -> bool:
+        return self.get_metadata_builder().supports_draft_decode_metadata_update
+
+    def update_draft_decode_metadata(
+        self,
+        attn_metadata: Mapping[str, Any],
+    ) -> None:
+        metadata = attn_metadata[self.layer_names[0]]
+        self.get_metadata_builder().update_draft_decode_metadata(metadata)
+
 
 def select_common_block_size(
     kv_manager_block_size: int,