From a311916a291c1fed3dbfb72e60f74cd778c8419d Mon Sep 17 00:00:00 2001 From: Yizhou <136800916+yiz-liu@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:06:22 +0800 Subject: [PATCH] [MRV2][Spec] Fuse AR speculator multi-step decodes back into one CUDA graph (#46849) Signed-off-by: Yizhou Liu --- docs/design/model_runner_v2.md | 8 + .../test_gpu_autoregressive_speculator.py | 166 ++++++++++++++++++ vllm/models/deepseek_v4/amd/rocm.py | 4 + vllm/v1/attention/backend.py | 13 ++ vllm/v1/attention/backends/flash_attn.py | 136 +++++++++----- vllm/v1/attention/backends/mla/sparse_swa.py | 35 ++++ vllm/v1/attention/backends/mla/triton_mla.py | 5 + vllm/v1/attention/backends/triton_attn.py | 5 + .../spec_decode/autoregressive/speculator.py | 158 ++++++++++++++++- vllm/v1/worker/utils.py | 11 ++ 10 files changed, 492 insertions(+), 49 deletions(-) diff --git a/docs/design/model_runner_v2.md b/docs/design/model_runner_v2.md index fb40d51ee7..cabc3526b1 100644 --- a/docs/design/model_runner_v2.md +++ b/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. diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 9f5f90a977..0d8c2943b0 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/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 diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 7c232fe9e5..9b38409d1f 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d4698520f1..32e932e3b8 100644 --- a/vllm/v1/attention/backend.py +++ b/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, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 97f30736fd..acdb1661dc 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/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) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c823f82be2..c11698335a 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/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]: diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 71a38e95d0..e55076700f 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/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. diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 3f9b9ed0c1..7814a8c0f8 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/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]] = [ diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 5ceb6c7558..7d357533e2 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/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, diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 9d056cc67c..55f93860ad 100644 --- a/vllm/v1/worker/utils.py +++ b/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,