From 166a8e954b6e2d90fa34ac5d167aaacef720dd80 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 16 Jul 2026 06:46:59 +0000 Subject: [PATCH] [Model] Add Inkling multi-depth MTP support [5/N] Extend the merged Inkling MTP=1 implementation to multiple checkpoint depths with multi-step speculative decoding and KV-cache plumbing. Co-authored-by: Bugen Zhao Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Jee Jee Li Co-authored-by: Roger Wang Co-authored-by: Yifan Qiao Co-authored-by: Claude Fable 5 Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- rust/src/chat/src/multimodal.rs | 7 +- .../test_speculative_draft_hf_overrides.py | 4 +- vllm/config/compilation.py | 17 +- vllm/config/speculative.py | 24 +- vllm/models/inkling/nvidia/model.py | 16 +- vllm/models/inkling/nvidia/moe.py | 6 +- vllm/models/inkling/nvidia/mtp.py | 55 +- vllm/utils/multi_stream_utils.py | 4 +- vllm/v1/core/kv_cache_coordinator.py | 23 +- vllm/v1/core/kv_cache_manager.py | 2 + vllm/v1/core/kv_cache_utils.py | 20 + vllm/v1/core/sched/scheduler.py | 5 + vllm/v1/core/single_type_kv_cache_manager.py | 15 +- vllm/v1/kv_cache_interface.py | 15 +- .../spec_decode/autoregressive/speculator.py | 528 ++++++++++++++++-- .../gpu/spec_decode/dflash/speculator.py | 6 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 45 +- 17 files changed, 672 insertions(+), 120 deletions(-) diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 9c4f57c0413..7d0045f67c4 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -173,7 +173,12 @@ impl ResolvedPlaceholder { let marker_token_id = context.tokenizer().token_to_id(&token).ok_or_else(|| { multimodal!("placeholder token `{token}` is not in the tokenizer vocabulary") })?; - let embed_token_id = raw.placeholder_token_id_for(&metadata, modality)? as u32; + let embed_token_id = raw.placeholder_token_id_for(&metadata, modality)?; + let embed_token_id = u32::try_from(embed_token_id).map_err(|_| { + multimodal!( + "placeholder token id `{embed_token_id}` for `{modality}` must be non-negative" + ) + })?; Ok(Self { token, diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 7e425d68eec..9652c5ea020 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -86,7 +86,7 @@ def test_arch_mapping_applies_before_callable_override(): @pytest.mark.cpu_test -def test_inkling_override_exposes_only_first_mtp_depth(): +def test_inkling_override_exposes_checkpoint_mtp_depths(): text_config = _make_hf_config( architectures=["InklingForCausalLM"], model_type="inkling_model", @@ -107,7 +107,7 @@ def test_inkling_override_exposes_only_first_mtp_depth(): assert out is text_config assert out.model_type == "inkling_mtp" assert out.architectures == ["InklingMTPModel"] - assert out.n_predict == 1 + assert out.n_predict == 8 assert out.num_nextn_predict_layers == 8 assert out.chain_hidden_post_norm is False assert out.local_layer_ids == [0, 2, 4] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 810c40131fc..ec0326c2858 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1422,16 +1422,15 @@ class CompilationConfig: "and make sure compilation mode is VLLM_COMPILE" ) - # MRV1 adjusts cudagraph sizes to be a multiple of uniform_decode_query_len - # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix: - # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536 - # Will be removed in the near future when we have separate cudagraph capture - # sizes for decode and mixed prefill-decode. - # MRV2 handles cudagraph capture sizing in cudagraph_utils.py - # and doesn't need below: https://github.com/vllm-project/vllm/pull/45953 + # Round cudagraph sizes to a multiple of uniform_decode_query_len so a + # uniform decode batch cannot dispatch to a "mixed" descriptor whose token + # count is not a multiple of the decode span, which corrupts per-request + # conv/attention state (https://github.com/vllm-project/vllm/issues/28207). + # Enabled for MRV2 as well: its cudagraph_utils.py sizing does not on its + # own keep uniform batches off mixed descriptors for multi-module MTP. + # Temporary until decode and mixed prefill-decode have separate sizes. if ( - not use_v2_model_runner - and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and uniform_decode_query_len > 1 ): self.adjust_cudagraph_sizes_for_spec_decode( diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cf6f0a5ae6b..361002e72ff 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -555,6 +555,8 @@ class SpeculativeConfig: ) if hf_config.model_type in ("inkling_mm_model", "inkling_model"): + # MTP weights live on the text backbone. Promote text_config and + # expose the nested MTP fields on the flat draft config. mtp_config = getattr(hf_config, "mtp_config", None) or {} hf_config = getattr(hf_config, "text_config", hf_config) checkpoint_depths = mtp_config.get("num_nextn_predict_layers", 0) @@ -563,12 +565,13 @@ class SpeculativeConfig: hf_config.model_type = "inkling_mtp" hf_config.update( { - # Inkling currently exposes only the first checkpoint depth. - "n_predict": 1, + "n_predict": checkpoint_depths, "num_nextn_predict_layers": checkpoint_depths, "chain_hidden_post_norm": mtp_config.get( "chain_hidden_post_norm", False ), + # The MTP depth blocks carry their own sliding-window + # pattern, which differs from the backbone's layer pattern. "local_layer_ids": mtp_config.get("local_layer_ids", []), "architectures": ["InklingMTPModel"], } @@ -974,14 +977,6 @@ class SpeculativeConfig: "`num_speculative_tokens` was not provided" ) - if ( - self.draft_model_config.hf_config.model_type == "inkling_mtp" - and self.num_speculative_tokens != 1 - ): - raise ValueError( - "Inkling MTP currently supports exactly one speculative token" - ) - if self.method == "dspark": # DSpark is a semi-autoregressive *block* drafter. A # speculative length smaller than the checkpoint's block @@ -1308,6 +1303,15 @@ class SpeculativeConfig: # TODO(ben): Refactor this so the naming is clearer return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark") + def num_speculative_prefill_steps(self) -> int: + if self.method == "mtp" and self.draft_model_config is not None: + n_predict = getattr( + self.draft_model_config.hf_config, "num_nextn_predict_layers", None + ) + if n_predict: + return min(int(n_predict), self.num_speculative_tokens) + return 1 + def use_dflash(self) -> bool: return self.method == "dflash" diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index 39f8ea5375d..def6db90c3c 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/vllm/models/inkling/nvidia/model.py @@ -227,6 +227,7 @@ class InklingDecoderLayer(nn.Module): # Caller folds mlp_output (pre-reduce, pre-sconv) into the next # fused sconv+add+rmsnorm. return hidden_states, (mlp_output, self.mlp_sconv) + # Standalone (MTP) tail: finish the sublayer without a norm. return _sconv_add_norm( mlp_output, hidden_states, self.mlp_sconv, None, positions )[1] @@ -236,10 +237,12 @@ class InklingReplicatedEmbedding(nn.Module): """Full-vocab embedding table replicated on every TP rank. Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a - 1/tp shard) for no masked lookup or per-lookup TP all-reduce, and keeps the - full table on-rank for the fused gather+norm kernel. Bit-exact vs - vocab-parallel: the all-reduce there only ever summed one real row against - exact zeros. The LM head stays vocab-sharded. + 1/tp shard) for no masked lookup and no per-lookup TP all-reduce — one + all-reduce per MTP draft step plus one per verify pass — and keeps the + full table on-rank for the fused gather+norm kernels (``embed_rmsnorm``, + ``embed_dual_rmsnorm_cat``). Bit-exact vs vocab-parallel: the all-reduce + there only ever summed one real row against exact zeros. The LM head + stays vocab-sharded. """ def __init__(self, num_embeddings: int, embedding_dim: int) -> None: @@ -611,7 +614,8 @@ class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal): def get_language_model(self) -> nn.Module: # This class IS the causal LM (the towers are side branches), so the # language model is self — callers expect a module exposing ``.model`` - # and ``.lm_head``. + # / ``.lm_head`` (e.g. the MTP/eagle loader shares embeddings via + # ``get_language_model().model.embed_tokens``). return self @@ -672,8 +676,6 @@ def _load_inkling_weights( yield name, weight - # The release checkpoint also carries auxiliary prediction-head weights; - # they are not part of the causal LM served by this implementation. loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) loaded |= loader.load_weights(_iter_loadable_weights()) diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 3bb281c963b..5189bf0cdc8 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -475,9 +475,9 @@ class InklingMoE(nn.Module): lambda: self.sink_experts(x, gammas), self._sink_events[0], self._sink_events[1], - self._sink_stream - if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - else None, + self._sink_stream, + enable_parallel=num_tokens + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD, ) self._routed_sel = None diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index 868ede3bed6..bc6f4047e3d 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -2,9 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). -Implements the first MTP depth from the reference ``mtp_model.py`` shipped with -the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` -input projection, and a full Inkling transformer block with a dense bf16 MLP. +Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP +depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj`` +(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the +same short convolutions as the backbone; its attention is full or sliding-window +per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared +``chain_norm`` is applied after every depth; its output is both the logits input +and the previous hidden state fed to the next depth. The draft shares the target's token embedding table and LM head (``load_eagle_model`` wires those references) and applies the backbone @@ -50,6 +54,16 @@ def _mtp_depth_from_name(name: str) -> int | None: return int(m.group(1)) if m else None +def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int: + num_layers = min(n_predict, num_spec) if num_spec else n_predict + if num_layers <= 0: + raise ValueError( + "Inkling MTP requires num_nextn_predict_layers and " + "num_speculative_tokens to select at least one depth layer." + ) + return num_layers + + class InklingMTPDepthLayer(nn.Module): """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" @@ -94,14 +108,30 @@ class InklingMultiTokenPredictor(nn.Module): vllm_config.speculative_config.draft_model_config.hf_config ) self.config = config - if vllm_config.speculative_config.num_speculative_tokens != 1: - raise ValueError( - "Inkling MTP currently supports exactly one speculative token" - ) + # The checkpoint ships num_nextn_predict_layers depth blocks, but only + # the first ``num_speculative_tokens`` are exercised (step i uses depth + # i). Build only those to save memory — each depth is a full Inkling block + # with its own (large) full-history sconv caches and KV cache. + n_predict = config.num_nextn_predict_layers + num_spec = vllm_config.speculative_config.num_speculative_tokens + self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec) self.chain_hidden_post_norm = config.chain_hidden_post_norm + + # Depth blocks whose attention is sliding-window (swa_* head config) + # rather than full; keyed by MTP depth via the checkpoint's + # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors + # InklingModel's local_ids split, but over MTP depths, not backbone + # layers. local_ids = set(config.local_layer_ids) + + # Keyed by depth index (str) to mirror the checkpoint layout. self.layers = nn.ModuleDict( - {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + { + str(idx): InklingMTPDepthLayer( + config, f"{prefix}.layers.{idx}", idx in local_ids + ) + for idx in range(self.num_mtp_layers) + } ) self.chain_norm = ( InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -199,9 +229,8 @@ class InklingMultiTokenPredictor(nn.Module): # auto-enumerated as a draft attention layer); its per-token metadata is # built by the speculator's build_attn_metadata and read from the # forward context, so nothing extra is threaded here. - if spec_step_idx != 0: - raise ValueError("Inkling MTP only supports spec_step_idx=0") - layer = self.layers["0"] + depth = spec_step_idx % self.num_mtp_layers + layer = self.layers[str(depth)] combined = self.fused_input_cat( layer, previous_hidden_states, input_ids, inputs_embeds ) @@ -353,8 +382,8 @@ def _load_inkling_mtp_weights( # Only consume the MTP weights; everything else belongs to the target. if ".mtp." not in name: continue - # Only the first checkpoint depth is used for MTP=1. - if depth is not None and depth != 0: + # Skip depth blocks beyond the ones we built (num_speculative_tokens). + if depth is not None and depth >= module.model.num_mtp_layers: continue # model.mtp.chain_norm.weight -> model.chain_norm.weight # model.mtp.layers.{i}.X -> model.layers.{i}.X diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index 2203221c5a1..b90e39b3305 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -23,6 +23,7 @@ def maybe_execute_in_parallel( event0: torch.cuda.Event, event1: torch.cuda.Event, aux_stream: torch.cuda.Stream | None = None, + enable_parallel: bool = True, ) -> tuple[Any, Any]: """Run two functions potentially in parallel on separate CUDA streams. @@ -40,11 +41,12 @@ def maybe_execute_in_parallel( event1: CUDA event recorded after fn1 so default stream can wait. aux_stream: The second CUDA stream for fn1. Multi-stream is disabled when aux_stream is None. + enable_parallel: Opt-in switch for the multi-stream path. Defaults to True. Returns: Tuple of (fn0_result, fn1_result). """ - if aux_stream is not None: + if aux_stream is not None and enable_parallel: event0.record() result0 = fn0() with torch.cuda.stream(aux_stream): diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 7d482192d54..ae6a08fbac5 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -86,6 +86,7 @@ class KVCacheCoordinator(ABC): for g in kv_cache_config.kv_cache_groups ) self.scheduler_block_size = scheduler_block_size + self.num_spec_prefill_steps = 1 self.block_pool = BlockPool( num_gpu_blocks=kv_cache_config.num_blocks, @@ -281,9 +282,14 @@ class KVCacheCoordinator(ABC): (including tokens that are already cached). """ for manager in self.single_type_managers: + # Only cache tokens with finalized KV. The last num_spec_prefill_steps - 1 + # tokens can be re-prefilled by speculative modules. + num_tokens_to_cache = max( + 0, num_computed_tokens - (self.num_spec_prefill_steps - 1) + ) manager.cache_blocks( request, - num_computed_tokens, + num_tokens_to_cache, retention_interval=self.retention_interval, ) @@ -661,9 +667,20 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): # EAGLE groups match one block past each aligned boundary and drop # it, so make that lookahead block eligible to be cached. if manager.use_eagle and aligned_num_computed_tokens > 0: + # Only cache tokens with finalized KV. The last + # num_spec_prefill_steps - 1 tokens can be re-prefilled by + # speculative modules. + num_finalized_computed_tokens = max( + 0, num_computed_tokens - (self.num_spec_prefill_steps - 1) + ) + aligned_num_finalized_computed_tokens = ( + num_finalized_computed_tokens + // self.scheduler_block_size + * self.scheduler_block_size + ) num_tokens_to_cache = min( - num_computed_tokens, - aligned_num_computed_tokens + manager.block_size, + num_finalized_computed_tokens, + aligned_num_finalized_computed_tokens + manager.block_size, ) # The manager already knows the fine hit granularity # (``scheduler_block_size``); retention is passed separately so it diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 3d3d8c1573a..78d04a68608 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -121,6 +121,7 @@ class KVCacheManager: max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, + num_spec_prefill_steps: int = 1, log_stats: bool = False, enable_kv_cache_events: bool = False, dcp_world_size: int = 1, @@ -158,6 +159,7 @@ class KVCacheManager: hash_block_size=hash_block_size, metrics_collector=self.metrics_collector, ) + self.coordinator.num_spec_prefill_steps = num_spec_prefill_steps self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.kv_cache_config = kv_cache_config diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index dc5ecf3f4d2..d0949ec88f5 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2086,6 +2086,26 @@ def get_kv_cache_configs( # Check if the KV cache specs are registered correctly. # This is to prevent that some layers are initialized with unregistered specs. KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + # Under multi-module MTP with prefix caching, tag every sliding-window spec + # with the store-side retention lag so pool sizing, the runtime admission + # cap, and block eviction all agree that these groups retain a few extra + # trailing blocks (single source of truth: the spec field). No-op (lag 0) + # otherwise, so non-MTP models are unchanged. Done before grouping so the lag + # flows through spec unification (which preserves it via ``replace``). + + # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # tag every SlidingWindowSpec with how many extra tokens to retain in the window. + extra_retained_tokens = ( + vllm_config.speculative_config.num_speculative_prefill_steps() - 1 + if vllm_config.speculative_config is not None + else 0 + ) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + # Get global KV cache groups. This also handles spec unification for # hybrid models when disable_hybrid_kv_cache_manager is enabled. # After this call, merged_kv_cache_specs may be modified in-place. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index d6f3c3ad0e3..0f334e08540 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -233,6 +233,7 @@ class Scheduler(SchedulerInterface): speculative_config = vllm_config.speculative_config self.use_eagle = False self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_spec_prefill_steps = 1 self.num_lookahead_tokens = 0 self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: @@ -245,6 +246,9 @@ class Scheduler(SchedulerInterface): if speculative_config.use_eagle(): self.use_eagle = True self.num_lookahead_tokens = self.num_spec_tokens + self.num_spec_prefill_steps = ( + speculative_config.num_speculative_prefill_steps() + ) if speculative_config.uses_draft_model(): self.num_lookahead_tokens = self.num_spec_tokens if speculative_config.use_dflash(): @@ -268,6 +272,7 @@ class Scheduler(SchedulerInterface): max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, + num_spec_prefill_steps=self.num_spec_prefill_steps, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 24404839094..51d8439a453 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -853,6 +853,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) self.sliding_window = kv_cache_spec.sliding_window + # Extra trailing tokens to retain below the window (never attended) so a + # multi-module MTP store-side lag can still reconstruct the window from + # cached blocks. + self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens @classmethod def _contiguous_blocks_for_hit( @@ -1048,13 +1052,22 @@ class SlidingWindowManager(SingleTypeKVCacheManager): attention computation since they are outside the sliding window. Thus, get_num_skipped_tokens(7) == 4. + The trailing edge of the window is extended by ``extra_retained_tokens`` + so that those extra trailing tokens' blocks are retained (but not + attended). This is needed for multi-module spec decoding which can + re-prefill the last num_spec_prefill_tokens - 1 tokens from the end + of the sequence, and thus needs to delay freeing/caching of blocks. + Args: num_computed_tokens: The number of tokens that have been computed. Returns: The number of tokens that will be skipped for attention computation. """ - return max(0, num_computed_tokens - self.sliding_window + 1) + return max( + 0, + num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens, + ) def get_num_common_prefix_blocks(self, running_request_id: str) -> int: """ diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d7f195e35b8..11c8d32fa48 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -547,6 +547,12 @@ class ChunkedLocalAttentionSpec(AttentionSpec): class SlidingWindowSpec(AttentionSpec): sliding_window: int head_size_v: int = None # type: ignore[assignment] + # The trailing edge of the window is extended by ``extra_retained_tokens`` + # so that those extra trailing tokens' blocks are retained (but not + # attended). This is needed for multi-module spec decoding which can + # re-prefill the last num_spec_prefill_tokens - 1 tokens from the end + # of the sequence, and thus needs to delay freeing/caching of blocks. + extra_retained_tokens: int = 0 def __post_init__(self): if self.head_size_v is None: @@ -588,8 +594,13 @@ class SlidingWindowSpec(AttentionSpec): """ # During chunked prefill, we hold KV for the last `sliding_window-1` # computed tokens plus the in-flight tokens (frees happen on the - # processed-token basis); never more than `max_model_len`. - num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) + # processed-token basis); never more than `max_model_len`. An additional + # `extra_retained_tokens` trailing tokens are kept alive below the + # window for multi-module spec decoding, and must be accounted here too. + num_tokens = min( + self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens, + max_model_len, + ) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 3cb478348a9..84410c94c48 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -10,6 +10,7 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -40,6 +41,17 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs( self.draft_model_config ) + # HACK: the Inkling MTP draft has no MM processor of its own (its draft + # config is flattened text-only), but it consumes the target's merged + # embeddings at draft prefill — treat it as MM-capable whenever the + # target is. + if ( + not self.supports_mm_inputs + and self.draft_model_config.hf_config.model_type == "inkling_mtp" + ): + self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs( + vllm_config.model_config + ) if self.supports_mm_inputs: self.inputs_embeds = torch.zeros( self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device @@ -48,6 +60,23 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None + self.num_spec_prefill_steps = ( + self.speculative_config.num_speculative_prefill_steps() + ) + self.cached_draft_input_ids = torch.zeros( + self.max_num_reqs, + self.num_spec_prefill_steps - 1, + dtype=torch.int64, + device=self.device, + ) + self.cached_target_hidden_states = torch.zeros( + self.max_num_reqs, + self.num_spec_prefill_steps - 1, + self.hidden_size, + dtype=self.dtype, + device=self.device, + ) + @property def advance_draft_positions(self) -> bool: """ @@ -92,23 +121,29 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): # For FULL graphs, the entire routine is recorded as one graph. # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. - # Draft prefill reuses the target model's attention metadata at - # runtime, so capture builds its dummy metadata through the target - # model runner's builders and buffers. + # When num_spec_prefill_steps > 1 (e.g. multi-module MTP), the + # speculator builds its attention metadata using its own builders + # and buffers. Otherwise, the target model's attention metadata + # can be reused, so capture builds its dummy metadata through the + # target model runner's builders and buffers. assert self.prefill_cudagraph_manager is not None if self.prefill_cudagraph_manager.use_breakable_cg: self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model) self.prefill_cudagraph_manager.capture( self._prefill, self.model_state, - self.target_input_buffers, + self.input_buffers + if self.num_spec_prefill_steps > 1 + else self.target_input_buffers, self.block_tables, - self.target_attn_groups, + self.attn_groups + if self.num_spec_prefill_steps > 1 + else self.target_attn_groups, self.kv_cache_config, progress_bar_desc="Capturing prefill CUDA graphs", ) - if self.num_speculative_steps == 1: + if self.num_speculative_steps <= self.num_spec_prefill_steps: return # Capture the decode draft generation routine (model forward + @@ -116,7 +151,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): # step. assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.capture( - self._generate_draft, + self._decode, self.model_state, self.input_buffers, self.block_tables, @@ -160,6 +195,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.draft_max_seq_len = min( max_seq_len + self.num_speculative_steps, self.max_model_len ) + skip_attn = dummy_run and skip_attn_for_dummy_run # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of input_ids and @@ -174,7 +210,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens].copy_(hidden_states) self._copy_request_inputs( num_reqs, @@ -183,16 +218,21 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): seeds, ) - # Get the input ids and last token indices for the speculator. + # Get the input ids, last token indices, and input hidden states for + # the speculator. prepare_prefill_inputs( self.last_token_indices, - self.current_draft_step, + self.hidden_states, self.input_buffers, + hidden_states, + self.cached_target_hidden_states, + self.cached_draft_input_ids, input_batch, num_sampled, num_rejected, last_sampled, next_prefill_tokens, + self.num_spec_prefill_steps, self.max_num_reqs, ) @@ -215,6 +255,38 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): need_eager=is_profile, ) + if not skip_attn and self.num_spec_prefill_steps > 1: + # Re-prefill shifts the draft positions/seq_lens, so the target's + # slot mappings and attention metadata can no longer be reused. + # Rebuild both from the draft buffers written by + # prepare_prefill_inputs, keeping the target's (non-uniform) query + # layout since the draft query_start_loc mirrors the target's. + index_mapping = self.idx_mapping[:num_reqs] + last_token_indices = self.last_token_indices[:num_reqs] + slot_mappings_tensor = self.block_tables.compute_slot_mappings( + index_mapping, + self.input_buffers.query_start_loc, + self.input_buffers.positions, + prefill_batch_desc.num_tokens, + ) + # Apply padding values to slots not corresponding to real draft + # tokens to prevent stale value writes. + pad_trailing_draft_slots( + slot_mappings_tensor, + self.input_buffers.query_start_loc, + last_token_indices, + num_reqs, + ) + slot_mappings = build_slot_mappings_by_layer( + slot_mappings_tensor, self.kv_cache_config + ) + attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=prefill_batch_desc.num_reqs or num_reqs, + num_tokens_padded=prefill_batch_desc.num_tokens, + query_start_loc_np=input_batch.query_start_loc_np, + ) + self._prepare_eplb_forward(input_batch.num_tokens) if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: @@ -235,9 +307,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): mm_inputs=mm_inputs, ) - if self.num_speculative_steps == 1: + if self.num_speculative_steps <= self.num_spec_prefill_steps: # Early exit. - return self.draft_tokens[:num_reqs, :1] + return self.draft_tokens[:num_reqs, : self.num_speculative_steps] # Prepare the inputs for the decode steps. prepare_decode_inputs( @@ -265,7 +337,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): # Generate the remaining num_speculative_steps - 1 draft tokens. self._multi_step_decode( num_reqs, - dummy_run and skip_attn_for_dummy_run, + skip_attn, decode_batch_desc, num_tokens_across_dp, ) @@ -281,6 +353,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + embed_mm: bool = False, + spec_module_idx: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: batch_descriptor = BatchDescriptor(num_tokens=num_tokens) with set_forward_context( @@ -293,7 +367,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): batch_descriptor=batch_descriptor, ): inputs_embeds = None - if self.supports_mm_inputs: + # Only draft prefill consumes (target-merged) MM embeddings; the + # decode steps always draft text tokens and keep the raw-ids path + # (per-callsite, not data-dependent, so each cudagraph family + # captures the branch it will replay). + if self.supports_mm_inputs and embed_mm: # Merge multimodal embeddings with input ids. mm_embeds, is_mm_embed = mm_inputs or (None, None) num_input_tokens = ( @@ -312,6 +390,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) + if spec_module_idx > 0: + # Pass the speculative module index to the model to indicate which + # module to use for the forward pass. + model_inputs["spec_step_idx"] = spec_module_idx if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable), # chosen inside run_pw_graph. @@ -342,33 +424,67 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, ) -> None: last_token_indices = self.last_token_indices[:num_reqs] - positions = self.input_buffers.positions[last_token_indices] + sample_positions = self.input_buffers.positions[last_token_indices] idx_mapping = self.idx_mapping[:num_reqs] - last_hidden_states, hidden_states = self._run_model( - num_tokens, - attn_metadata, - slot_mappings, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - mm_inputs=mm_inputs, - ) - sample_hidden_states = last_hidden_states[last_token_indices] + if self.num_spec_prefill_steps > 1: + # Snapshot the trailing draft tokens/hidden states (for the next + # prefill's re-prefill gap) before the step loop overwrites them. + cache_prefill_state( + self.input_buffers, + self.hidden_states, + self.cached_draft_input_ids, + self.cached_target_hidden_states, + last_token_indices, + idx_mapping, + num_reqs, + self.num_spec_prefill_steps, + ) + + for step in range(self.num_spec_prefill_steps): + # Update the current draft step. + self.current_draft_step.fill_(step) + + # Run the model forward pass. + last_hidden_states, hidden_states = self._run_model( + num_tokens, + attn_metadata, + slot_mappings, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + mm_inputs=mm_inputs, + embed_mm=True, + spec_module_idx=step, + ) + + # Sample draft tokens for the current step. + sample_hidden_states = last_hidden_states[last_token_indices] + draft_tokens = self.sample_draft( + sample_hidden_states, + sample_positions, + idx_mapping, + self.temperature, + self.seeds, + self.current_draft_step, + self.draft_logits, + ) + + self.draft_tokens[:num_reqs, step] = draft_tokens + if step < self.num_spec_prefill_steps - 1: + self.hidden_states[:num_tokens] = hidden_states + update_prefill_inputs( + draft_tokens, + self.input_buffers, + last_token_indices, + num_reqs, + ) + sample_positions += 1 - self.draft_tokens[:num_reqs, 0] = self.sample_draft( - sample_hidden_states, - positions, - idx_mapping, - self.temperature, - self.seeds, - self.current_draft_step, - self.draft_logits, - ) if last_hidden_states is hidden_states: self.hidden_states[:num_reqs] = sample_hidden_states else: self.hidden_states[:num_reqs] = hidden_states[last_token_indices] - self.input_buffers.positions[:num_reqs] = positions + self.input_buffers.positions[:num_reqs] = sample_positions def _multi_step_decode( self, @@ -383,7 +499,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): attn_metadata = None slot_mappings_by_layer = None - for step in range(1, self.num_speculative_steps): + for step in range(self.num_spec_prefill_steps, self.num_speculative_steps): # Rebuild every step when positions advance, or just once # on the first step when positions are constant (Gemma4 MTP). if not skip_attn and (self.advance_draft_positions or step == 1): @@ -396,7 +512,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, self.kv_cache_config ) - attn_metadata = self._build_draft_attn_metadata( + attn_metadata = self._build_uniform_draft_attn_metadata( num_reqs=num_reqs, num_reqs_padded=batch_desc.num_reqs or num_reqs, num_tokens_padded=batch_desc.num_tokens, @@ -410,7 +526,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.run_fullgraph(batch_desc) else: - self._generate_draft( + self._decode( num_reqs, batch_desc.num_tokens, attn_metadata, @@ -419,7 +535,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): cudagraph_runtime_mode=batch_desc.cg_mode, ) - def _generate_draft( + def _decode( self, num_reqs: int, num_tokens_padded: int, @@ -471,13 +587,14 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): @triton.jit def _prepare_prefill_inputs_kernel( last_token_indices_ptr, - draft_current_step_ptr, draft_input_ids_ptr, draft_positions_ptr, draft_query_start_loc_ptr, draft_seq_lens_ptr, target_input_ids_ptr, target_positions_ptr, + cached_draft_input_ids_ptr, + cached_draft_input_ids_stride0, idx_mapping_ptr, last_sampled_ptr, next_prefill_tokens_ptr, @@ -486,6 +603,7 @@ def _prepare_prefill_inputs_kernel( query_start_loc_ptr, seq_lens_ptr, max_num_reqs, + num_spec_prefill_steps, BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0) @@ -500,6 +618,7 @@ def _prepare_prefill_inputs_kernel( # Get the true query length and next token after accounting for rejected tokens. num_rejected = tl.load(num_rejected_ptr + req_idx) query_len -= num_rejected + num_reprefills = 0 if num_spec_prefill_steps == 1 else max(0, num_rejected - 1) num_sampled = tl.load(num_sampled_ptr + req_idx) if num_sampled > 0: @@ -509,31 +628,57 @@ def _prepare_prefill_inputs_kernel( # Get the next prefill token. next_token = tl.load(next_prefill_tokens_ptr + req_state_idx) - # Shift target_input_ids by one. + # Shift target_input_ids by one + the number of tokens to be re-prefilled. for i in range(1, query_len, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < query_len input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask) - tl.store(draft_input_ids_ptr + query_start + block - 1, input_ids, mask=mask) + tl.store( + draft_input_ids_ptr + query_start + block - 1 + num_reprefills, + input_ids, + mask=mask, + ) - last_token_index = query_start + query_len - 1 + last_token_index = query_start + query_len - 1 + num_reprefills tl.store(last_token_indices_ptr + req_idx, last_token_index) tl.store(draft_input_ids_ptr + last_token_index, next_token) - # Copy positions. + # Copy positions, shifted over by the number of tokens to be re-prefilled. for i in range(0, query_len, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < query_len target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask) - tl.store(draft_positions_ptr + query_start + block, target_pos, mask=mask) + tl.store( + draft_positions_ptr + query_start + block + num_reprefills, + target_pos, + mask=mask, + ) + + # Fill the re-prefill gap with the cached token ids from the previous + # decode step. These tokens sit immediately before the query's first + # token, so their positions are contiguous and derived here. + first_position = tl.load(target_positions_ptr + query_start) + for i in range(num_reprefills): + cache_read_slot = num_spec_prefill_steps - 1 - num_reprefills + i + cached_token_id = tl.load( + cached_draft_input_ids_ptr + + req_state_idx * cached_draft_input_ids_stride0 + + cache_read_slot + ) + tl.store(draft_input_ids_ptr + query_start + i, cached_token_id) + tl.store( + draft_positions_ptr + query_start + i, + first_position - num_reprefills + i, + ) # Copy query start locations. tl.store(draft_query_start_loc_ptr + req_idx, query_start) - # Copy sequence lengths. - tl.store(draft_seq_lens_ptr + req_idx, seq_len) + # Copy sequence lengths. Re-prefilled tokens are packed into the query + # window without widening it, so the effective KV length shrinks by + # num_reprefills to keep RoPE positions aligned with the attention's + # implicit (seq_len - query_len) causal positions. + tl.store(draft_seq_lens_ptr + req_idx, seq_len - num_reprefills) if req_idx == (num_reqs - 1): - # Reset the current draft step to 0. - tl.store(draft_current_step_ptr, 0) # Pad query_start_loc for CUDA graphs. for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) @@ -551,11 +696,126 @@ def _prepare_prefill_inputs_kernel( tl.store(last_token_indices_ptr + block, 0, mask=mask) +@triton.jit +def _prepare_prefill_hidden_states_kernel( + draft_input_hidden_states_ptr, + draft_input_hidden_states_stride0, + target_hidden_states_ptr, + target_hidden_states_stride0, + cached_target_hidden_states_ptr, + cached_target_hidden_states_stride0, + cached_target_hidden_states_stride1, + idx_mapping_ptr, + num_rejected_ptr, + query_start_loc_ptr, + num_spec_prefill_steps, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + block_idx = tl.program_id(1) + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = block < hidden_size + + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + + num_rejected = tl.load(num_rejected_ptr + req_idx) + query_len -= num_rejected + num_reprefills = 0 if num_spec_prefill_steps == 1 else max(0, num_rejected - 1) + + # Fill the re-prefill gap with the cached hidden states from the previous + # decode step, mirroring the token ids and positions inserted by + # _prepare_prefill_inputs_kernel. + for i in range(num_reprefills): + cache_read_slot = num_spec_prefill_steps - 1 - num_reprefills + i + cached_hidden_state = tl.load( + cached_target_hidden_states_ptr + + req_state_idx * cached_target_hidden_states_stride0 + + cache_read_slot * cached_target_hidden_states_stride1 + + block, + mask=mask, + ) + tl.store( + draft_input_hidden_states_ptr + + (query_start + i) * draft_input_hidden_states_stride0 + + block, + cached_hidden_state, + mask=mask, + ) + + for i in range(query_len): + hidden_state = tl.load( + target_hidden_states_ptr + + (query_start + i) * target_hidden_states_stride0 + + block, + mask=mask, + ) + tl.store( + draft_input_hidden_states_ptr + + (query_start + num_reprefills + i) * draft_input_hidden_states_stride0 + + block, + hidden_state, + mask=mask, + ) + + +@triton.jit +def _pad_trailing_draft_slots_kernel( + slot_mappings_ptr, + slot_mappings_stride0, + query_start_loc_ptr, + last_token_indices_ptr, + PAD_ID, + BLOCK_SIZE: tl.constexpr, +): + group_idx = tl.program_id(0) + req_idx = tl.program_id(1) + # Slots computed from stale token positions in the range + # [last_token_index + 1, query_end) can result in writes to blocks. + # Pad these slot values so that attention kernels ignore them. + start = tl.load(last_token_indices_ptr + req_idx) + 1 + end = tl.load(query_start_loc_ptr + req_idx + 1) + base = slot_mappings_ptr + group_idx * slot_mappings_stride0 + for i in range(start, end, BLOCK_SIZE): + offs = i + tl.arange(0, BLOCK_SIZE) + mask = offs < end + tl.store(base + offs, PAD_ID, mask=mask) + + +def pad_trailing_draft_slots( + # [num_groups, num_tokens_padded] + slot_mappings: torch.Tensor, + # [num_reqs + 1] + query_start_loc: torch.Tensor, + # [num_reqs] + last_token_indices: torch.Tensor, + num_reqs: int, +) -> None: + num_groups = slot_mappings.shape[0] + _pad_trailing_draft_slots_kernel[(num_groups, num_reqs)]( + slot_mappings, + slot_mappings.stride(0), + query_start_loc, + last_token_indices, + PAD_SLOT_ID, + BLOCK_SIZE=256, + ) + + def prepare_prefill_inputs( # [num_reqs] last_token_indices: torch.Tensor, - current_draft_step: torch.Tensor, + draft_input_hidden_states: torch.Tensor, input_buffers: InputBuffers, + target_hidden_states: torch.Tensor, + # [max_num_reqs, num_spec_prefill_steps - 1, hidden_size] + cached_target_hidden_states: torch.Tensor | None, + # [max_num_reqs, num_spec_prefill_steps - 1] + cached_draft_input_ids: torch.Tensor | None, input_batch: InputBatch, # [num_reqs] num_sampled: torch.Tensor, @@ -565,18 +825,21 @@ def prepare_prefill_inputs( last_sampled: torch.Tensor, # [max_num_reqs] next_prefill_tokens: torch.Tensor, + num_spec_prefill_steps, max_num_reqs, ) -> torch.Tensor: num_reqs = input_batch.num_reqs + hidden_size = target_hidden_states.shape[-1] _prepare_prefill_inputs_kernel[(num_reqs,)]( last_token_indices, - current_draft_step, input_buffers.input_ids, input_buffers.positions, input_buffers.query_start_loc, input_buffers.seq_lens, input_batch.input_ids, input_batch.positions, + cached_draft_input_ids, + cached_draft_input_ids.stride(0) if cached_draft_input_ids is not None else 0, input_batch.idx_mapping, last_sampled, next_prefill_tokens, @@ -585,11 +848,176 @@ def prepare_prefill_inputs( input_batch.query_start_loc, input_batch.seq_lens, max_num_reqs, + num_spec_prefill_steps, BLOCK_SIZE=1024, ) + if num_spec_prefill_steps > 1: + hidden_block_size = 1024 + num_dim_blocks = triton.cdiv(hidden_size, hidden_block_size) + _prepare_prefill_hidden_states_kernel[(num_reqs, num_dim_blocks)]( + draft_input_hidden_states, + draft_input_hidden_states.stride(0), + target_hidden_states, + target_hidden_states.stride(0), + cached_target_hidden_states, + cached_target_hidden_states.stride(0) + if cached_target_hidden_states is not None + else 0, + cached_target_hidden_states.stride(1) + if cached_target_hidden_states is not None + else 0, + input_batch.idx_mapping, + num_rejected, + input_batch.query_start_loc, + num_spec_prefill_steps, + hidden_size, + BLOCK_SIZE=hidden_block_size, + ) + else: + num_tokens = input_batch.num_tokens + draft_input_hidden_states[:num_tokens].copy_(target_hidden_states[:num_tokens]) return last_token_indices +@triton.jit +def _cache_prefill_state_kernel( + draft_input_ids_ptr, + draft_input_hidden_states_ptr, + draft_input_hidden_states_stride0, + cached_draft_input_ids_ptr, + cached_draft_input_ids_stride0, + cached_target_hidden_states_ptr, + cached_target_hidden_states_stride0, + cached_target_hidden_states_stride1, + idx_mapping_ptr, + last_token_indices_ptr, + query_start_loc_ptr, + num_spec_prefill_steps, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + block_idx = tl.program_id(1) + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = block < hidden_size + + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + if req_state_idx < 0: + # Skip cudagraph padded requests. + return + + query_start = tl.load(query_start_loc_ptr + req_idx) + last_token_index = tl.load(last_token_indices_ptr + req_idx) + + # Snapshot the last num_spec_prefill_steps - 1 input draft tokens/hidden + # states, indexed by request state, so the next prefill can re-prefill + # rejected positions. + cache_window_size = num_spec_prefill_steps - 1 + window_start = last_token_index - cache_window_size + 1 + for i in range(max(window_start, query_start), last_token_index + 1): + cache_write_slot = i - window_start + if block_idx == 0: + input_id = tl.load(draft_input_ids_ptr + i) + tl.store( + cached_draft_input_ids_ptr + + req_state_idx * cached_draft_input_ids_stride0 + + cache_write_slot, + input_id, + ) + hidden_state = tl.load( + draft_input_hidden_states_ptr + + i * draft_input_hidden_states_stride0 + + block, + mask=mask, + ) + tl.store( + cached_target_hidden_states_ptr + + req_state_idx * cached_target_hidden_states_stride0 + + cache_write_slot * cached_target_hidden_states_stride1 + + block, + hidden_state, + mask=mask, + ) + + +def cache_prefill_state( + input_buffers: InputBuffers, + # [num_tokens, hidden_size] + draft_input_hidden_states: torch.Tensor, + # [max_num_reqs, num_spec_prefill_steps - 1] + cached_draft_input_ids: torch.Tensor, + # [max_num_reqs, num_spec_prefill_steps - 1, hidden_size] + cached_target_hidden_states: torch.Tensor, + # [num_reqs] + last_token_indices: torch.Tensor, + # [num_reqs] + idx_mapping: torch.Tensor, + num_reqs: int, + num_spec_prefill_steps: int, +) -> None: + hidden_size = draft_input_hidden_states.shape[-1] + hidden_block_size = 1024 + _cache_prefill_state_kernel[ + (num_reqs, triton.cdiv(hidden_size, hidden_block_size)) + ]( + input_buffers.input_ids, + draft_input_hidden_states, + draft_input_hidden_states.stride(0), + cached_draft_input_ids, + cached_draft_input_ids.stride(0), + cached_target_hidden_states, + cached_target_hidden_states.stride(0), + cached_target_hidden_states.stride(1), + idx_mapping, + last_token_indices, + input_buffers.query_start_loc, + num_spec_prefill_steps, + hidden_size, + BLOCK_SIZE=hidden_block_size, + ) + + +@triton.jit +def _update_prefill_inputs_kernel( + input_ids_ptr, + query_start_loc_ptr, + last_token_indices_ptr, + draft_tokens_ptr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + query_start = tl.load(query_start_loc_ptr + req_idx) + # Use the post-rejection last token index so the shift and insertion align + # with the position the draft token was sampled from. + last_token_index = tl.load(last_token_indices_ptr + req_idx) + query_len = last_token_index - query_start + 1 + + # Shift input token ids to the left by one position and + # insert the last sampled draft token. + for i in range(1, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + input_ids = tl.load(input_ids_ptr + query_start + block, mask=mask) + tl.store(input_ids_ptr + query_start + block - 1, input_ids, mask=mask) + draft_token = tl.load(draft_tokens_ptr + req_idx) + tl.store(input_ids_ptr + last_token_index, draft_token) + + +def update_prefill_inputs( + draft_tokens: torch.Tensor, + input_buffers: InputBuffers, + last_token_indices: torch.Tensor, + num_reqs: int, +) -> None: + _update_prefill_inputs_kernel[(num_reqs,)]( + input_buffers.input_ids, + input_buffers.query_start_loc, + last_token_indices, + draft_tokens, + BLOCK_SIZE=1024, + ) + + @triton.jit def _prepare_decode_inputs_kernel( draft_tokens_ptr, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 3cc88c49b2f..45e83f4b71d 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -264,7 +264,7 @@ class DFlashSpeculator(DraftModelSpeculator): num_reqs, self.num_speculative_steps ) - def _build_draft_attn_metadata( + def _build_uniform_draft_attn_metadata( self, num_reqs: int, num_reqs_padded: int, @@ -275,7 +275,7 @@ class DFlashSpeculator(DraftModelSpeculator): if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead - return super()._build_draft_attn_metadata( + return super()._build_uniform_draft_attn_metadata( num_reqs, num_reqs_padded, num_tokens_padded, @@ -423,7 +423,7 @@ class DFlashSpeculator(DraftModelSpeculator): # Rebuild the draft attention metadata even when replaying the FULL # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_draft_attn_metadata( + draft_attn_metadata = self._build_uniform_draft_attn_metadata( num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, num_tokens_padded=num_tokens_padded, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index f460f58fe09..5c8116ec34b 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from typing import Any +import numpy as np import torch import torch.nn as nn @@ -122,9 +123,7 @@ class DraftModelSpeculator(BaseSpeculator): dtype=torch.int64, device=device, ) - self.arange = torch.arange( - self.max_num_reqs + 1, dtype=torch.int32, device="cpu" - ) + self.np_arange = np.arange(self.max_num_reqs + 1, dtype=np.int32) self.draft_logits: torch.Tensor | None = None if self.speculative_config.draft_sample_method == "probabilistic": @@ -205,7 +204,7 @@ class DraftModelSpeculator(BaseSpeculator): self.target_input_buffers = target_input_buffers self.target_attn_groups = target_attn_groups - def _build_draft_attn_metadata( + def _build_uniform_draft_attn_metadata( self, num_reqs: int, num_reqs_padded: int, @@ -213,13 +212,30 @@ class DraftModelSpeculator(BaseSpeculator): num_query_per_req: int = 1, causal: bool | Mapping[int, bool] = True, ) -> dict[str, Any] | None: - # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req. - # Clamp keeps the series non-decreasing past num_reqs, which some - # attention backends require. - query_start_loc_cpu = ( - torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs) - * num_query_per_req + return self._build_draft_attn_metadata( + num_reqs, + num_reqs_padded, + num_tokens_padded, + self.np_arange[: num_reqs_padded + 1] * num_query_per_req, + causal=causal, ) + + def _build_draft_attn_metadata( + self, + num_reqs: int, + num_reqs_padded: int, + num_tokens_padded: int, + query_start_loc_np: np.ndarray, + causal: bool | Mapping[int, bool] = True, + ) -> dict[str, Any]: + # Fresh tensor: no aliasing with the caller's array, sized to the padded + # request count, with the tail clamped to the last real cumulative value. + query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) + query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( + query_start_loc_np[: num_reqs + 1] + ) + query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] + max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()) block_tables = [ x[:num_reqs_padded] for x in self.block_tables.input_block_tables ] @@ -232,7 +248,7 @@ class DraftModelSpeculator(BaseSpeculator): : num_reqs_padded + 1 ], query_start_loc_cpu=query_start_loc_cpu, - max_query_len=num_query_per_req, + max_query_len=max_query_len, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], max_seq_len=self.draft_max_seq_len, block_tables=block_tables, @@ -313,7 +329,6 @@ class DraftModelSpeculator(BaseSpeculator): self.temperature.copy_(temperature) self.seeds.copy_(seeds) self.idx_mapping[:num_reqs].copy_(idx_mapping) - if self.draft_logits is not None: - # idx_mapping for CG padded requests points to -1, which is ignored - # during sampling to prevent writing stale values to draft logits. - self.idx_mapping[num_reqs:].fill_(-1) + # idx_mapping for CG padded requests points to -1, which is ignored + # during sampling to prevent writing stale values to draft logits. + self.idx_mapping[num_reqs:].fill_(-1)