@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dec13a33b7

[Model Runner V2][Spec Decode] Add multi-layer MTP speculator (#48892)

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

Giancarlo Delfin committed Jul 30, 2026 at 15:54 UTC dec13a33b712e686b1fa26796d51aba4a2e2b3f5
9 files changed +1227 -32
vllm/config/speculative.py
+8
@@ -1410,6 +1410,14 @@ class SpeculativeConfig:
1410 def use_ngram_gpu(self) -> bool:
1411 return self.method == "ngram_gpu"
1412
1413 + def use_multi_module_mtp(self) -> bool:
1414 + if self.method != "mtp" or self.draft_model_config is None:
1415 + return False
1416 + num_mtp_layers = getattr(
1417 + self.draft_model_config.hf_config, "num_nextn_predict_layers", 1
1418 + )
1419 + return min(num_mtp_layers, self.num_speculative_tokens) > 1
1420 +
1421 def __repr__(self) -> str:
1422 method = self.method
1423 model = (
vllm/models/inkling/nvidia/mtp.py
+42 -13
@@ -2,9 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Inkling MTP (Multi-Token Prediction) draft model (NVIDIA).
4
5 -Implements the first MTP depth from the reference ``mtp_model.py`` shipped with
6 -the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H``
7 -input projection, and a full Inkling transformer block with a dense bf16 MLP.
5 +Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP
6 +depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj``
7 +(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the
8 +same short convolutions as the backbone; its attention is full or sliding-window
9 +per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared
10 +``chain_norm`` is applied after every depth; its output is both the logits input
11 +and the previous hidden state fed to the next depth.
12
13 The draft shares the target's token embedding table and LM head
14 (``load_eagle_model`` wires those references) and applies the backbone
@@ -53,6 +57,16 @@ def _mtp_depth_from_name(name: str) -> int | None:
57 return int(m.group(1)) if m else None
58
59
60 +def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int:
61 + num_layers = min(n_predict, num_spec) if num_spec else n_predict
62 + if num_layers <= 0:
63 + raise ValueError(
64 + "Inkling MTP requires num_nextn_predict_layers and "
65 + "num_speculative_tokens to select at least one depth layer."
66 + )
67 + return num_layers
68 +
69 +
70 class InklingMTPDepthLayer(nn.Module):
71 """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block."""
72
@@ -96,14 +110,30 @@ class InklingMultiTokenPredictor(nn.Module):
110 vllm_config.speculative_config.draft_model_config.hf_config
111 )
112 self.config = config
99 - if vllm_config.speculative_config.num_speculative_tokens != 1:
100 - raise ValueError(
101 - "Inkling MTP currently supports exactly one speculative token"
102 - )
113 + # The checkpoint ships num_nextn_predict_layers depth blocks, but only
114 + # the first ``num_speculative_tokens`` are exercised (step i uses depth
115 + # i). Build only those to save memory — each depth is a full Inkling block
116 + # with its own (large) full-history sconv caches and KV cache.
117 + n_predict = config.num_nextn_predict_layers
118 + num_spec = vllm_config.speculative_config.num_speculative_tokens
119 + self.num_mtp_layers = _select_mtp_depth_count(n_predict, num_spec)
120 self.chain_hidden_post_norm = config.chain_hidden_post_norm
121 +
122 + # Depth blocks whose attention is sliding-window (swa_* head config)
123 + # rather than full; keyed by MTP depth via the checkpoint's
124 + # mtp_config.local_layer_ids (promoted onto the draft config). Mirrors
125 + # InklingModel's local_ids split, but over MTP depths, not backbone
126 + # layers.
127 local_ids = set(config.local_layer_ids)
128 +
129 + # Keyed by depth index (str) to mirror the checkpoint layout.
130 self.layers = nn.ModuleDict(
106 - {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)}
131 + {
132 + str(idx): InklingMTPDepthLayer(
133 + config, f"{prefix}.layers.{idx}", idx in local_ids
134 + )
135 + for idx in range(self.num_mtp_layers)
136 + }
137 )
138 self.chain_norm = (
139 InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -201,9 +231,8 @@ class InklingMultiTokenPredictor(nn.Module):
231 # auto-enumerated as a draft attention layer); its per-token metadata is
232 # built by the speculator's build_attn_metadata and read from the
233 # forward context, so nothing extra is threaded here.
204 - if spec_step_idx != 0:
205 - raise ValueError("Inkling MTP only supports spec_step_idx=0")
206 - layer = self.layers["0"]
234 + depth = spec_step_idx % self.num_mtp_layers
235 + layer = self.layers[str(depth)]
236 combined = self.fused_input_cat(
237 layer, previous_hidden_states, input_ids, inputs_embeds
238 )
@@ -355,8 +384,8 @@ def _load_inkling_mtp_weights(
384 # Only consume the MTP weights; everything else belongs to the target.
385 if ".mtp." not in name:
386 continue
358 - # Only the first checkpoint depth is used for MTP=1.
359 - if depth is not None and depth != 0:
387 + # Skip depth blocks beyond the ones we built (num_speculative_tokens).
388 + if depth is not None and depth >= module.model.num_mtp_layers:
389 continue
390 # model.mtp.chain_norm.weight -> model.chain_norm.weight
391 # model.mtp.layers.{i}.X -> model.layers.{i}.X
vllm/v1/worker/gpu/input_batch.py
+21 -4
@@ -191,6 +191,8 @@ class InputBatch:
191 def _prepare_prefill_inputs_kernel(
192 input_ids_ptr,
193 next_prefill_tokens_ptr,
194 + next_prefill_tokens_stride,
195 + num_lookahead,
196 idx_mapping_ptr,
197 query_start_loc_ptr,
198 all_token_ids_ptr,
@@ -198,6 +200,7 @@ def _prepare_prefill_inputs_kernel(
200 prefill_lens_ptr,
201 num_computed_tokens_ptr,
202 BLOCK_SIZE: tl.constexpr,
203 + LOOKAHEAD_BLOCK: tl.constexpr,
204 ):
205 batch_idx = tl.program_id(0)
206 req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
@@ -218,10 +221,20 @@ def _prepare_prefill_inputs_kernel(
221 tokens = tl.load(request_ptr + num_computed + block, mask=mask)
222 tl.store(input_ids_ptr + query_start + block, tokens, mask=mask)
223
221 - next_pos = num_computed + query_len
222 - if next_pos < prefill_len:
223 - next_token = tl.load(request_ptr + next_pos)
224 - tl.store(next_prefill_tokens_ptr + req_state_idx, next_token)
224 + # Store the next num_lookahead prefill tokens.
225 + lookahead = tl.arange(0, LOOKAHEAD_BLOCK)
226 + pos = num_computed + query_len + lookahead
227 + in_lookahead = lookahead < num_lookahead
228 + tokens = tl.load(
229 + request_ptr + pos, mask=in_lookahead & (pos < prefill_len), other=0
230 + )
231 + tl.store(
232 + next_prefill_tokens_ptr
233 + + lookahead * next_prefill_tokens_stride
234 + + req_state_idx,
235 + tokens,
236 + mask=in_lookahead,
237 + )
238
239
240 def prepare_prefill_inputs(
@@ -234,9 +247,12 @@ def prepare_prefill_inputs(
247 num_computed_tokens: torch.Tensor,
248 ) -> None:
249 num_reqs = idx_mapping.shape[0]
250 + num_lookahead = next_prefill_tokens.shape[0]
251 _prepare_prefill_inputs_kernel[(num_reqs,)](
252 input_ids,
253 next_prefill_tokens,
254 + next_prefill_tokens.stride(0),
255 + num_lookahead,
256 idx_mapping,
257 query_start_loc,
258 all_token_ids,
@@ -244,6 +260,7 @@ def prepare_prefill_inputs(
260 prefill_len,
261 num_computed_tokens,
262 BLOCK_SIZE=1024,
263 + LOOKAHEAD_BLOCK=triton.next_power_of_2(num_lookahead),
264 )
265
266
vllm/v1/worker/gpu/model_runner.py
+13 -1
@@ -223,6 +223,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
223 self.is_pooling_model = self.model_config.runner_type == "pooling"
224 self.pooling_runner: PoolingRunner | None = None
225
226 + # Multi-module MTP feeds its modules the next num_speculative_steps prefill
227 + # tokens during chunked prefill. Other speculators only read the immediate
228 + # next one.
229 + num_prefill_lookahead = (
230 + self.num_speculative_steps
231 + if self.speculative_config is not None
232 + and self.speculative_config.use_multi_module_mtp()
233 + else 1
234 + )
235 # General request states.
236 self.req_states = RequestState(
237 max_num_reqs=self.max_num_reqs,
@@ -231,6 +240,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
240 num_speculative_steps=self.num_speculative_steps,
241 vocab_size=self.vocab_size,
242 device=self.device,
243 + num_prefill_lookahead=num_prefill_lookahead,
244 )
245 self.input_buffers = InputBuffers(
246 max_num_reqs=self.max_num_reqs,
@@ -628,7 +638,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
638 torch.zeros(
639 input_batch.num_tokens,
640 dtype=torch.bool,
631 - device=self.device,
641 + device="cpu",
642 ),
643 )
644
@@ -1521,6 +1531,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1531 # NOTE: This is done here because postprocess updates
1532 # num_computed_prefill_tokens.
1533 # The EAGLE/MTP drafter reads one position ahead of the target.
1534 + # TODO(TheEpicDolphin): Gather MM embeddings for all speculative
1535 + # steps during multi-module MTP.
1536 mm_inputs = self.model_state.gather_mm_embeddings(
1537 input_batch, draft_lookahead=1
1538 )
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+3
@@ -3,6 +3,7 @@
3 from collections.abc import Mapping
4 from typing import Any
5
6 +import numpy as np
7 import torch
8 import torch.nn as nn
9
@@ -282,6 +283,7 @@ class DFlashSpeculator(DraftModelSpeculator):
283 step: int,
284 num_query_per_req: int | None = None,
285 causal: bool | Mapping[int, bool] = False,
286 + query_start_loc_np: np.ndarray | None = None,
287 ) -> dict[str, Any] | None:
288 if not self.draft_attn_layer_names:
289 return None
@@ -294,6 +296,7 @@ class DFlashSpeculator(DraftModelSpeculator):
296 step=step,
297 num_query_per_req=self.num_query_per_req,
298 causal=causal,
299 + query_start_loc_np=query_start_loc_np,
300 )
301
302 @torch.inference_mode()
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py new
+1106
@@ -0,0 +1,1106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +import torch.nn as nn
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.config.compilation import CUDAGraphMode
10 +from vllm.forward_context import BatchDescriptor, set_forward_context
11 +from vllm.logger import init_logger
12 +from vllm.multimodal import MULTIMODAL_REGISTRY
13 +from vllm.triton_utils import tl, triton
14 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
15 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
16 +from vllm.v1.worker.gpu.cudagraph_utils import (
17 + get_uniform_token_count,
18 +)
19 +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
20 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21 +from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
22 + SpeculatorCudaGraphManager,
23 +)
24 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
25 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
26 +
27 +logger = init_logger(__name__)
28 +
29 +
30 +class MultiModuleMTPSpeculator(DraftModelSpeculator):
31 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
32 + super().__init__(vllm_config, device)
33 +
34 + self.hidden_states = torch.zeros(
35 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
36 + )
37 + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
38 + self.last_token_indices = torch.zeros(
39 + self.max_num_reqs, dtype=torch.int64, device=device
40 + )
41 +
42 + self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
43 + self.draft_model_config
44 + )
45 + # HACK: the Inkling MTP draft has no MM processor of its own (its draft
46 + # config is flattened text-only), but it consumes the target's merged
47 + # embeddings at draft prefill — treat it as MM-capable whenever the
48 + # target is.
49 + if (
50 + not self.supports_mm_inputs
51 + and self.draft_model_config.hf_config.model_type == "inkling_mtp"
52 + ):
53 + self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
54 + vllm_config.model_config
55 + )
56 + self.inputs_embeds: torch.Tensor | None = None
57 + if self.supports_mm_inputs:
58 + self.inputs_embeds = torch.zeros(
59 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
60 + )
61 +
62 + # Input id overrides for the last num_speculative_steps - 1 draft steps.
63 + # Used by chunked-prefilling requests to swap in the future prefill token
64 + # for the sampled draft token. Non-chunked-prefilling requests fill the -1
65 + # value, indicating that the sampled draft token should be used.
66 + self.draft_input_id_overrides = torch.full(
67 + (self.max_num_reqs, self.num_speculative_steps - 1),
68 + -1,
69 + dtype=torch.int64,
70 + device=device,
71 + )
72 +
73 + # Cached input ids, embeddings, and target hidden states from the last
74 + # decode step. Used to re-prefill tokens to update stale KV cache slots
75 + # in later MTP modules.
76 + self.cached_draft_input_ids = torch.zeros(
77 + self.max_num_reqs,
78 + self.num_speculative_steps - 1,
79 + dtype=torch.int64,
80 + device=self.device,
81 + )
82 + self.cached_draft_input_embeds: torch.Tensor | None = None
83 + if self.supports_mm_inputs:
84 + self.cached_draft_input_embeds = torch.zeros(
85 + self.max_num_reqs,
86 + self.num_speculative_steps - 1,
87 + self.hidden_size,
88 + dtype=self.dtype,
89 + device=self.device,
90 + )
91 + self.cached_target_hidden_states = torch.zeros(
92 + self.max_num_reqs,
93 + self.num_speculative_steps - 1,
94 + self.hidden_size,
95 + dtype=self.dtype,
96 + device=self.device,
97 + )
98 +
99 + self.cudagraph_manager: SpeculatorCudaGraphManager | None = None
100 +
101 + def load_draft_model(
102 + self,
103 + target_model: nn.Module,
104 + target_attn_layer_names: set[str],
105 + ) -> nn.Module:
106 + return load_eagle_model(target_model, self.vllm_config)
107 +
108 + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
109 + # TODO(TheEpicDolphin): Support piecewise cudagraph for multi-module MTP.
110 + if cudagraph_mode.has_piecewise_cudagraphs():
111 + cudagraph_mode = (
112 + CUDAGraphMode.FULL_DECODE_ONLY
113 + if cudagraph_mode.has_full_cudagraphs()
114 + else CUDAGraphMode.NONE
115 + )
116 + self.cudagraph_manager = SpeculatorCudaGraphManager(
117 + self.vllm_config,
118 + self.device,
119 + cudagraph_mode,
120 + self.num_speculative_steps + 1,
121 + )
122 +
123 + def capture(self) -> None:
124 + logger.info("Capturing model for multi-module MTP speculator...")
125 + # Reset indices to zeros to prevent stale values from prior
126 + # dummy runs to cause out-of-bounds indexing during capture.
127 + self.last_token_indices.zero_()
128 + assert self.cudagraph_manager is not None
129 + if self.cudagraph_manager.use_breakable_cg:
130 + self.cudagraph_manager.init_breakable_cg_runner(self.model)
131 + self.cudagraph_manager.capture(
132 + self._generate_drafts,
133 + self.model_state,
134 + self.input_buffers,
135 + self.block_tables,
136 + self.attn_groups,
137 + self.kv_cache_config,
138 + progress_bar_desc="Capturing multi-module MTP CUDA graphs",
139 + )
140 +
141 + @torch.inference_mode()
142 + def propose(
143 + self,
144 + input_batch: InputBatch,
145 + attn_metadata: dict[str, Any],
146 + slot_mappings: dict[str, torch.Tensor],
147 + # [num_tokens, hidden_size]
148 + last_hidden_states: torch.Tensor,
149 + # num_layers x [num_tokens, hidden_size]
150 + aux_hidden_states: list[torch.Tensor] | None,
151 + # [num_reqs]
152 + num_sampled: torch.Tensor,
153 + # [num_reqs]
154 + num_rejected: torch.Tensor,
155 + # [max_num_reqs]
156 + last_sampled: torch.Tensor,
157 + # [num_prefill_lookahead, max_num_reqs]
158 + next_prefill_tokens: torch.Tensor,
159 + # [max_num_reqs]
160 + temperature: torch.Tensor,
161 + # [max_num_reqs]
162 + seeds: torch.Tensor,
163 + num_tokens_across_dp: torch.Tensor | None = None,
164 + dummy_run: bool = False,
165 + skip_attn_for_dummy_run: bool = False,
166 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
167 + is_profile: bool = False,
168 + ) -> torch.Tensor:
169 + num_reqs = input_batch.num_reqs
170 + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound
171 + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
172 + self.draft_max_seq_len = min(max_seq_len, self.max_model_len)
173 +
174 + self._copy_request_inputs(
175 + num_reqs,
176 + input_batch.idx_mapping,
177 + temperature,
178 + seeds,
179 + )
180 +
181 + num_tokens = input_batch.num_tokens
182 + max_query_len = input_batch.num_scheduled_tokens.max()
183 + self.input_buffers.query_start_loc[: num_reqs + 1].copy_(
184 + input_batch.query_start_loc[: num_reqs + 1]
185 + )
186 +
187 + self._prepare_inputs(
188 + last_hidden_states,
189 + input_batch,
190 + num_tokens,
191 + max_query_len,
192 + num_sampled,
193 + num_rejected,
194 + last_sampled,
195 + next_prefill_tokens,
196 + mm_inputs,
197 + )
198 +
199 + # When all requests are decoding (no true prefills), each has
200 + # num_speculative_steps + 1 tokens, enabling FULL graph replay.
201 + uniform_token_count = get_uniform_token_count(
202 + num_reqs,
203 + num_tokens,
204 + max_query_len,
205 + )
206 + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
207 + self.cudagraph_manager,
208 + num_reqs,
209 + input_batch.num_tokens_after_padding,
210 + uniform_token_count,
211 + dp_size=self.dp_size,
212 + dp_rank=self.dp_rank,
213 + need_eager=is_profile,
214 + )
215 +
216 + # Rebuild the slot mappings and attention metadata.
217 + skip_attn = dummy_run and skip_attn_for_dummy_run
218 + if not skip_attn:
219 + # Build the slot mappings and attention metadata.
220 + slot_mappings_tensor = self.block_tables.compute_slot_mappings(
221 + self.idx_mapping[:num_reqs],
222 + self.input_buffers.query_start_loc,
223 + self.input_buffers.positions,
224 + batch_desc.num_tokens,
225 + )
226 + # Apply padding values to slots not corresponding to real draft
227 + # tokens to prevent stale value writes.
228 + pad_trailing_draft_slots(
229 + slot_mappings_tensor,
230 + self.input_buffers.query_start_loc,
231 + self.last_token_indices[:num_reqs],
232 + num_reqs,
233 + )
234 + slot_mappings = build_slot_mappings_by_layer(
235 + slot_mappings_tensor, self.kv_cache_config
236 + )
237 + draft_attn_metadata = self._build_draft_attn_metadata(
238 + num_reqs=num_reqs,
239 + num_reqs_padded=batch_desc.num_reqs or num_reqs,
240 + num_tokens_padded=batch_desc.num_tokens,
241 + query_start_loc_np=input_batch.query_start_loc_np,
242 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
243 + step=0,
244 + )
245 + assert draft_attn_metadata is not None
246 + attn_metadata = draft_attn_metadata
247 +
248 + self._prepare_eplb_forward(num_tokens)
249 +
250 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
251 + assert self.cudagraph_manager is not None
252 + self.cudagraph_manager.run_fullgraph(batch_desc)
253 + else:
254 + self._generate_drafts(
255 + num_reqs,
256 + batch_desc.num_tokens,
257 + attn_metadata,
258 + slot_mappings,
259 + num_tokens_across_dp=num_tokens_across_dp,
260 + cudagraph_runtime_mode=batch_desc.cg_mode,
261 + )
262 + return self.draft_tokens[:num_reqs]
263 +
264 + @torch.inference_mode()
265 + def _run_model(
266 + self,
267 + num_tokens: int,
268 + attn_metadata: dict[str, Any] | None,
269 + slot_mappings: dict[str, torch.Tensor] | None,
270 + num_tokens_across_dp: torch.Tensor | None,
271 + spec_module_idx: int,
272 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
273 + ) -> tuple[torch.Tensor, torch.Tensor]:
274 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
275 + with set_forward_context(
276 + attn_metadata,
277 + self.vllm_config,
278 + num_tokens=num_tokens,
279 + cudagraph_runtime_mode=cudagraph_runtime_mode,
280 + num_tokens_across_dp=num_tokens_across_dp,
281 + slot_mapping=slot_mappings,
282 + batch_descriptor=batch_descriptor,
283 + ):
284 + model_inputs = dict(
285 + input_ids=self.input_buffers.input_ids[:num_tokens],
286 + positions=self.input_buffers.positions[:num_tokens],
287 + hidden_states=self.hidden_states[:num_tokens],
288 + inputs_embeds=(
289 + self.inputs_embeds[:num_tokens]
290 + if self.inputs_embeds is not None
291 + else None
292 + ),
293 + spec_step_idx=spec_module_idx,
294 + )
295 + if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
296 + # PIECEWISE cudagraph (compiled PW or breakable), chosen inside
297 + # run_pw_graph.
298 + assert self.cudagraph_manager is not None
299 + ret_hidden_states = self.cudagraph_manager.run_pw_graph(
300 + self.model, model_inputs
301 + )
302 + else:
303 + # Eager (NONE): call the raw model directly.
304 + ret_hidden_states = self.model(**model_inputs)
305 + # Some MTP models declare a single-tensor contract but return
306 + # (logits_hidden, feedback_hidden) for final-norm correctness.
307 + if isinstance(ret_hidden_states, tuple):
308 + last_hidden_states, hidden_states = ret_hidden_states
309 + else:
310 + last_hidden_states = ret_hidden_states
311 + hidden_states = ret_hidden_states
312 + return last_hidden_states, hidden_states
313 +
314 + def _prepare_inputs(
315 + self,
316 + target_hidden_states: torch.Tensor,
317 + input_batch: InputBatch,
318 + num_tokens: int,
319 + max_query_len: int,
320 + # [num_reqs]
321 + num_sampled: torch.Tensor,
322 + # [num_reqs]
323 + num_rejected: torch.Tensor,
324 + # [max_num_reqs]
325 + last_sampled: torch.Tensor,
326 + # [num_prefill_lookahead, max_num_reqs]
327 + next_prefill_tokens: torch.Tensor,
328 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
329 + ) -> None:
330 + num_reqs = input_batch.num_reqs
331 + prepare_input_buffers(
332 + num_reqs,
333 + input_batch,
334 + self.cached_draft_input_ids,
335 + num_sampled,
336 + num_rejected,
337 + last_sampled,
338 + next_prefill_tokens,
339 + self.draft_input_id_overrides,
340 + self.input_buffers,
341 + self.last_token_indices,
342 + self.max_num_reqs,
343 + self.num_speculative_steps,
344 + )
345 +
346 + # Compute the input embeddings with the MM embeddings merged in.
347 + # TODO(TheEpicDolphin): When the batch has no MM content (is_mm_embed
348 + # all False), skip this embed/copy and run the model with
349 + # inputs_embeds=None. Requirements:
350 + # 1. Eager steps can switch freely, but FULL cudagraphs bake the
351 + # embeds/no-embeds path at capture, so extending to decode steps
352 + # needs a uses_input_embeds property in the graph descriptors,
353 + # resulting in 2x captures.
354 + # 2. The skip condition must also verify the request's cached
355 + # re-prefill window (cached_draft_input_embeds) holds no MM
356 + # embeddings, or the rejection re-prefill gap would be re-embedded
357 + # from token ids incorrectly.
358 + if self.inputs_embeds is not None:
359 + mm_embeds, is_mm_embed = mm_inputs or (None, None)
360 + self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
361 + self.input_buffers.input_ids[:num_tokens],
362 + multimodal_embeddings=mm_embeds,
363 + is_multimodal=is_mm_embed,
364 + )
365 +
366 + prepare_input_hidden_states_and_embeddings(
367 + num_reqs,
368 + max_query_len,
369 + self.hidden_states,
370 + target_hidden_states,
371 + self.cached_target_hidden_states,
372 + self.inputs_embeds,
373 + self.cached_draft_input_embeds,
374 + input_batch,
375 + self.input_buffers,
376 + num_rejected,
377 + self.num_speculative_steps,
378 + )
379 +
380 + def _generate_drafts(
381 + self,
382 + num_reqs: int,
383 + num_tokens: int,
384 + attn_metadata: dict[str, Any] | None,
385 + slot_mappings: dict[str, torch.Tensor] | None,
386 + num_tokens_across_dp: torch.Tensor | None,
387 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
388 + ) -> None:
389 + last_token_indices = self.last_token_indices[:num_reqs]
390 + sample_positions = self.input_buffers.positions[last_token_indices]
391 + idx_mapping = self.idx_mapping[:num_reqs]
392 +
393 + # Cache the trailing token's ids, hidden states (and embeddings for
394 + # MM models), which are needed if the trailing tokens are re-prefilled
395 + # during the next decode step.
396 + cache_inputs(
397 + self.input_buffers,
398 + self.inputs_embeds,
399 + self.hidden_states,
400 + self.cached_draft_input_ids,
401 + self.cached_draft_input_embeds,
402 + self.cached_target_hidden_states,
403 + last_token_indices,
404 + idx_mapping,
405 + num_reqs,
406 + self.num_speculative_steps,
407 + use_input_embeds=self.inputs_embeds is not None,
408 + )
409 +
410 + for step in range(self.num_speculative_steps):
411 + # Update the current draft step.
412 + self.current_draft_step.fill_(step)
413 +
414 + # Run the model forward pass.
415 + last_hidden_states, hidden_states = self._run_model(
416 + num_tokens,
417 + attn_metadata,
418 + slot_mappings,
419 + num_tokens_across_dp=num_tokens_across_dp,
420 + spec_module_idx=step,
421 + cudagraph_runtime_mode=cudagraph_runtime_mode,
422 + )
423 +
424 + # Sample draft tokens for the current step.
425 + sample_hidden_states = last_hidden_states[last_token_indices]
426 + draft_tokens = self.sample_draft(
427 + sample_hidden_states,
428 + sample_positions,
429 + idx_mapping,
430 + self.temperature,
431 + self.seeds,
432 + self.current_draft_step,
433 + self.draft_logits,
434 + )
435 +
436 + self.draft_tokens[:num_reqs, step] = draft_tokens
437 + if step < self.num_speculative_steps - 1:
438 + self.hidden_states[:num_tokens] = hidden_states
439 + # Mid-prefill requests append the known future prefill token
440 + # instead of the sampled draft.
441 + overrides = self.draft_input_id_overrides[:num_reqs, step]
442 + next_input_tokens = torch.where(overrides >= 0, overrides, draft_tokens)
443 + # Shift the draft inputs left by one and append the next
444 + # input token id/embeddings.
445 + draft_embeds = (
446 + self.model.embed_input_ids(next_input_tokens)
447 + if self.inputs_embeds is not None
448 + else None
449 + )
450 + update_draft_inputs(
451 + next_input_tokens,
452 + draft_embeds,
453 + self.input_buffers,
454 + self.inputs_embeds,
455 + last_token_indices,
456 + idx_mapping,
457 + num_reqs,
458 + )
459 + sample_positions += 1
460 +
461 +
462 +@triton.jit
463 +def _prepare_input_buffers_kernel(
464 + last_token_indices_ptr,
465 + draft_input_ids_ptr,
466 + draft_positions_ptr,
467 + draft_seq_lens_ptr,
468 + target_input_ids_ptr,
469 + target_positions_ptr,
470 + cached_draft_input_ids_ptr,
471 + cached_draft_input_ids_stride0,
472 + draft_input_id_overrides_ptr,
473 + draft_input_id_overrides_stride0,
474 + idx_mapping_ptr,
475 + last_sampled_ptr,
476 + next_prefill_tokens_ptr,
477 + next_prefill_tokens_stride0,
478 + num_sampled_ptr,
479 + num_rejected_ptr,
480 + target_seq_lens_ptr,
481 + query_start_loc_ptr,
482 + max_num_reqs,
483 + num_speculative_steps,
484 + BLOCK_SIZE: tl.constexpr,
485 +):
486 + req_idx = tl.program_id(0)
487 + num_reqs = tl.num_programs(0)
488 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
489 +
490 + query_start = tl.load(query_start_loc_ptr + req_idx)
491 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
492 + query_len = query_end - query_start
493 + seq_len = tl.load(target_seq_lens_ptr + req_idx)
494 +
495 + # Get the number of rejected tokens, and the number of trailing tokens from
496 + # the last decode step that need to be re-prefilled to update the stale
497 + # KV cache slots in later MTP modules.
498 + num_rejected = tl.load(num_rejected_ptr + req_idx)
499 + num_reprefill_tokens = max(0, num_rejected - 1)
500 +
501 + # The rejected tokens are excluded as inputs to the draft model.
502 + num_input_tokens = query_len - num_rejected
503 + # Re-prefilled tokens are packed into the query window without widening
504 + # it, so the effective KV length shrinks by num_reprefill_tokens.
505 + # Example (r = rejected draft token, x = unused trailing slot, later padded)
506 + # re-prefilling t1 and t2 from the last decode step:
507 + # t0 t1 t2 [ t3 t4 r r r ] => t0 [ t1 t2 t3 t4 . ]
508 + seq_len -= num_reprefill_tokens
509 +
510 + # Write the updated sequence length.
511 + tl.store(draft_seq_lens_ptr + req_idx, seq_len)
512 +
513 + # Get the next draft input token.
514 + num_sampled = tl.load(num_sampled_ptr + req_idx)
515 + if num_sampled > 0:
516 + next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
517 + else:
518 + # Chunked prefill. Seed with the next prefill token.
519 + next_token = tl.load(next_prefill_tokens_ptr + req_state_idx)
520 +
521 + # For chunked-prefilling requests, copy the future prefill tokens that
522 + # will later override the input token ids for the last
523 + # num_speculative_steps - 1 draft steps.
524 + for i in range(1, num_speculative_steps):
525 + future_token = tl.load(
526 + next_prefill_tokens_ptr + i * next_prefill_tokens_stride0 + req_state_idx
527 + )
528 + override = tl.where(num_sampled > 0, -1, future_token).to(tl.int64)
529 + tl.store(
530 + draft_input_id_overrides_ptr
531 + + req_idx * draft_input_id_overrides_stride0
532 + + i
533 + - 1,
534 + override,
535 + )
536 +
537 + # Copy the target's input ids (read at the target offset) shifted left by 1,
538 + # and right by the number of re-prefills, into the draft buffer.
539 + for i in range(1, num_input_tokens, BLOCK_SIZE):
540 + block = i + tl.arange(0, BLOCK_SIZE)
541 + mask = block < num_input_tokens
542 + input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask)
543 + tl.store(
544 + draft_input_ids_ptr + query_start + num_reprefill_tokens - 1 + block,
545 + input_ids,
546 + mask=mask,
547 + )
548 + last_token_index = query_start + num_reprefill_tokens + num_input_tokens - 1
549 + tl.store(last_token_indices_ptr + req_idx, last_token_index)
550 + tl.store(draft_input_ids_ptr + last_token_index, next_token)
551 +
552 + # Copy the target's positions, shifted over by the number of tokens to be
553 + # re-prefilled.
554 + for i in range(0, num_input_tokens, BLOCK_SIZE):
555 + block = i + tl.arange(0, BLOCK_SIZE)
556 + mask = block < num_input_tokens
557 + target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask)
558 + tl.store(
559 + draft_positions_ptr + query_start + num_reprefill_tokens + block,
560 + target_pos,
561 + mask=mask,
562 + )
563 +
564 + # Fill the re-prefill gap with the cached token ids from the previous
565 + # decode step. These tokens sit immediately before the query's first
566 + # token, so their positions are contiguous and derived here.
567 + first_position = tl.load(target_positions_ptr + query_start)
568 + for i in range(num_reprefill_tokens):
569 + cache_read_slot = num_speculative_steps - 1 - num_reprefill_tokens + i
570 + cached_token_id = tl.load(
571 + cached_draft_input_ids_ptr
572 + + req_state_idx * cached_draft_input_ids_stride0
573 + + cache_read_slot
574 + )
575 + tl.store(draft_input_ids_ptr + query_start + i, cached_token_id)
576 + tl.store(
577 + draft_positions_ptr + query_start + i,
578 + first_position - num_reprefill_tokens + i,
579 + )
580 +
581 + if req_idx == (num_reqs - 1):
582 + # Pad query_start_loc for CUDA graphs.
583 + for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
584 + block = i + tl.arange(0, BLOCK_SIZE)
585 + mask = block < max_num_reqs + 1
586 + tl.store(query_start_loc_ptr + block, query_end, mask=mask)
587 + # Pad seq_lens for CUDA graphs.
588 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
589 + block = i + tl.arange(0, BLOCK_SIZE)
590 + mask = block < max_num_reqs
591 + tl.store(draft_seq_lens_ptr + block, 0, mask=mask)
592 + # Pad last_token_indices for CUDA graphs.
593 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
594 + block = i + tl.arange(0, BLOCK_SIZE)
595 + mask = block < max_num_reqs
596 + tl.store(last_token_indices_ptr + block, 0, mask=mask)
597 +
598 +
599 +def prepare_input_buffers(
600 + num_reqs: int,
601 + input_batch: InputBatch,
602 + # [max_num_reqs, num_speculative_steps - 1]
603 + cached_draft_input_ids: torch.Tensor,
604 + # [num_reqs]
605 + num_sampled: torch.Tensor,
606 + # [num_reqs]
607 + num_rejected: torch.Tensor,
608 + # [max_num_reqs]
609 + last_sampled: torch.Tensor,
610 + # [num_prefill_lookahead, max_num_reqs]
611 + next_prefill_tokens: torch.Tensor,
612 + # [max_num_reqs, num_speculative_steps - 1]
613 + draft_input_id_overrides: torch.Tensor,
614 + input_buffers: InputBuffers,
615 + # [max_num_reqs]
616 + last_token_indices: torch.Tensor,
617 + max_num_reqs: int,
618 + num_speculative_steps: int,
619 +) -> None:
620 + _prepare_input_buffers_kernel[(num_reqs,)](
621 + last_token_indices,
622 + input_buffers.input_ids,
623 + input_buffers.positions,
624 + input_buffers.seq_lens,
625 + input_batch.input_ids,
626 + input_batch.positions,
627 + cached_draft_input_ids,
628 + cached_draft_input_ids.stride(0) if cached_draft_input_ids is not None else 0,
629 + draft_input_id_overrides,
630 + draft_input_id_overrides.stride(0),
631 + input_batch.idx_mapping,
632 + last_sampled,
633 + next_prefill_tokens,
634 + next_prefill_tokens.stride(0),
635 + num_sampled,
636 + num_rejected,
637 + input_batch.seq_lens,
638 + input_buffers.query_start_loc,
639 + max_num_reqs,
640 + num_speculative_steps,
641 + BLOCK_SIZE=1024,
642 + )
643 +
644 +
645 +@triton.jit
646 +def _prepare_input_hidden_states_and_embeddings_kernel(
647 + draft_input_hidden_states_ptr,
648 + draft_input_hidden_states_stride0,
649 + target_hidden_states_ptr,
650 + target_hidden_states_stride0,
651 + cached_target_hidden_states_ptr,
652 + cached_target_hidden_states_stride0,
653 + cached_target_hidden_states_stride1,
654 + input_embeds_ptr,
655 + input_embeds_stride0,
656 + cached_draft_input_embeds_ptr,
657 + cached_draft_input_embeds_stride0,
658 + cached_draft_input_embeds_stride1,
659 + idx_mapping_ptr,
660 + num_rejected_ptr,
661 + query_start_loc_ptr,
662 + num_speculative_steps,
663 + hidden_size,
664 + BLOCK_SIZE_Q: tl.constexpr,
665 + BLOCK_SIZE_H: tl.constexpr,
666 + USE_INPUT_EMBEDS: tl.constexpr,
667 +):
668 + req_idx = tl.program_id(0)
669 + query_block_idx = tl.program_id(1)
670 + dim_block_idx = tl.program_id(2)
671 + dim_block = dim_block_idx * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
672 + dim_mask = dim_block < hidden_size
673 +
674 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
675 + query_start = tl.load(query_start_loc_ptr + req_idx)
676 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
677 + query_len = query_end - query_start
678 +
679 + # Get the number of rejected tokens, and the number of trailing tokens from
680 + # the last decode step that need to be re-prefilled to update the stale
681 + # KV cache slots in later MTP modules.
682 + num_rejected = tl.load(num_rejected_ptr + req_idx)
683 + num_reprefill_hidden_states = max(0, num_rejected - 1)
684 +
685 + # The rejected tokens are excluded as inputs to the draft model.
686 + num_input_hidden_states = query_len - num_rejected
687 +
688 + # Copy the output target hidden states as inputs to the first MTP module,
689 + # shifted right by the re-prefill gap. Each program copies one
690 + # (BLOCK_SIZE_Q, BLOCK_SIZE_H) tile of this request's query.
691 + query_block = query_block_idx * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)
692 + query_mask = (query_block < num_input_hidden_states)[:, None] & dim_mask[None, :]
693 + hidden_state = tl.load(
694 + target_hidden_states_ptr
695 + + (query_start + query_block)[:, None] * target_hidden_states_stride0
696 + + dim_block[None, :],
697 + mask=query_mask,
698 + )
699 + tl.store(
700 + draft_input_hidden_states_ptr
701 + + (query_start + num_reprefill_hidden_states + query_block)[:, None]
702 + * draft_input_hidden_states_stride0
703 + + dim_block[None, :],
704 + hidden_state,
705 + mask=query_mask,
706 + )
707 +
708 + if query_block_idx == 0:
709 + # Fill the re-prefill gap with the cached hidden states and embeddings
710 + # (for MM models) from the previous decode step, mirroring the token
711 + # ids and positions inserted by _prepare_input_buffers_kernel. The gap
712 + # is at most num_speculative_steps - 1 tokens, so a serial loop in the
713 + # first query block suffices.
714 + for i in range(num_reprefill_hidden_states):
715 + cache_read_slot = (
716 + num_speculative_steps - 1 - num_reprefill_hidden_states + i
717 + )
718 + cached_hidden_state = tl.load(
719 + cached_target_hidden_states_ptr
720 + + req_state_idx * cached_target_hidden_states_stride0
721 + + cache_read_slot * cached_target_hidden_states_stride1
722 + + dim_block,
723 + mask=dim_mask,
724 + )
725 + tl.store(
726 + draft_input_hidden_states_ptr
727 + + (query_start + i) * draft_input_hidden_states_stride0
728 + + dim_block,
729 + cached_hidden_state,
730 + mask=dim_mask,
731 + )
732 + if USE_INPUT_EMBEDS:
733 + cached_embed = tl.load(
734 + cached_draft_input_embeds_ptr
735 + + req_state_idx * cached_draft_input_embeds_stride0
736 + + cache_read_slot * cached_draft_input_embeds_stride1
737 + + dim_block,
738 + mask=dim_mask,
739 + )
740 + tl.store(
741 + input_embeds_ptr
742 + + (query_start + i) * input_embeds_stride0
743 + + dim_block,
744 + cached_embed,
745 + mask=dim_mask,
746 + )
747 +
748 +
749 +def prepare_input_hidden_states_and_embeddings(
750 + num_reqs: int,
751 + # Upper bound on the draft query length of any request in the batch.
752 + max_query_len: int,
753 + # [num_tokens, hidden_size]
754 + hidden_states: torch.Tensor,
755 + # [num_tokens, hidden_size]
756 + target_hidden_states: torch.Tensor,
757 + # [max_num_reqs, num_speculative_steps - 1, hidden_size]
758 + cached_target_hidden_states: torch.Tensor | None,
759 + # [num_tokens, hidden_size]
760 + input_embeds: torch.Tensor | None,
761 + # [max_num_reqs, num_speculative_steps - 1, hidden_size]
762 + cached_draft_input_embeds: torch.Tensor | None,
763 + input_batch: InputBatch,
764 + input_buffers: InputBuffers,
765 + # [num_reqs]
766 + num_rejected: torch.Tensor,
767 + num_speculative_steps: int,
768 +) -> None:
769 + use_input_embeds = input_embeds is not None
770 + hidden_size = target_hidden_states.shape[-1]
771 + query_block_size = 16
772 + hidden_block_size = 256
773 + grid = (
774 + num_reqs,
775 + triton.cdiv(max_query_len, query_block_size),
776 + triton.cdiv(hidden_size, hidden_block_size),
777 + )
778 + _prepare_input_hidden_states_and_embeddings_kernel[grid](
779 + hidden_states,
780 + hidden_states.stride(0),
781 + target_hidden_states,
782 + target_hidden_states.stride(0),
783 + cached_target_hidden_states,
784 + cached_target_hidden_states.stride(0)
785 + if cached_target_hidden_states is not None
786 + else 0,
787 + cached_target_hidden_states.stride(1)
788 + if cached_target_hidden_states is not None
789 + else 0,
790 + input_embeds,
791 + input_embeds.stride(0) if input_embeds is not None else 0,
792 + cached_draft_input_embeds,
793 + cached_draft_input_embeds.stride(0)
794 + if cached_draft_input_embeds is not None
795 + else 0,
796 + cached_draft_input_embeds.stride(1)
797 + if cached_draft_input_embeds is not None
798 + else 0,
799 + input_batch.idx_mapping,
800 + num_rejected,
801 + input_buffers.query_start_loc,
802 + num_speculative_steps,
803 + hidden_size,
804 + BLOCK_SIZE_Q=query_block_size,
805 + BLOCK_SIZE_H=hidden_block_size,
806 + USE_INPUT_EMBEDS=use_input_embeds,
807 + )
808 +
809 +
810 +@triton.jit
811 +def _pad_trailing_draft_slots_kernel(
812 + slot_mappings_ptr,
813 + slot_mappings_stride0,
814 + query_start_loc_ptr,
815 + last_token_indices_ptr,
816 + PAD_ID,
817 + BLOCK_SIZE: tl.constexpr,
818 +):
819 + group_idx = tl.program_id(0)
820 + req_idx = tl.program_id(1)
821 + # Slots computed from stale token positions in the range
822 + # [last_token_index + 1, query_end) can result in writes to blocks.
823 + # Pad these slot values so that attention kernels ignore them.
824 + start = tl.load(last_token_indices_ptr + req_idx) + 1
825 + end = tl.load(query_start_loc_ptr + req_idx + 1)
826 + base = slot_mappings_ptr + group_idx * slot_mappings_stride0
827 + for i in range(start, end, BLOCK_SIZE):
828 + offs = i + tl.arange(0, BLOCK_SIZE)
829 + mask = offs < end
830 + tl.store(base + offs, PAD_ID, mask=mask)
831 +
832 +
833 +def pad_trailing_draft_slots(
834 + # [num_groups, num_tokens_padded]
835 + slot_mappings: torch.Tensor,
836 + # [num_reqs + 1]
837 + query_start_loc: torch.Tensor,
838 + # [num_reqs]
839 + last_token_indices: torch.Tensor,
840 + num_reqs: int,
841 +) -> None:
842 + num_groups = slot_mappings.shape[0]
843 + _pad_trailing_draft_slots_kernel[(num_groups, num_reqs)](
844 + slot_mappings,
845 + slot_mappings.stride(0),
846 + query_start_loc,
847 + last_token_indices,
848 + PAD_SLOT_ID,
849 + BLOCK_SIZE=256,
850 + )
851 +
852 +
853 +@triton.jit
854 +def _cache_inputs_kernel(
855 + draft_input_ids_ptr,
856 + draft_input_embeds_ptr,
857 + draft_input_embeds_stride0,
858 + draft_input_hidden_states_ptr,
859 + draft_input_hidden_states_stride0,
860 + cached_draft_input_ids_ptr,
861 + cached_draft_input_ids_stride0,
862 + cached_draft_input_embeds_ptr,
863 + cached_draft_input_embeds_stride0,
864 + cached_draft_input_embeds_stride1,
865 + cached_target_hidden_states_ptr,
866 + cached_target_hidden_states_stride0,
867 + cached_target_hidden_states_stride1,
868 + idx_mapping_ptr,
869 + last_token_indices_ptr,
870 + query_start_loc_ptr,
871 + num_speculative_steps,
872 + hidden_size,
873 + BLOCK_SIZE: tl.constexpr,
874 + USE_INPUT_EMBEDS: tl.constexpr,
875 +):
876 + req_idx = tl.program_id(0)
877 + block_idx = tl.program_id(1)
878 + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
879 + mask = block < hidden_size
880 +
881 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
882 + if req_state_idx < 0:
883 + # Skip cudagraph padded requests.
884 + return
885 +
886 + query_start = tl.load(query_start_loc_ptr + req_idx)
887 + last_token_index = tl.load(last_token_indices_ptr + req_idx)
888 +
889 + # Snapshot the last num_speculative_steps - 1 input draft token ids/hidden
890 + # states and embeddings (for MM models), indexed by request state. These
891 + # may be needed to re-prefill the tokens during the next decode step.
892 + cache_window_size = num_speculative_steps - 1
893 + window_start = last_token_index - cache_window_size + 1
894 + for i in range(max(window_start, query_start), last_token_index + 1):
895 + cache_write_slot = i - window_start
896 + if block_idx == 0:
897 + input_id = tl.load(draft_input_ids_ptr + i)
898 + tl.store(
899 + cached_draft_input_ids_ptr
900 + + req_state_idx * cached_draft_input_ids_stride0
901 + + cache_write_slot,
902 + input_id,
903 + )
904 + if USE_INPUT_EMBEDS:
905 + input_embeds = tl.load(
906 + draft_input_embeds_ptr + i * draft_input_embeds_stride0 + block,
907 + mask=mask,
908 + )
909 + tl.store(
910 + cached_draft_input_embeds_ptr
911 + + req_state_idx * cached_draft_input_embeds_stride0
912 + + cache_write_slot * cached_draft_input_embeds_stride1
913 + + block,
914 + input_embeds,
915 + mask=mask,
916 + )
917 + hidden_state = tl.load(
918 + draft_input_hidden_states_ptr
919 + + i * draft_input_hidden_states_stride0
920 + + block,
921 + mask=mask,
922 + )
923 + tl.store(
924 + cached_target_hidden_states_ptr
925 + + req_state_idx * cached_target_hidden_states_stride0
926 + + cache_write_slot * cached_target_hidden_states_stride1
927 + + block,
928 + hidden_state,
929 + mask=mask,
930 + )
931 +
932 +
933 +def cache_inputs(
934 + input_buffers: InputBuffers,
935 + # [num_tokens, hidden_size]
936 + draft_input_embeds: torch.Tensor | None,
937 + # [num_tokens, hidden_size]
938 + draft_input_hidden_states: torch.Tensor,
939 + # [max_num_reqs, num_speculative_steps - 1]
940 + cached_draft_input_ids: torch.Tensor,
941 + # [max_num_reqs, num_speculative_steps - 1, hidden_size]
942 + cached_draft_input_embeds: torch.Tensor | None,
943 + # [max_num_reqs, num_speculative_steps - 1, hidden_size]
944 + cached_target_hidden_states: torch.Tensor,
945 + # [num_reqs]
946 + last_token_indices: torch.Tensor,
947 + # [num_reqs]
948 + idx_mapping: torch.Tensor,
949 + num_reqs: int,
950 + num_speculative_steps: int,
951 + use_input_embeds: bool,
952 +) -> None:
953 + hidden_size = draft_input_hidden_states.shape[-1]
954 + hidden_block_size = 1024
955 + _cache_inputs_kernel[(num_reqs, triton.cdiv(hidden_size, hidden_block_size))](
956 + input_buffers.input_ids,
957 + draft_input_embeds,
958 + draft_input_embeds.stride(0) if draft_input_embeds is not None else 0,
959 + draft_input_hidden_states,
960 + draft_input_hidden_states.stride(0),
961 + cached_draft_input_ids,
962 + cached_draft_input_ids.stride(0),
963 + cached_draft_input_embeds,
964 + cached_draft_input_embeds.stride(0)
965 + if cached_draft_input_embeds is not None
966 + else 0,
967 + cached_draft_input_embeds.stride(1)
968 + if cached_draft_input_embeds is not None
969 + else 0,
970 + cached_target_hidden_states,
971 + cached_target_hidden_states.stride(0),
972 + cached_target_hidden_states.stride(1),
973 + idx_mapping,
974 + last_token_indices,
975 + input_buffers.query_start_loc,
976 + num_speculative_steps,
977 + hidden_size,
978 + BLOCK_SIZE=hidden_block_size,
979 + USE_INPUT_EMBEDS=use_input_embeds,
980 + )
981 +
982 +
983 +@triton.jit
984 +def _shift_input_ids_kernel(
985 + input_ids_ptr,
986 + idx_mapping_ptr,
987 + query_start_loc_ptr,
988 + last_token_indices_ptr,
989 + draft_tokens_ptr,
990 + BLOCK_SIZE: tl.constexpr,
991 +):
992 + req_idx = tl.program_id(0)
993 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
994 + if req_state_idx < 0:
995 + # Skip cudagraph padded requests.
996 + return
997 +
998 + query_start = tl.load(query_start_loc_ptr + req_idx)
999 + # Use the post-rejection last token index so the shift and insertion align
1000 + # with the position the draft token was sampled from.
1001 + last_token_index = tl.load(last_token_indices_ptr + req_idx)
1002 + query_len = last_token_index - query_start + 1
1003 +
1004 + # Shift input token ids to the left by one position and
1005 + # insert the last sampled draft token.
1006 + for i in range(1, query_len, BLOCK_SIZE):
1007 + block = i + tl.arange(0, BLOCK_SIZE)
1008 + mask = block < query_len
1009 + input_ids = tl.load(input_ids_ptr + query_start + block, mask=mask)
1010 + tl.store(input_ids_ptr + query_start + block - 1, input_ids, mask=mask)
1011 + draft_token = tl.load(draft_tokens_ptr + req_idx)
1012 + tl.store(input_ids_ptr + last_token_index, draft_token)
1013 +
1014 +
1015 +@triton.jit
1016 +def _shift_input_embeds_kernel(
1017 + input_embeds_ptr,
1018 + input_embeds_stride0,
1019 + draft_embeds_ptr,
1020 + draft_embeds_stride0,
1021 + idx_mapping_ptr,
1022 + query_start_loc_ptr,
1023 + last_token_indices_ptr,
1024 + hidden_size,
1025 + BLOCK_SIZE_Q: tl.constexpr,
1026 + BLOCK_SIZE_H: tl.constexpr,
1027 +):
1028 + req_idx = tl.program_id(0)
1029 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
1030 + if req_state_idx < 0:
1031 + # Skip cudagraph padded requests.
1032 + return
1033 +
1034 + block_idx = tl.program_id(1)
1035 + dim_block = block_idx * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
1036 + dim_mask = dim_block < hidden_size
1037 +
1038 + query_start = tl.load(query_start_loc_ptr + req_idx)
1039 + last_token_index = tl.load(last_token_indices_ptr + req_idx)
1040 + query_len = last_token_index - query_start + 1
1041 +
1042 + # Shift input token embeddings to the left by one position and
1043 + # insert the last sampled draft token's embeddings.
1044 + for i in range(1, query_len, BLOCK_SIZE_Q):
1045 + query_block = i + tl.arange(0, BLOCK_SIZE_Q)
1046 + query_mask = query_block < query_len
1047 + mask = query_mask[:, None] & dim_mask[None, :]
1048 + input_embed = tl.load(
1049 + input_embeds_ptr
1050 + + (query_start + query_block)[:, None] * input_embeds_stride0
1051 + + dim_block[None, :],
1052 + mask=mask,
1053 + )
1054 + tl.store(
1055 + input_embeds_ptr
1056 + + (query_start + query_block - 1)[:, None] * input_embeds_stride0
1057 + + dim_block[None, :],
1058 + input_embed,
1059 + mask=mask,
1060 + )
1061 + draft_embed = tl.load(
1062 + draft_embeds_ptr + req_idx * draft_embeds_stride0 + dim_block,
1063 + mask=dim_mask,
1064 + )
1065 + tl.store(
1066 + input_embeds_ptr + last_token_index * input_embeds_stride0 + dim_block,
1067 + draft_embed,
1068 + mask=dim_mask,
1069 + )
1070 +
1071 +
1072 +def update_draft_inputs(
1073 + draft_tokens: torch.Tensor,
1074 + draft_embeds: torch.Tensor | None,
1075 + input_buffers: InputBuffers,
1076 + input_embeds: torch.Tensor | None,
1077 + last_token_indices: torch.Tensor,
1078 + idx_mapping: torch.Tensor,
1079 + num_reqs: int,
1080 +) -> None:
1081 + _shift_input_ids_kernel[(num_reqs,)](
1082 + input_buffers.input_ids,
1083 + idx_mapping,
1084 + input_buffers.query_start_loc,
1085 + last_token_indices,
1086 + draft_tokens,
1087 + BLOCK_SIZE=1024,
1088 + )
1089 + if input_embeds is not None:
1090 + assert draft_embeds is not None
1091 + hidden_size = input_embeds.shape[-1]
1092 + hidden_block_size = 256
1093 + _shift_input_embeds_kernel[
1094 + (num_reqs, triton.cdiv(hidden_size, hidden_block_size))
1095 + ](
1096 + input_embeds,
1097 + input_embeds.stride(0),
1098 + draft_embeds,
1099 + draft_embeds.stride(0),
1100 + idx_mapping,
1101 + input_buffers.query_start_loc,
1102 + last_token_indices,
1103 + hidden_size,
1104 + BLOCK_SIZE_Q=16,
1105 + BLOCK_SIZE_H=hidden_block_size,
1106 + )
vllm/v1/worker/gpu/spec_decode/speculator.py
+27 -13
@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
4 from collections.abc import Mapping
5 from typing import Any
6
7 +import numpy as np
8 import torch
9 import torch.nn as nn
10
@@ -51,7 +52,7 @@ class BaseSpeculator(ABC):
52 num_rejected: torch.Tensor,
53 # [max_num_reqs]
54 last_sampled: torch.Tensor,
54 - # [max_num_reqs]
55 + # [num_prefill_lookahead, max_num_reqs]
56 next_prefill_tokens: torch.Tensor,
57 # [max_num_reqs]
58 temperature: torch.Tensor,
@@ -214,14 +215,28 @@ class DraftModelSpeculator(BaseSpeculator):
215 step: int,
216 num_query_per_req: int = 1,
217 causal: bool | Mapping[int, bool] = True,
218 + query_start_loc_np: np.ndarray | None = None,
219 ) -> dict[str, Any] | None:
218 - # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
219 - # Clamp keeps the series non-decreasing past num_reqs, which some
220 - # attention backends require.
221 - query_start_loc_cpu = (
222 - torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
223 - * num_query_per_req
224 - )
220 + if query_start_loc_np is not None:
221 + # Non-uniform query layout (e.g. multi-module MTP's mixed
222 + # prefill/decode queries); num_query_per_req is ignored.
223 + query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
224 + query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
225 + query_start_loc_np[: num_reqs + 1]
226 + )
227 + query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
228 + max_query_len = int(
229 + (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()
230 + )
231 + else:
232 + # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
233 + # Clamp keeps the series non-decreasing past num_reqs, which some
234 + # attention backends require.
235 + query_start_loc_cpu = (
236 + torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
237 + * num_query_per_req
238 + )
239 + max_query_len = num_query_per_req
240 block_tables = [
241 x[:num_reqs_padded] for x in self.block_tables.input_block_tables
242 ]
@@ -243,7 +258,7 @@ class DraftModelSpeculator(BaseSpeculator):
258 : num_reqs_padded + 1
259 ],
260 query_start_loc_cpu=query_start_loc_cpu,
246 - max_query_len=num_query_per_req,
261 + max_query_len=max_query_len,
262 seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
263 max_seq_len=self.draft_max_seq_len,
264 block_tables=block_tables,
@@ -325,7 +340,6 @@ class DraftModelSpeculator(BaseSpeculator):
340 self.temperature.copy_(temperature)
341 self.seeds.copy_(seeds)
342 self.idx_mapping[:num_reqs].copy_(idx_mapping)
328 - if self.draft_logits is not None:
329 - # idx_mapping for CG padded requests points to -1, which is ignored
330 - # during sampling to prevent writing stale values to draft logits.
331 - self.idx_mapping[num_reqs:].fill_(-1)
343 + # idx_mapping for CG padded requests points to -1, which is ignored
344 + # during sampling to prevent writing stale values to draft logits.
345 + self.idx_mapping[num_reqs:].fill_(-1)
vllm/v1/worker/gpu/states.py
+5 -1
@@ -15,6 +15,7 @@ class RequestState:
15 num_speculative_steps: int,
16 vocab_size: int,
17 device: torch.device,
18 + num_prefill_lookahead: int = 1,
19 ):
20 self.max_num_reqs = max_num_reqs
21 self.max_model_len = max_model_len
@@ -77,7 +78,10 @@ class RequestState:
78 )
79
80 self.next_prefill_tokens = torch.zeros(
80 - self.max_num_reqs, dtype=torch.int32, device=device
81 + num_prefill_lookahead,
82 + self.max_num_reqs,
83 + dtype=torch.int32,
84 + device=device,
85 )
86
87 @property