[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 <i@bugenzhao.com>
Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <mozf@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

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,

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]

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(

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"
 

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())
 

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
 

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

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

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

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

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.

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,

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

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

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,

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,

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)