@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ceb0111a90
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 01:00 UTC

[Model Runner V2][Spec Decode] Add Gemma4 MTP support (#43241)

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

Giancarlo Delfin committed Jun 3, 2026 at 19:51 UTC ceb0111a90acd204a5444a384c4327cd0280ee4d
14 files changed +1243 -942
vllm/v1/attention/backends/flashinfer.py
+3 -5
@@ -62,7 +62,6 @@ from vllm.v1.attention.backends.utils import (
62 KVCacheLayoutType,
63 get_dcp_local_seq_lens,
64 get_kv_cache_layout,
65 - get_num_attention_heads_from_layers,
65 get_per_layer_parameters,
66 infer_global_hyperparameters,
67 split_decodes_and_prefills,
@@ -608,10 +607,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
607 self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
608 )
609
611 - # Compatible with models with non-uniform per-layer head counts.
612 - self.num_qo_heads = get_num_attention_heads_from_layers(
613 - vllm_config, layer_names
614 - ) or self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)
610 + self.num_qo_heads = self.model_config.get_num_attention_heads(
611 + self.vllm_config.parallel_config
612 + )
613
614 self.num_kv_heads = self.kv_cache_spec.num_kv_heads
615 self.head_dim = self.kv_cache_spec.head_size
vllm/v1/attention/backends/triton_attn.py
+4 -8
@@ -30,10 +30,7 @@ from vllm.v1.attention.backend import (
30 CommonAttentionMetadata,
31 MultipleOf,
32 )
33 -from vllm.v1.attention.backends.utils import (
34 - get_kv_cache_layout,
35 - get_num_attention_heads_from_layers,
36 -)
33 +from vllm.v1.attention.backends.utils import get_kv_cache_layout
34 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
35 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
36 triton_reshape_and_cache_flash,
@@ -142,10 +139,9 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
139 self.block_size = kv_cache_spec.block_size
140
141 model_config = vllm_config.model_config
145 - # Compatible with models with non-uniform per-layer head counts.
146 - self.num_heads_q = get_num_attention_heads_from_layers(
147 - vllm_config, layer_names
148 - ) or model_config.get_num_attention_heads(vllm_config.parallel_config)
142 + self.num_heads_q = model_config.get_num_attention_heads(
143 + vllm_config.parallel_config
144 + )
145 self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config)
146 self.headdim = model_config.get_head_size()
147
vllm/v1/attention/backends/utils.py
-26
@@ -136,32 +136,6 @@ def get_per_layer_parameters(
136 return per_layer_params
137
138
139 -def get_num_attention_heads_from_layers(
140 - vllm_config: VllmConfig, layer_names: list[str]
141 -) -> int | None:
142 - """Per-TP-rank ``num_heads`` shared by the named Attention layers.
143 -
144 - Use in metadata builders whose plan-time allocations depend on the
145 - head count: the model-wide ``get_num_attention_heads()`` is wrong
146 - for models with non-uniform per-layer head counts. All layers in
147 - one attention group must agree on ``num_heads``; this is asserted.
148 - Returns ``None`` when no matching Attention layer is found.
149 - """
150 - attn_layers = get_layers_from_vllm_config(
151 - vllm_config,
152 - AttentionLayerBase, # type: ignore[type-abstract]
153 - layer_names,
154 - )
155 - if not attn_layers:
156 - return None
157 - heads = {layer.impl.num_heads for layer in attn_layers.values()}
158 - assert len(heads) == 1, (
159 - f"All layers in one attention group must share num_heads; "
160 - f"got {heads} for {layer_names}."
161 - )
162 - return heads.pop()
163 -
164 -
139 def infer_global_hyperparameters(
140 per_layer_params: dict[str, PerLayerParameters],
141 ) -> PerLayerParameters:
vllm/v1/worker/gpu/model_runner.py
+3 -2
@@ -103,6 +103,7 @@ from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
103 set_eagle3_aux_hidden_state_layers,
104 )
105 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler
106 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
107 from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
108 from vllm.v1.worker.gpu.states import RequestState
109 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
@@ -307,7 +308,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
308 if self.use_aux_hidden_state_outputs:
309 assert self.speculative_config is not None
310 set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
310 - if self.speculator is not None:
311 + if isinstance(self.speculator, DraftModelSpeculator):
312 self.speculator.load_model(self.model)
313 eplb_models_added = self.eplb.maybe_register_speculator(
314 self.speculator, self.speculative_config, load_dummy_weights
@@ -457,7 +458,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
458 self.speculator.init_cudagraph_manager(cudagraph_mode)
459
460 check_attention_cp_compatibility(self.vllm_config)
460 - if self.speculator is not None:
461 + if isinstance(self.speculator, DraftModelSpeculator):
462 # HACK(woosuk)
463 self.speculator.set_attn(
464 self.model_state, self.kv_cache_config, self.block_tables
vllm/v1/worker/gpu/spec_decode/__init__.py
+16 -3
@@ -8,8 +8,21 @@ from vllm.config import VllmConfig
8 def init_speculator(vllm_config: VllmConfig, device: torch.device):
9 speculative_config = vllm_config.speculative_config
10 assert speculative_config is not None
11 - if speculative_config.use_eagle():
12 - from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
11 + if speculative_config.use_gemma4_mtp():
12 + from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
13 + Gemma4Speculator,
14 + )
15 +
16 + return Gemma4Speculator(vllm_config, device)
17 + elif speculative_config.method == "mtp":
18 + from vllm.v1.worker.gpu.spec_decode.mtp.speculator import MTPSpeculator
19 +
20 + return MTPSpeculator(vllm_config, device)
21 + elif speculative_config.use_eagle():
22 + from vllm.v1.worker.gpu.spec_decode.eagle.speculator import (
23 + EagleSpeculator,
24 + )
25
26 return EagleSpeculator(vllm_config, device)
15 - raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
27 + else:
28 + raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
vllm/v1/worker/gpu/spec_decode/autoregressive/__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/autoregressive/cudagraph_utils.py renamed
+4 -5
@@ -19,8 +19,8 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
19 from vllm.v1.worker.utils import AttentionGroup
20
21
22 -class PrefillEagleCudaGraphManager(CudaGraphManager):
23 - """Eagle CudaGraphManager for prefill, using pre-built attention states
22 +class PrefillSpeculatorCudaGraphManager(CudaGraphManager):
23 + """CudaGraphManager for draft prefill, using pre-built attention states
24 from the target model's capture."""
25
26 def capture(
@@ -56,9 +56,8 @@ class PrefillEagleCudaGraphManager(CudaGraphManager):
56 super().capture(create_forward_fn, progress_bar_desc)
57
58
59 -class DecodeEagleCudaGraphManager(CudaGraphManager):
60 - """Eagle CudaGraphManager for decode draft generation, building its own
61 - attention metadata from scratch."""
59 +class DecodeSpeculatorCudaGraphManager(CudaGraphManager):
60 + """CudaGraphManager for draft decode, building its own attention metadata."""
61
62 def capture(
63 self,
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py new
+795
@@ -0,0 +1,795 @@
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 +
7 +from vllm.config import VllmConfig
8 +from vllm.config.compilation import CUDAGraphMode
9 +from vllm.forward_context import BatchDescriptor, set_forward_context
10 +from vllm.logger import init_logger
11 +from vllm.multimodal import MULTIMODAL_REGISTRY
12 +from vllm.triton_utils import tl, triton
13 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
14 +from vllm.v1.worker.gpu.cudagraph_utils import (
15 + AttentionStatePair,
16 + BatchExecutionDescriptor,
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.sample.gumbel import gumbel_sample
22 +from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
23 + DecodeSpeculatorCudaGraphManager,
24 + PrefillSpeculatorCudaGraphManager,
25 +)
26 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class AutoRegressiveSpeculator(DraftModelSpeculator):
32 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
33 + super().__init__(vllm_config, device)
34 +
35 + self.hidden_states = torch.zeros(
36 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
37 + )
38 + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
39 + self.last_token_indices = torch.zeros(
40 + self.max_num_reqs, dtype=torch.int64, device=device
41 + )
42 +
43 + self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
44 + self.draft_model_config
45 + )
46 + if self.supports_mm_inputs:
47 + self.inputs_embeds = torch.zeros(
48 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
49 + )
50 +
51 + self.prefill_cudagraph_manager: PrefillSpeculatorCudaGraphManager | None = None
52 + self.decode_cudagraph_manager: DecodeSpeculatorCudaGraphManager | None = None
53 +
54 + @property
55 + def advance_draft_positions(self) -> bool:
56 + """
57 + Whether to increment positions and seq_lens between draft steps.
58 +
59 + True for Eagle/standard MTP (each step produces new KV).
60 + False for Gemma4 MTP (Q-only, shares target KV, constant positions).
61 + """
62 + return True
63 +
64 + @property
65 + def model_returns_tuple(self) -> bool:
66 + """
67 + Whether the draft model's forward() returns a tuple.
68 +
69 + True: returns (last_hidden_states, hidden_states) — Eagle, Gemma4 MTP.
70 + False: returns a single tensor used for both — standard MTP (DeepSeek).
71 + """
72 + return True
73 +
74 + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
75 + # Initialize cudagraph manager for draft prefill (draft position 0).
76 + self.prefill_cudagraph_manager = PrefillSpeculatorCudaGraphManager(
77 + self.vllm_config,
78 + self.device,
79 + cudagraph_mode,
80 + self.num_speculative_steps + 1,
81 + )
82 +
83 + # PIECEWISE cudagraphs are not supported for draft decodes.
84 + if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
85 + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
86 + else:
87 + cudagraph_mode = CUDAGraphMode.NONE
88 +
89 + # Initialize cudagraph manager for draft decodes (draft positions > 0).
90 + self.decode_cudagraph_manager = DecodeSpeculatorCudaGraphManager(
91 + self.vllm_config,
92 + self.device,
93 + cudagraph_mode,
94 + decode_query_len=1,
95 + )
96 +
97 + def capture(
98 + self,
99 + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],
100 + ) -> None:
101 + logger.info("Capturing model for speculator...")
102 + # Reset indices to zeros to prevent stale values from prior
103 + # dummy runs to cause out-of-bounds indexing during capture.
104 + self.last_token_indices.zero_()
105 +
106 + # Capture the prefill routine (model forward + compute_logits +
107 + # sample).
108 + # For FULL graphs, the entire routine is recorded as one graph.
109 + # For PIECEWISE, only the model's compiled regions are captured
110 + # and the rest (compute_logits, gumbel_sample) runs eagerly.
111 + assert self.prefill_cudagraph_manager is not None
112 + if self.prefill_cudagraph_manager.use_breakable_cg:
113 + self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
114 + self.prefill_cudagraph_manager.capture(
115 + self._prefill,
116 + attn_states,
117 + progress_bar_desc="Capturing prefill CUDA graphs",
118 + )
119 +
120 + if self.num_speculative_steps == 1:
121 + return
122 +
123 + # Capture the decode draft generation routine (model forward +
124 + # sample + update_draft_inputs) for a single
125 + # step.
126 + assert self.decode_cudagraph_manager is not None
127 + self.decode_cudagraph_manager.capture(
128 + self._generate_draft,
129 + self.model_state,
130 + self.input_buffers,
131 + self.block_tables,
132 + self.attn_groups,
133 + self.kv_cache_config,
134 + progress_bar_desc="Capturing decode CUDA graphs",
135 + )
136 +
137 + @torch.inference_mode()
138 + def propose(
139 + self,
140 + input_batch: InputBatch,
141 + attn_metadata: dict[str, Any],
142 + slot_mappings: dict[str, torch.Tensor],
143 + # [num_tokens, hidden_size]
144 + last_hidden_states: torch.Tensor,
145 + # num_layers x [num_tokens, hidden_size]
146 + aux_hidden_states: list[torch.Tensor] | None,
147 + # [num_reqs]
148 + num_sampled: torch.Tensor,
149 + # [num_reqs]
150 + num_rejected: torch.Tensor,
151 + # [max_num_reqs]
152 + last_sampled: torch.Tensor,
153 + # [max_num_reqs]
154 + next_prefill_tokens: torch.Tensor,
155 + # [max_num_reqs]
156 + temperature: torch.Tensor,
157 + # [max_num_reqs]
158 + seeds: torch.Tensor,
159 + num_tokens_across_dp: torch.Tensor | None = None,
160 + dummy_run: bool = False,
161 + skip_attn_for_dummy_run: bool = False,
162 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
163 + is_profile: bool = False,
164 + ) -> torch.Tensor:
165 + num_tokens = input_batch.num_tokens_after_padding
166 + num_reqs = input_batch.num_reqs
167 + max_query_len = input_batch.num_scheduled_tokens.max()
168 + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
169 + self.draft_max_seq_len = min(
170 + max_seq_len + self.num_speculative_steps, self.max_model_len
171 + )
172 +
173 + # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
174 + # number of rejected tokens, we maintain the size of input_ids and
175 + # hidden_states the same as the target model's. This means, we pad each
176 + # request's query length to include any rejected positions. By doing so,
177 + # we can also reuse the attention metadata (e.g., query_start_loc,
178 + # seq_lens) of the target model.
179 + if aux_hidden_states:
180 + assert self.method == "eagle3"
181 + hidden_states = self.model.combine_hidden_states(
182 + torch.cat(aux_hidden_states, dim=-1)
183 + )
184 + else:
185 + hidden_states = last_hidden_states
186 + self.hidden_states[:num_tokens].copy_(hidden_states)
187 +
188 + self._copy_request_inputs(
189 + num_reqs,
190 + input_batch.idx_mapping,
191 + temperature,
192 + seeds,
193 + )
194 +
195 + # Get the input ids and last token indices for the speculator.
196 + prepare_prefill_inputs(
197 + self.last_token_indices,
198 + self.current_draft_step,
199 + self.input_buffers,
200 + input_batch,
201 + num_sampled,
202 + num_rejected,
203 + last_sampled,
204 + next_prefill_tokens,
205 + self.max_num_reqs,
206 + )
207 +
208 + # When all requests are decoding (no true prefills), each has
209 + # num_speculative_steps + 1 tokens, enabling FULL graph replay.
210 + uniform_token_count = get_uniform_token_count(
211 + num_reqs,
212 + # Use the actual number of tokens without padding added by
213 + # the target model during FULL cudagraph.
214 + input_batch.num_tokens,
215 + max_query_len,
216 + )
217 + prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
218 + self.prefill_cudagraph_manager,
219 + num_reqs,
220 + num_tokens,
221 + uniform_token_count,
222 + dp_size=self.dp_size,
223 + dp_rank=self.dp_rank,
224 + need_eager=is_profile,
225 + )
226 +
227 + if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
228 + # Replay the full graph for draft prefill.
229 + assert self.prefill_cudagraph_manager is not None
230 + self.prefill_cudagraph_manager.run_fullgraph(prefill_batch_desc)
231 + else:
232 + # The target model's attention metadata and slot mappings
233 + # can directly be used for draft prefill, because of the
234 + # identical batch shape and KV cache layout.
235 + self._prefill(
236 + num_reqs,
237 + prefill_batch_desc.num_tokens,
238 + attn_metadata,
239 + slot_mappings,
240 + num_tokens_across_dp=num_tokens_across_dp,
241 + cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
242 + mm_inputs=mm_inputs,
243 + )
244 +
245 + if self.num_speculative_steps == 1:
246 + # Early exit.
247 + return self.draft_tokens[:num_reqs, :1]
248 +
249 + # Prepare the inputs for the decode steps.
250 + prepare_decode_inputs(
251 + self.draft_tokens[:num_reqs, 0],
252 + input_batch.seq_lens,
253 + num_rejected,
254 + self.input_buffers,
255 + self.max_model_len,
256 + self.max_num_reqs,
257 + advance_draft_positions=self.advance_draft_positions,
258 + )
259 +
260 + # Each request produces exactly 1 token per draft generation step,
261 + # enabling FULL graph replay.
262 + decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
263 + self.decode_cudagraph_manager,
264 + num_reqs,
265 + num_reqs,
266 + uniform_token_count=1,
267 + dp_size=self.dp_size,
268 + dp_rank=self.dp_rank,
269 + need_eager=is_profile,
270 + )
271 +
272 + # Generate the remaining num_speculative_steps - 1 draft tokens.
273 + self._multi_step_decode(
274 + num_reqs,
275 + dummy_run and skip_attn_for_dummy_run,
276 + decode_batch_desc,
277 + num_tokens_across_dp,
278 + )
279 +
280 + return self.draft_tokens[:num_reqs]
281 +
282 + def sample_draft(
283 + self,
284 + hidden_states: torch.Tensor,
285 + positions: torch.Tensor,
286 + idx_mapping: torch.Tensor,
287 + temperature: torch.Tensor,
288 + seeds: torch.Tensor,
289 + draft_step: torch.Tensor,
290 + draft_logits: torch.Tensor | None,
291 + ) -> torch.Tensor:
292 + logits = self.model.compute_logits(hidden_states)
293 + if draft_logits is not None:
294 + # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
295 + # used for draft and target sampling.
296 + return gumbel_sample(
297 + logits,
298 + idx_mapping,
299 + temperature,
300 + seeds,
301 + positions + 1,
302 + apply_temperature=True,
303 + output_processed_logits=draft_logits,
304 + output_processed_logits_col=draft_step,
305 + use_fp64=self.use_fp64_gumbel,
306 + )
307 + else:
308 + return logits.argmax(dim=-1)
309 +
310 + @torch.inference_mode()
311 + def _run_model(
312 + self,
313 + num_tokens: int,
314 + attn_metadata: dict[str, Any] | None,
315 + slot_mappings: dict[str, torch.Tensor] | None,
316 + num_tokens_across_dp: torch.Tensor | None,
317 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
318 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
319 + ) -> tuple[torch.Tensor, torch.Tensor]:
320 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
321 + with set_forward_context(
322 + attn_metadata,
323 + self.vllm_config,
324 + num_tokens=num_tokens,
325 + cudagraph_runtime_mode=cudagraph_runtime_mode,
326 + num_tokens_across_dp=num_tokens_across_dp,
327 + slot_mapping=slot_mappings,
328 + batch_descriptor=batch_descriptor,
329 + ):
330 + inputs_embeds = None
331 + if self.supports_mm_inputs:
332 + # Merge multimodal embeddings with input ids.
333 + mm_embeds, is_mm_embed = mm_inputs or (None, None)
334 + num_input_tokens = (
335 + is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens
336 + )
337 + self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids(
338 + self.input_buffers.input_ids[:num_input_tokens],
339 + multimodal_embeddings=mm_embeds,
340 + is_multimodal=is_mm_embed,
341 + )
342 + inputs_embeds = self.inputs_embeds[:num_tokens]
343 +
344 + model_inputs = dict(
345 + input_ids=self.input_buffers.input_ids[:num_tokens],
346 + positions=self.input_buffers.positions[:num_tokens],
347 + hidden_states=self.hidden_states[:num_tokens],
348 + inputs_embeds=inputs_embeds,
349 + )
350 + if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
351 + # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable),
352 + # chosen inside run_pw_graph.
353 + assert self.prefill_cudagraph_manager is not None
354 + ret_hidden_states = self.prefill_cudagraph_manager.run_pw_graph(
355 + self.model, model_inputs
356 + )
357 + else:
358 + # Eager (NONE): call the raw model directly.
359 + ret_hidden_states = self.model(**model_inputs)
360 + if self.model_returns_tuple:
361 + last_hidden_states, hidden_states = ret_hidden_states
362 + else:
363 + last_hidden_states = ret_hidden_states
364 + hidden_states = ret_hidden_states
365 + return last_hidden_states, hidden_states
366 +
367 + def _prefill(
368 + self,
369 + num_reqs: int,
370 + num_tokens: int,
371 + attn_metadata: dict[str, Any] | None,
372 + slot_mappings: dict[str, torch.Tensor] | None,
373 + num_tokens_across_dp: torch.Tensor | None,
374 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
375 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
376 + ) -> None:
377 + last_token_indices = self.last_token_indices[:num_reqs]
378 + positions = self.input_buffers.positions[last_token_indices]
379 + idx_mapping = self.idx_mapping[:num_reqs]
380 +
381 + last_hidden_states, hidden_states = self._run_model(
382 + num_tokens,
383 + attn_metadata,
384 + slot_mappings,
385 + num_tokens_across_dp=num_tokens_across_dp,
386 + cudagraph_runtime_mode=cudagraph_runtime_mode,
387 + mm_inputs=mm_inputs,
388 + )
389 + sample_hidden_states = last_hidden_states[last_token_indices]
390 +
391 + self.draft_tokens[:num_reqs, 0] = self.sample_draft(
392 + sample_hidden_states,
393 + positions,
394 + idx_mapping,
395 + self.temperature,
396 + self.seeds,
397 + self.current_draft_step,
398 + self.draft_logits,
399 + )
400 + self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
401 + self.input_buffers.positions[:num_reqs] = positions
402 +
403 + def _multi_step_decode(
404 + self,
405 + num_reqs: int,
406 + skip_attn: bool,
407 + batch_desc: BatchExecutionDescriptor,
408 + num_tokens_across_dp: torch.Tensor | None,
409 + ) -> None:
410 + positions = self.input_buffers.positions[:num_reqs]
411 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
412 + idx_mapping = self.idx_mapping[:num_reqs]
413 +
414 + attn_metadata = None
415 + slot_mappings_by_layer = None
416 + for step in range(1, self.num_speculative_steps):
417 + # Rebuild every step when positions advance, or just once
418 + # on the first step when positions are constant (Gemma4 MTP).
419 + if not skip_attn and (self.advance_draft_positions or step == 1):
420 + slot_mappings = self.block_tables.compute_slot_mappings(
421 + idx_mapping,
422 + query_start_loc,
423 + positions,
424 + batch_desc.num_tokens,
425 + )
426 + slot_mappings_by_layer = build_slot_mappings_by_layer(
427 + slot_mappings, self.kv_cache_config
428 + )
429 + attn_metadata = self._build_draft_attn_metadata(
430 + num_reqs=num_reqs,
431 + num_reqs_padded=batch_desc.num_reqs or num_reqs,
432 + num_tokens_padded=batch_desc.num_tokens,
433 + )
434 +
435 + # Update the current draft step.
436 + self.current_draft_step.fill_(step)
437 +
438 + # Generate draft tokens for the current step.
439 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
440 + assert self.decode_cudagraph_manager is not None
441 + self.decode_cudagraph_manager.run_fullgraph(batch_desc)
442 + else:
443 + self._generate_draft(
444 + num_reqs,
445 + batch_desc.num_tokens,
446 + attn_metadata,
447 + slot_mappings_by_layer,
448 + num_tokens_across_dp=num_tokens_across_dp,
449 + cudagraph_runtime_mode=batch_desc.cg_mode,
450 + )
451 +
452 + def _generate_draft(
453 + self,
454 + num_reqs: int,
455 + num_tokens_padded: int,
456 + attn_metadata: dict[str, Any] | None,
457 + slot_mappings: dict[str, torch.Tensor] | None,
458 + num_tokens_across_dp: torch.Tensor | None,
459 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
460 + ) -> None:
461 + idx_mapping = self.idx_mapping[:num_reqs]
462 + positions = self.input_buffers.positions[:num_reqs]
463 + # Run the draft model forward pass.
464 + last_hidden_states, hidden_states = self._run_model(
465 + num_tokens_padded,
466 + attn_metadata,
467 + slot_mappings,
468 + num_tokens_across_dp,
469 + cudagraph_runtime_mode,
470 + )
471 + last_hidden_states = last_hidden_states[:num_reqs]
472 +
473 + # Sample the draft tokens.
474 + draft_tokens = self.sample_draft(
475 + last_hidden_states,
476 + positions,
477 + idx_mapping,
478 + self.temperature,
479 + self.seeds,
480 + self.current_draft_step,
481 + self.draft_logits,
482 + )
483 +
484 + # Update the inputs for the next step.
485 + update_draft_inputs(
486 + draft_tokens,
487 + self.current_draft_step,
488 + hidden_states,
489 + self.draft_tokens,
490 + self.hidden_states,
491 + self.input_buffers,
492 + num_reqs,
493 + self.max_model_len,
494 + self.num_speculative_steps,
495 + advance_draft_positions=self.advance_draft_positions,
496 + )
497 +
498 +
499 +@triton.jit
500 +def _prepare_prefill_inputs_kernel(
501 + last_token_indices_ptr,
502 + draft_current_step_ptr,
503 + draft_input_ids_ptr,
504 + draft_positions_ptr,
505 + draft_query_start_loc_ptr,
506 + draft_seq_lens_ptr,
507 + target_input_ids_ptr,
508 + target_positions_ptr,
509 + idx_mapping_ptr,
510 + last_sampled_ptr,
511 + next_prefill_tokens_ptr,
512 + num_sampled_ptr,
513 + num_rejected_ptr,
514 + query_start_loc_ptr,
515 + seq_lens_ptr,
516 + max_num_reqs,
517 + BLOCK_SIZE: tl.constexpr,
518 +):
519 + req_idx = tl.program_id(0)
520 + num_reqs = tl.num_programs(0)
521 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
522 +
523 + query_start = tl.load(query_start_loc_ptr + req_idx)
524 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
525 + query_len = query_end - query_start
526 + seq_len = tl.load(seq_lens_ptr + req_idx)
527 +
528 + # Get the true query length and next token after accounting for rejected tokens.
529 + num_rejected = tl.load(num_rejected_ptr + req_idx)
530 + query_len -= num_rejected
531 +
532 + num_sampled = tl.load(num_sampled_ptr + req_idx)
533 + if num_sampled > 0:
534 + next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
535 + else:
536 + # Chunked prefilling.
537 + # Get the next prefill token.
538 + next_token = tl.load(next_prefill_tokens_ptr + req_state_idx)
539 +
540 + # Shift target_input_ids by one.
541 + for i in range(1, query_len, BLOCK_SIZE):
542 + block = i + tl.arange(0, BLOCK_SIZE)
543 + mask = block < query_len
544 + input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask)
545 + tl.store(draft_input_ids_ptr + query_start + block - 1, input_ids, mask=mask)
546 +
547 + last_token_index = query_start + query_len - 1
548 + tl.store(last_token_indices_ptr + req_idx, last_token_index)
549 + tl.store(draft_input_ids_ptr + last_token_index, next_token)
550 +
551 + # Copy positions.
552 + for i in range(0, query_len, BLOCK_SIZE):
553 + block = i + tl.arange(0, BLOCK_SIZE)
554 + mask = block < query_len
555 + target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask)
556 + tl.store(draft_positions_ptr + query_start + block, target_pos, mask=mask)
557 +
558 + # Copy query start locations.
559 + tl.store(draft_query_start_loc_ptr + req_idx, query_start)
560 + # Copy sequence lengths.
561 + tl.store(draft_seq_lens_ptr + req_idx, seq_len)
562 + if req_idx == (num_reqs - 1):
563 + # Reset the current draft step to 0.
564 + tl.store(draft_current_step_ptr, 0)
565 + # Pad query_start_loc for CUDA graphs.
566 + for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
567 + block = i + tl.arange(0, BLOCK_SIZE)
568 + mask = block < max_num_reqs + 1
569 + tl.store(draft_query_start_loc_ptr + block, query_end, mask=mask)
570 + # Pad seq_lens for CUDA graphs.
571 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
572 + block = i + tl.arange(0, BLOCK_SIZE)
573 + mask = block < max_num_reqs
574 + tl.store(draft_seq_lens_ptr + block, 0, mask=mask)
575 + # Pad last_token_indices for CUDA graphs.
576 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
577 + block = i + tl.arange(0, BLOCK_SIZE)
578 + mask = block < max_num_reqs
579 + tl.store(last_token_indices_ptr + block, 0, mask=mask)
580 +
581 +
582 +def prepare_prefill_inputs(
583 + # [num_reqs]
584 + last_token_indices: torch.Tensor,
585 + current_draft_step: torch.Tensor,
586 + input_buffers: InputBuffers,
587 + input_batch: InputBatch,
588 + # [num_reqs]
589 + num_sampled: torch.Tensor,
590 + # [num_reqs]
591 + num_rejected: torch.Tensor,
592 + # [max_num_reqs]
593 + last_sampled: torch.Tensor,
594 + # [max_num_reqs]
595 + next_prefill_tokens: torch.Tensor,
596 + max_num_reqs,
597 +) -> torch.Tensor:
598 + num_reqs = input_batch.num_reqs
599 + _prepare_prefill_inputs_kernel[(num_reqs,)](
600 + last_token_indices,
601 + current_draft_step,
602 + input_buffers.input_ids,
603 + input_buffers.positions,
604 + input_buffers.query_start_loc,
605 + input_buffers.seq_lens,
606 + input_batch.input_ids,
607 + input_batch.positions,
608 + input_batch.idx_mapping,
609 + last_sampled,
610 + next_prefill_tokens,
611 + num_sampled,
612 + num_rejected,
613 + input_batch.query_start_loc,
614 + input_batch.seq_lens,
615 + max_num_reqs,
616 + BLOCK_SIZE=1024,
617 + )
618 + return last_token_indices
619 +
620 +
621 +@triton.jit
622 +def _prepare_decode_inputs_kernel(
623 + draft_tokens_ptr,
624 + draft_tokens_stride,
625 + target_seq_lens_ptr,
626 + num_rejected_ptr,
627 + input_ids_ptr,
628 + positions_ptr,
629 + query_start_loc_ptr,
630 + seq_lens_ptr,
631 + max_model_len,
632 + max_num_reqs,
633 + BLOCK_SIZE: tl.constexpr,
634 + ADVANCE_DRAFT_POSITIONS: tl.constexpr,
635 +):
636 + req_idx = tl.program_id(0)
637 + num_reqs = tl.num_programs(0) - 1
638 + if req_idx == num_reqs:
639 + # Compute query_start_loc. Pad it with the last query_start_loc
640 + # for CUDA graphs.
641 + for i in range(0, max_num_reqs + 1, BLOCK_SIZE):
642 + block = i + tl.arange(0, BLOCK_SIZE)
643 + q = tl.where(block < num_reqs, block, num_reqs)
644 + mask = block < max_num_reqs + 1
645 + tl.store(query_start_loc_ptr + block, q, mask=mask)
646 + # Pad seq_lens for CUDA graphs.
647 + for i in range(req_idx, max_num_reqs, BLOCK_SIZE):
648 + block = i + tl.arange(0, BLOCK_SIZE)
649 + mask = block < max_num_reqs
650 + tl.store(seq_lens_ptr + block, 0, mask=mask)
651 + return
652 +
653 + # draft token -> input id.
654 + draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
655 + tl.store(input_ids_ptr + req_idx, draft_token)
656 +
657 + if ADVANCE_DRAFT_POSITIONS:
658 + # Compute position and seq_lens.
659 + # NOTE(woosuk): To prevent out-of-range access, we clamp these values
660 + # if they reach the max model length.
661 + position = tl.load(positions_ptr + req_idx)
662 + position = tl.minimum(position + 1, max_model_len - 1)
663 + tl.store(positions_ptr + req_idx, position)
664 +
665 + target_seq_len = tl.load(target_seq_lens_ptr + req_idx)
666 + num_rejected = tl.load(num_rejected_ptr + req_idx)
667 + seq_len = target_seq_len - num_rejected
668 + seq_len = tl.minimum(seq_len + 1, max_model_len)
669 + tl.store(seq_lens_ptr + req_idx, seq_len)
670 +
671 +
672 +def prepare_decode_inputs(
673 + draft_tokens: torch.Tensor,
674 + target_seq_lens: torch.Tensor,
675 + num_rejected: torch.Tensor,
676 + input_buffers: InputBuffers,
677 + max_model_len: int,
678 + max_num_reqs: int,
679 + advance_draft_positions: bool = True,
680 +):
681 + num_reqs = draft_tokens.shape[0]
682 + _prepare_decode_inputs_kernel[(num_reqs + 1,)](
683 + draft_tokens,
684 + draft_tokens.stride(0),
685 + target_seq_lens,
686 + num_rejected,
687 + input_buffers.input_ids,
688 + input_buffers.positions,
689 + input_buffers.query_start_loc,
690 + input_buffers.seq_lens,
691 + max_model_len,
692 + max_num_reqs,
693 + BLOCK_SIZE=1024,
694 + ADVANCE_DRAFT_POSITIONS=advance_draft_positions,
695 + )
696 +
697 +
698 +@triton.jit
699 +def _update_draft_inputs_kernel(
700 + output_draft_tokens_ptr,
701 + output_draft_tokens_stride,
702 + next_input_hidden_states_ptr,
703 + next_input_hidden_states_stride,
704 + input_ids_ptr,
705 + positions_ptr,
706 + seq_lens_ptr,
707 + draft_tokens_ptr,
708 + current_draft_step_ptr,
709 + hidden_states_ptr,
710 + hidden_states_stride,
711 + hidden_size,
712 + max_model_len,
713 + num_speculative_steps,
714 + BLOCK_SIZE: tl.constexpr,
715 + ADVANCE_DRAFT_POSITIONS: tl.constexpr,
716 +):
717 + req_idx = tl.program_id(0)
718 +
719 + # Write the sampled draft token into self.draft_tokens[req_idx, step].
720 + draft_token = tl.load(draft_tokens_ptr + req_idx)
721 + step = tl.load(current_draft_step_ptr)
722 + tl.store(
723 + output_draft_tokens_ptr + req_idx * output_draft_tokens_stride + step,
724 + draft_token,
725 + )
726 +
727 + if step >= num_speculative_steps - 1:
728 + # This is the final step. Skip updating draft forward inputs.
729 + return
730 +
731 + # Write the sampled draft token into the input ids tensor for the next
732 + # forward pass.
733 + tl.store(input_ids_ptr + req_idx, draft_token)
734 +
735 + # Copy hidden states into the input hidden states tensor for the next
736 + # forward pass.
737 + for i in range(0, hidden_size, BLOCK_SIZE):
738 + block = i + tl.arange(0, BLOCK_SIZE)
739 + mask = block < hidden_size
740 + hidden_states = tl.load(
741 + hidden_states_ptr + req_idx * hidden_states_stride + block,
742 + mask=mask,
743 + )
744 + tl.store(
745 + next_input_hidden_states_ptr
746 + + req_idx * next_input_hidden_states_stride
747 + + block,
748 + hidden_states,
749 + mask=mask,
750 + )
751 +
752 + if ADVANCE_DRAFT_POSITIONS:
753 + # Increment position and seq_lens.
754 + # NOTE(woosuk): To prevent out-of-range access, we clamp these values
755 + # if they reach the max model length.
756 + position = tl.load(positions_ptr + req_idx)
757 + position = tl.minimum(position + 1, max_model_len - 1)
758 + tl.store(positions_ptr + req_idx, position)
759 +
760 + seq_len = tl.load(seq_lens_ptr + req_idx)
761 + seq_len = tl.minimum(seq_len + 1, max_model_len)
762 + tl.store(seq_lens_ptr + req_idx, seq_len)
763 +
764 +
765 +def update_draft_inputs(
766 + draft_tokens: torch.Tensor,
767 + current_draft_step: torch.Tensor,
768 + hidden_states: torch.Tensor,
769 + output_draft_tokens: torch.Tensor,
770 + next_input_hidden_states: torch.Tensor,
771 + input_buffers: InputBuffers,
772 + num_reqs: int,
773 + max_model_len: int,
774 + num_speculative_steps: int,
775 + advance_draft_positions: bool = True,
776 +):
777 + _, hidden_size = hidden_states.shape
778 + _update_draft_inputs_kernel[(num_reqs,)](
779 + output_draft_tokens,
780 + output_draft_tokens.stride(0),
781 + next_input_hidden_states,
782 + next_input_hidden_states.stride(0),
783 + input_buffers.input_ids,
784 + input_buffers.positions,
785 + input_buffers.seq_lens,
786 + draft_tokens,
787 + current_draft_step,
788 + hidden_states,
789 + hidden_states.stride(0),
790 + hidden_size,
791 + max_model_len,
792 + num_speculative_steps,
793 + BLOCK_SIZE=1024,
794 + ADVANCE_DRAFT_POSITIONS=advance_draft_positions,
795 + )
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+8 -893
@@ -1,903 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Any
3
5 -import torch
4 import torch.nn as nn
5
8 -from vllm.config import VllmConfig, get_layers_from_vllm_config
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.model_executor.layers.attention_layer_base import AttentionLayerBase
13 -from vllm.multimodal import MULTIMODAL_REGISTRY
14 -from vllm.triton_utils import tl, triton
15 -from vllm.v1.kv_cache_interface import KVCacheConfig
16 -from vllm.v1.worker.gpu.attn_utils import (
17 - build_attn_metadata,
18 - build_slot_mappings_by_layer,
19 - init_attn_backend,
20 -)
21 -from vllm.v1.worker.gpu.block_table import BlockTables
22 -from vllm.v1.worker.gpu.cudagraph_utils import (
23 - AttentionStatePair,
24 - BatchExecutionDescriptor,
25 - get_uniform_token_count,
26 -)
27 -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
28 -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
29 -from vllm.v1.worker.gpu.model_states.interface import ModelState
30 -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
31 -from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import (
32 - DecodeEagleCudaGraphManager,
33 - PrefillEagleCudaGraphManager,
6 +from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
7 + AutoRegressiveSpeculator,
8 )
9 from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
10
37 -logger = init_logger(__name__)
38 -
39 -
40 -class EagleSpeculator:
41 - def __init__(self, vllm_config: VllmConfig, device: torch.device):
42 - self.vllm_config = vllm_config
43 - self.device = device
44 -
45 - self.speculative_config = vllm_config.speculative_config
46 - assert self.speculative_config is not None
47 - self.method = self.speculative_config.method
48 - self.num_speculative_steps = self.speculative_config.num_speculative_tokens
49 - self.draft_model_config = self.speculative_config.draft_model_config
50 -
51 - self.scheduler_config = vllm_config.scheduler_config
52 - self.max_num_reqs = self.scheduler_config.max_num_seqs
53 - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
54 - self.max_model_len = vllm_config.model_config.max_model_len
55 - self.draft_max_seq_len = self.max_model_len
56 - # We need to get the hidden size from the draft model config because
57 - # the draft model's hidden size can be different from the target model's
58 - # hidden size (e.g., Llama 3.3 70B).
59 - self.hidden_size = self.draft_model_config.get_hidden_size()
60 - # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP
61 - # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual).
62 - # Non-HC models default to hc_mult=1 and are unaffected.
63 - hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
64 - self.hidden_size = self.hidden_size * hc_mult
65 - self.vocab_size = self.draft_model_config.get_vocab_size()
66 - self.dtype = vllm_config.model_config.dtype
67 - self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
68 -
69 - # DP configuration
70 - self.dp_size = vllm_config.parallel_config.data_parallel_size
71 - self.dp_rank = vllm_config.parallel_config.data_parallel_rank
72 -
73 - self.input_buffers = InputBuffers(
74 - max_num_reqs=self.max_num_reqs,
75 - max_num_tokens=self.max_num_tokens,
76 - device=device,
77 - )
78 - self.hidden_states = torch.zeros(
79 - self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
80 - )
81 - self.idx_mapping = torch.zeros(
82 - self.max_num_reqs, dtype=torch.int32, device=device
83 - )
84 - self.temperature = torch.zeros(
85 - self.max_num_reqs, dtype=torch.float32, device=device
86 - )
87 - self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device)
88 - self.draft_tokens = torch.zeros(
89 - self.max_num_reqs,
90 - self.num_speculative_steps,
91 - dtype=torch.int64,
92 - device=device,
93 - )
94 - self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
95 - self.last_token_indices = torch.zeros(
96 - self.max_num_reqs, dtype=torch.int64, device=device
97 - )
98 - self.arange = torch.arange(
99 - self.max_num_reqs + 1, dtype=torch.int32, device="cpu"
100 - )
101 -
102 - self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
103 - self.draft_model_config
104 - )
105 - if self.supports_mm_inputs:
106 - self.inputs_embeds = torch.zeros(
107 - self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
108 - )
109 -
110 - self.draft_logits: torch.Tensor | None = None
111 - if self.speculative_config.draft_sample_method == "probabilistic":
112 - self.draft_logits = torch.zeros(
113 - self.max_num_reqs,
114 - self.num_speculative_steps,
115 - self.vocab_size,
116 - dtype=torch.float32,
117 - device=device,
118 - )
119 -
120 - self.prefill_cudagraph_manager: PrefillEagleCudaGraphManager | None = None
121 - self.decode_cudagraph_manager: DecodeEagleCudaGraphManager | None = None
122 -
123 - def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
124 - cudagraph_mode = self.vllm_config.compilation_config.cudagraph_mode
125 - # Initialize cudagraph manager for draft prefill (draft position 0).
126 - self.prefill_cudagraph_manager = PrefillEagleCudaGraphManager(
127 - self.vllm_config,
128 - self.device,
129 - cudagraph_mode,
130 - self.num_speculative_steps + 1,
131 - )
132 -
133 - # PIECEWISE cudagraphs are not supported for eagle draft decodes.
134 - # PIECEWISE pads num_tokens to the next capture size without padding
135 - # num_reqs, which can cause attention backends to read past the
136 - # valid per-request metadata (e.g. FlashInfer's kv_indptr buffer).
137 - if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
138 - cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
139 - else:
140 - cudagraph_mode = CUDAGraphMode.NONE
141 -
142 - # Initialize cudagraph manager for draft decodes (draft positions > 0).
143 - self.decode_cudagraph_manager = DecodeEagleCudaGraphManager(
144 - self.vllm_config,
145 - self.device,
146 - cudagraph_mode,
147 - decode_query_len=1,
148 - )
149 -
150 - def load_model(self, target_model: nn.Module) -> None:
151 - target_attn_layer_names = get_layers_from_vllm_config(
152 - self.vllm_config,
153 - AttentionLayerBase, # type: ignore[type-abstract]
154 - ).keys()
155 -
156 - self.model = load_eagle_model(target_model, self.vllm_config)
157 -
158 - all_attn_layers = get_layers_from_vllm_config(
159 - self.vllm_config,
160 - AttentionLayerBase, # type: ignore[type-abstract]
161 - ).keys()
162 - self.draft_attn_layer_names = set(all_attn_layers) - set(
163 - target_attn_layer_names
164 - )
165 -
166 - def set_attn(
167 - self,
168 - model_state: ModelState,
169 - kv_cache_config: KVCacheConfig,
170 - block_tables: BlockTables,
171 - ) -> None:
172 - self.model_state = model_state
173 - self.kv_cache_config = kv_cache_config
174 - self.attn_groups, _, _ = init_attn_backend(
175 - kv_cache_config,
176 - self.vllm_config,
177 - self.device,
178 - active_layer_names=self.draft_attn_layer_names,
179 - )
180 - self.block_tables = block_tables
181 -
182 - @torch.inference_mode()
183 - def run_model(
184 - self,
185 - num_tokens: int,
186 - attn_metadata: dict[str, Any] | None,
187 - slot_mappings: dict[str, torch.Tensor] | None,
188 - num_tokens_across_dp: torch.Tensor | None,
189 - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
190 - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
191 - ) -> tuple[torch.Tensor, torch.Tensor]:
192 - batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
193 - with set_forward_context(
194 - attn_metadata,
195 - self.vllm_config,
196 - num_tokens=num_tokens,
197 - cudagraph_runtime_mode=cudagraph_runtime_mode,
198 - num_tokens_across_dp=num_tokens_across_dp,
199 - slot_mapping=slot_mappings,
200 - batch_descriptor=batch_descriptor,
201 - ):
202 - inputs_embeds = None
203 - if self.supports_mm_inputs:
204 - # Merge multimodal embeddings with input ids.
205 - mm_embeds, is_mm_embed = mm_inputs or (None, None)
206 - num_input_tokens = (
207 - is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens
208 - )
209 - self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids(
210 - self.input_buffers.input_ids[:num_input_tokens],
211 - multimodal_embeddings=mm_embeds,
212 - is_multimodal=is_mm_embed,
213 - )
214 - inputs_embeds = self.inputs_embeds[:num_tokens]
215 -
216 - model_inputs = dict(
217 - input_ids=self.input_buffers.input_ids[:num_tokens],
218 - positions=self.input_buffers.positions[:num_tokens],
219 - hidden_states=self.hidden_states[:num_tokens],
220 - inputs_embeds=inputs_embeds,
221 - )
222 - if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
223 - # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable),
224 - # chosen inside run_pw_graph.
225 - assert self.prefill_cudagraph_manager is not None
226 - ret_hidden_states = self.prefill_cudagraph_manager.run_pw_graph(
227 - self.model, model_inputs
228 - )
229 - else:
230 - # Eager (NONE): call the raw model directly.
231 - ret_hidden_states = self.model(**model_inputs)
232 - if self.method == "mtp":
233 - last_hidden_states = ret_hidden_states
234 - hidden_states = ret_hidden_states
235 - else:
236 - last_hidden_states, hidden_states = ret_hidden_states
237 - return last_hidden_states, hidden_states
238 -
239 - def _sample_draft(
240 - self,
241 - logits: torch.Tensor,
242 - idx_mapping: torch.Tensor,
243 - pos: torch.Tensor,
244 - draft_step: torch.Tensor,
245 - draft_logits: torch.Tensor | None,
246 - ) -> torch.Tensor:
247 - if draft_logits is not None:
248 - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
249 - # used for draft and target sampling.
250 - return gumbel_sample(
251 - logits,
252 - idx_mapping,
253 - self.temperature,
254 - self.seeds,
255 - pos + 1,
256 - apply_temperature=True,
257 - output_processed_logits=draft_logits,
258 - output_processed_logits_col=draft_step,
259 - use_fp64=self.use_fp64_gumbel,
260 - )
261 - else:
262 - return logits.argmax(dim=-1)
11
264 - def prefill(
12 +class EagleSpeculator(AutoRegressiveSpeculator):
13 + def load_draft_model(
14 self,
266 - num_reqs: int,
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 - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
272 - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
273 - ) -> None:
274 - last_token_indices = self.last_token_indices[:num_reqs]
275 - pos = self.input_buffers.positions[last_token_indices]
276 - idx_mapping = self.idx_mapping[:num_reqs]
277 -
278 - last_hidden_states, hidden_states = self.run_model(
279 - num_tokens,
280 - attn_metadata,
281 - slot_mappings,
282 - num_tokens_across_dp=num_tokens_across_dp,
283 - cudagraph_runtime_mode=cudagraph_runtime_mode,
284 - mm_inputs=mm_inputs,
285 - )
286 - sample_hidden_states = last_hidden_states[last_token_indices]
287 - logits = self.model.compute_logits(sample_hidden_states)
288 -
289 - self.draft_tokens[:num_reqs, 0] = self._sample_draft(
290 - logits,
291 - idx_mapping,
292 - pos,
293 - self.current_draft_step,
294 - self.draft_logits,
295 - )
296 - self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
297 - self.input_buffers.positions[:num_reqs] = pos
298 -
299 - def multi_step_decode(
300 - self,
301 - num_reqs: int,
302 - skip_attn: bool,
303 - batch_desc: BatchExecutionDescriptor,
304 - num_tokens_across_dp: torch.Tensor | None,
305 - ) -> None:
306 - positions = self.input_buffers.positions[:num_reqs]
307 - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
308 - idx_mapping = self.idx_mapping[:num_reqs]
309 -
310 - for step in range(1, self.num_speculative_steps):
311 - attn_metadata = None
312 - slot_mappings_by_layer = None
313 - if not skip_attn:
314 - # Build attention metadata and slot mappings for each draft
315 - # decode step. It is necessary to rebuild the attention
316 - # metadata even when replaying the FULL graph so that any
317 - # attention metadata builder state is updated.
318 - slot_mappings = self.block_tables.compute_slot_mappings(
319 - idx_mapping,
320 - query_start_loc,
321 - positions,
322 - batch_desc.num_tokens,
323 - )
324 - slot_mappings_by_layer = build_slot_mappings_by_layer(
325 - slot_mappings, self.kv_cache_config
326 - )
327 - attn_metadata = self._build_draft_attn_metadata(
328 - num_reqs=num_reqs,
329 - num_reqs_padded=batch_desc.num_reqs or num_reqs,
330 - num_tokens_padded=batch_desc.num_tokens,
331 - )
332 -
333 - # Update the current draft step.
334 - self.current_draft_step.fill_(step)
335 -
336 - # Generate draft tokens for the current step.
337 - if batch_desc.cg_mode == CUDAGraphMode.FULL:
338 - assert self.decode_cudagraph_manager is not None
339 - self.decode_cudagraph_manager.run_fullgraph(batch_desc)
340 - else:
341 - self.generate_draft(
342 - num_reqs,
343 - batch_desc.num_tokens,
344 - attn_metadata,
345 - slot_mappings_by_layer,
346 - num_tokens_across_dp=num_tokens_across_dp,
347 - cudagraph_runtime_mode=batch_desc.cg_mode,
348 - )
349 -
350 - def generate_draft(
351 - self,
352 - num_reqs: int,
353 - num_tokens_padded: int,
354 - attn_metadata: dict[str, Any] | None,
355 - slot_mappings: dict[str, torch.Tensor] | None,
356 - num_tokens_across_dp: torch.Tensor | None,
357 - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
358 - ) -> None:
359 - idx_mapping = self.idx_mapping[:num_reqs]
360 - positions = self.input_buffers.positions[:num_reqs]
361 - # Run the eagle model forward pass.
362 - last_hidden_states, hidden_states = self.run_model(
363 - num_tokens_padded,
364 - attn_metadata,
365 - slot_mappings,
366 - num_tokens_across_dp,
367 - cudagraph_runtime_mode,
368 - )
369 - last_hidden_states = last_hidden_states[:num_reqs]
370 -
371 - # Sample the draft tokens.
372 - logits = self.model.compute_logits(last_hidden_states)
373 - draft_tokens = self._sample_draft(
374 - logits,
375 - idx_mapping,
376 - positions,
377 - self.current_draft_step,
378 - self.draft_logits,
379 - )
380 -
381 - # Update the inputs for the next step.
382 - update_eagle_draft_inputs(
383 - draft_tokens,
384 - self.current_draft_step,
385 - hidden_states,
386 - self.draft_tokens,
387 - self.hidden_states,
388 - self.input_buffers,
389 - num_reqs,
390 - self.max_model_len,
391 - self.num_speculative_steps,
392 - )
393 -
394 - def _build_draft_attn_metadata(
395 - self,
396 - num_reqs: int,
397 - num_reqs_padded: int,
398 - num_tokens_padded: int,
399 - ) -> dict[str, Any] | None:
400 - if not self.draft_attn_layer_names:
401 - return None
402 -
403 - query_start_loc_cpu = torch.clamp(
404 - self.arange[: num_reqs_padded + 1], max=num_reqs
405 - )
406 - block_tables = [
407 - x[:num_reqs_padded] for x in self.block_tables.input_block_tables
408 - ]
409 - slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
410 - attn_metadata = build_attn_metadata(
411 - attn_groups=self.attn_groups,
412 - num_reqs=num_reqs_padded,
413 - num_tokens=num_tokens_padded,
414 - query_start_loc_gpu=self.input_buffers.query_start_loc[
415 - : num_reqs_padded + 1
416 - ],
417 - query_start_loc_cpu=query_start_loc_cpu,
418 - max_query_len=1,
419 - seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
420 - max_seq_len=self.draft_max_seq_len,
421 - block_tables=block_tables,
422 - slot_mappings=slot_mappings,
423 - kv_cache_config=self.kv_cache_config,
424 - )
425 - return attn_metadata
426 -
427 - def capture(
428 - self,
429 - attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],
430 - ) -> None:
431 - logger.info("Capturing model for Eagle speculator...")
432 - # Reset indices to zeros to prevent stale values from prior
433 - # dummy runs to cause out-of-bounds indexing during capture.
434 - self.last_token_indices.zero_()
435 -
436 - # Capture the prefill routine (model forward + compute_logits +
437 - # sample).
438 - # For FULL graphs, the entire routine is recorded as one graph.
439 - # For PIECEWISE, only the model's compiled regions are captured
440 - # and the rest (compute_logits, gumbel_sample) runs eagerly.
441 - assert self.prefill_cudagraph_manager is not None
442 - if self.prefill_cudagraph_manager.use_breakable_cg:
443 - self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
444 - self.prefill_cudagraph_manager.capture(
445 - self.prefill,
446 - attn_states,
447 - progress_bar_desc="Capturing eagle prefill CUDA graphs",
448 - )
449 -
450 - if self.num_speculative_steps == 1:
451 - return
452 -
453 - # Capture the decode draft generation routine (model forward +
454 - # compute_logits + sample + update_eagle_inputs) for a single
455 - # step.
456 - assert self.decode_cudagraph_manager is not None
457 - self.decode_cudagraph_manager.capture(
458 - self.generate_draft,
459 - self.model_state,
460 - self.input_buffers,
461 - self.block_tables,
462 - self.attn_groups,
463 - self.kv_cache_config,
464 - progress_bar_desc="Capturing eagle decode CUDA graphs",
465 - )
466 -
467 - @torch.inference_mode()
468 - def propose(
469 - self,
470 - input_batch: InputBatch,
471 - attn_metadata: dict[str, Any],
472 - slot_mappings: dict[str, torch.Tensor],
473 - # [num_tokens, hidden_size]
474 - last_hidden_states: torch.Tensor,
475 - # num_layers x [num_tokens, hidden_size]
476 - aux_hidden_states: list[torch.Tensor] | None,
477 - # [num_reqs]
478 - num_sampled: torch.Tensor,
479 - # [num_reqs]
480 - num_rejected: torch.Tensor,
481 - # [max_num_reqs]
482 - last_sampled: torch.Tensor,
483 - # [max_num_reqs]
484 - next_prefill_tokens: torch.Tensor,
485 - # [max_num_reqs]
486 - temperature: torch.Tensor,
487 - # [max_num_reqs]
488 - seeds: torch.Tensor,
489 - num_tokens_across_dp: torch.Tensor | None = None,
490 - dummy_run: bool = False,
491 - skip_attn_for_dummy_run: bool = False,
492 - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
493 - is_profile: bool = False,
494 - ) -> torch.Tensor:
495 - num_tokens = input_batch.num_tokens_after_padding
496 - num_reqs = input_batch.num_reqs
497 - max_query_len = input_batch.num_scheduled_tokens.max()
498 - max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
499 - self.draft_max_seq_len = min(
500 - max_seq_len + self.num_speculative_steps, self.max_model_len
501 - )
502 -
503 - # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
504 - # number of rejected tokens, we maintain the size of eagle's input_ids and
505 - # hidden_states the same as the target model's. This means, we pad each
506 - # request's query length to include any rejected positions. By doing so,
507 - # we can also reuse the attention metadata (e.g., query_start_loc,
508 - # seq_lens) of the target model.
509 - if aux_hidden_states:
510 - assert self.method == "eagle3"
511 - hidden_states = self.model.combine_hidden_states(
512 - torch.cat(aux_hidden_states, dim=-1)
513 - )
514 - else:
515 - hidden_states = last_hidden_states
516 - self.hidden_states[:num_tokens].copy_(hidden_states)
517 -
518 - # Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
519 - # NOTE(woosuk): For draft sampling, we only consider the temperature
520 - # and ignore the other sampling parameters such as top_k and top_p,
521 - # for simplicity and performance.
522 - # While this may slightly degrade the acceptance rate, it does not
523 - # affect the output distribution after rejection sampling.
524 - self.temperature.copy_(temperature)
525 - self.seeds.copy_(seeds)
526 - self.idx_mapping[:num_reqs].copy_(input_batch.idx_mapping)
527 -
528 - # Get the input ids and last token indices for the speculator.
529 - prepare_eagle_inputs(
530 - self.last_token_indices,
531 - self.current_draft_step,
532 - self.input_buffers,
533 - input_batch,
534 - num_sampled,
535 - num_rejected,
536 - last_sampled,
537 - next_prefill_tokens,
538 - self.max_num_reqs,
539 - )
540 -
541 - # When all requests are decoding (no true prefills), each has
542 - # num_speculative_steps + 1 tokens, enabling FULL graph replay.
543 - uniform_token_count = get_uniform_token_count(
544 - num_reqs,
545 - # Use the actual number of tokens without padding added by
546 - # the target model during FULL cudagraph.
547 - input_batch.num_tokens,
548 - max_query_len,
549 - )
550 - prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
551 - self.prefill_cudagraph_manager,
552 - num_reqs,
553 - num_tokens,
554 - uniform_token_count,
555 - dp_size=self.dp_size,
556 - dp_rank=self.dp_rank,
557 - need_eager=is_profile,
558 - )
559 -
560 - if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
561 - # Replay the full graph for draft prefill.
562 - assert self.prefill_cudagraph_manager is not None
563 - self.prefill_cudagraph_manager.run_fullgraph(prefill_batch_desc)
564 - else:
565 - # The target model's attention metadata and slot mappings
566 - # can directly be used for draft prefill, because of the
567 - # identical batch shape and KV cache layout.
568 - self.prefill(
569 - num_reqs,
570 - prefill_batch_desc.num_tokens,
571 - attn_metadata,
572 - slot_mappings,
573 - num_tokens_across_dp=num_tokens_across_dp,
574 - cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
575 - mm_inputs=mm_inputs,
576 - )
577 -
578 - if self.num_speculative_steps == 1:
579 - # Early exit.
580 - return self.draft_tokens[:num_reqs, :1]
581 -
582 - # Prepare the inputs for the decode steps.
583 - prepare_eagle_decode(
584 - self.draft_tokens[:num_reqs, 0],
585 - input_batch.seq_lens,
586 - num_rejected,
587 - self.input_buffers,
588 - self.max_model_len,
589 - self.max_num_reqs,
590 - )
591 -
592 - # Each request produces exactly 1 token per draft generation step,
593 - # enabling FULL graph replay.
594 - decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
595 - self.decode_cudagraph_manager,
596 - num_reqs,
597 - num_reqs,
598 - uniform_token_count=1,
599 - dp_size=self.dp_size,
600 - dp_rank=self.dp_rank,
601 - need_eager=is_profile,
602 - )
603 -
604 - # Generate the remaining num_speculative_steps - 1 draft tokens.
605 - self.multi_step_decode(
606 - num_reqs,
607 - dummy_run and skip_attn_for_dummy_run,
608 - decode_batch_desc,
609 - num_tokens_across_dp,
610 - )
611 -
612 - return self.draft_tokens[:num_reqs]
613 -
614 -
615 -@triton.jit
616 -def _prepare_eagle_inputs_kernel(
617 - last_token_indices_ptr,
618 - eagle_current_draft_step_ptr,
619 - eagle_input_ids_ptr,
620 - eagle_positions_ptr,
621 - eagle_query_start_loc_ptr,
622 - eagle_seq_lens_ptr,
623 - target_input_ids_ptr,
624 - target_positions_ptr,
625 - idx_mapping_ptr,
626 - last_sampled_ptr,
627 - next_prefill_tokens_ptr,
628 - num_sampled_ptr,
629 - num_rejected_ptr,
630 - query_start_loc_ptr,
631 - seq_lens_ptr,
632 - max_num_reqs,
633 - BLOCK_SIZE: tl.constexpr,
634 -):
635 - req_idx = tl.program_id(0)
636 - num_reqs = tl.num_programs(0)
637 - req_state_idx = tl.load(idx_mapping_ptr + req_idx)
638 -
639 - query_start = tl.load(query_start_loc_ptr + req_idx)
640 - query_end = tl.load(query_start_loc_ptr + req_idx + 1)
641 - query_len = query_end - query_start
642 - seq_len = tl.load(seq_lens_ptr + req_idx)
643 -
644 - # Get the true query length and next token after accounting for rejected tokens.
645 - num_rejected = tl.load(num_rejected_ptr + req_idx)
646 - query_len -= num_rejected
647 -
648 - num_sampled = tl.load(num_sampled_ptr + req_idx)
649 - if num_sampled > 0:
650 - next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
651 - else:
652 - # Chunked prefilling.
653 - # Get the next prefill token.
654 - next_token = tl.load(next_prefill_tokens_ptr + req_state_idx)
655 -
656 - # Shift target_input_ids by one.
657 - for i in range(1, query_len, BLOCK_SIZE):
658 - block = i + tl.arange(0, BLOCK_SIZE)
659 - mask = block < query_len
660 - input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask)
661 - tl.store(eagle_input_ids_ptr + query_start + block - 1, input_ids, mask=mask)
662 -
663 - last_token_index = query_start + query_len - 1
664 - tl.store(last_token_indices_ptr + req_idx, last_token_index)
665 - tl.store(eagle_input_ids_ptr + last_token_index, next_token)
666 -
667 - # Copy positions.
668 - for i in range(0, query_len, BLOCK_SIZE):
669 - block = i + tl.arange(0, BLOCK_SIZE)
670 - mask = block < query_len
671 - target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask)
672 - tl.store(eagle_positions_ptr + query_start + block, target_pos, mask=mask)
673 -
674 - # Copy query start locations.
675 - tl.store(eagle_query_start_loc_ptr + req_idx, query_start)
676 - # Copy sequence lengths.
677 - tl.store(eagle_seq_lens_ptr + req_idx, seq_len)
678 - if req_idx == (num_reqs - 1):
679 - # Reset the current draft step to 0.
680 - tl.store(eagle_current_draft_step_ptr, 0)
681 - # Pad query_start_loc for CUDA graphs.
682 - for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
683 - block = i + tl.arange(0, BLOCK_SIZE)
684 - mask = block < max_num_reqs + 1
685 - tl.store(eagle_query_start_loc_ptr + block, query_end, mask=mask)
686 - # Pad seq_lens for CUDA graphs.
687 - for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
688 - block = i + tl.arange(0, BLOCK_SIZE)
689 - mask = block < max_num_reqs
690 - tl.store(eagle_seq_lens_ptr + block, 0, mask=mask)
691 - # Pad last_token_indices for CUDA graphs.
692 - for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
693 - block = i + tl.arange(0, BLOCK_SIZE)
694 - mask = block < max_num_reqs
695 - tl.store(last_token_indices_ptr + block, 0, mask=mask)
696 -
697 -
698 -def prepare_eagle_inputs(
699 - # [num_reqs]
700 - last_token_indices: torch.Tensor,
701 - current_draft_step: torch.Tensor,
702 - input_buffers: InputBuffers,
703 - input_batch: InputBatch,
704 - # [num_reqs]
705 - num_sampled: torch.Tensor,
706 - # [num_reqs]
707 - num_rejected: torch.Tensor,
708 - # [max_num_reqs]
709 - last_sampled: torch.Tensor,
710 - # [max_num_reqs]
711 - next_prefill_tokens: torch.Tensor,
712 - max_num_reqs,
713 -) -> torch.Tensor:
714 - num_reqs = input_batch.num_reqs
715 - _prepare_eagle_inputs_kernel[(num_reqs,)](
716 - last_token_indices,
717 - current_draft_step,
718 - input_buffers.input_ids,
719 - input_buffers.positions,
720 - input_buffers.query_start_loc,
721 - input_buffers.seq_lens,
722 - input_batch.input_ids,
723 - input_batch.positions,
724 - input_batch.idx_mapping,
725 - last_sampled,
726 - next_prefill_tokens,
727 - num_sampled,
728 - num_rejected,
729 - input_batch.query_start_loc,
730 - input_batch.seq_lens,
731 - max_num_reqs,
732 - BLOCK_SIZE=1024,
733 - )
734 - return last_token_indices
735 -
736 -
737 -@triton.jit
738 -def _prepare_eagle_decode_kernel(
739 - draft_tokens_ptr,
740 - draft_tokens_stride,
741 - target_seq_lens_ptr,
742 - num_rejected_ptr,
743 - input_ids_ptr,
744 - positions_ptr,
745 - query_start_loc_ptr,
746 - seq_lens_ptr,
747 - max_model_len,
748 - max_num_reqs,
749 - BLOCK_SIZE: tl.constexpr,
750 -):
751 - req_idx = tl.program_id(0)
752 - num_reqs = tl.num_programs(0) - 1
753 - if req_idx == num_reqs:
754 - # Compute query_start_loc. Pad it with the last query_start_loc
755 - # for CUDA graphs.
756 - for i in range(0, max_num_reqs + 1, BLOCK_SIZE):
757 - block = i + tl.arange(0, BLOCK_SIZE)
758 - q = tl.where(block < num_reqs, block, num_reqs)
759 - mask = block < max_num_reqs + 1
760 - tl.store(query_start_loc_ptr + block, q, mask=mask)
761 - # Pad seq_lens for CUDA graphs.
762 - for i in range(req_idx, max_num_reqs, BLOCK_SIZE):
763 - block = i + tl.arange(0, BLOCK_SIZE)
764 - mask = block < max_num_reqs
765 - tl.store(seq_lens_ptr + block, 0, mask=mask)
766 - return
767 -
768 - # draft token -> input id.
769 - draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
770 - tl.store(input_ids_ptr + req_idx, draft_token)
771 -
772 - # Compute position and seq_lens.
773 - # NOTE(woosuk): To prevent out-of-range access, we clamp these values
774 - # if they reach the max model length.
775 - position = tl.load(positions_ptr + req_idx)
776 - position = tl.minimum(position + 1, max_model_len - 1)
777 - tl.store(positions_ptr + req_idx, position)
778 -
779 - target_seq_len = tl.load(target_seq_lens_ptr + req_idx)
780 - num_rejected = tl.load(num_rejected_ptr + req_idx)
781 - seq_len = target_seq_len - num_rejected
782 - seq_len = tl.minimum(seq_len + 1, max_model_len)
783 - tl.store(seq_lens_ptr + req_idx, seq_len)
784 -
785 -
786 -def prepare_eagle_decode(
787 - draft_tokens: torch.Tensor,
788 - target_seq_lens: torch.Tensor,
789 - num_rejected: torch.Tensor,
790 - input_buffers: InputBuffers,
791 - max_model_len: int,
792 - max_num_reqs: int,
793 -):
794 - num_reqs = draft_tokens.shape[0]
795 - _prepare_eagle_decode_kernel[(num_reqs + 1,)](
796 - draft_tokens,
797 - draft_tokens.stride(0),
798 - target_seq_lens,
799 - num_rejected,
800 - input_buffers.input_ids,
801 - input_buffers.positions,
802 - input_buffers.query_start_loc,
803 - input_buffers.seq_lens,
804 - max_model_len,
805 - max_num_reqs,
806 - BLOCK_SIZE=1024,
807 - )
808 -
809 -
810 -@triton.jit
811 -def _update_eagle_draft_inputs_kernel(
812 - output_draft_tokens_ptr,
813 - output_draft_tokens_stride,
814 - next_input_hidden_states_ptr,
815 - next_input_hidden_states_stride,
816 - input_ids_ptr,
817 - positions_ptr,
818 - seq_lens_ptr,
819 - draft_tokens_ptr,
820 - current_draft_step_ptr,
821 - hidden_states_ptr,
822 - hidden_states_stride,
823 - hidden_size,
824 - max_model_len,
825 - num_speculative_steps,
826 - BLOCK_SIZE: tl.constexpr,
827 -):
828 - req_idx = tl.program_id(0)
829 -
830 - # Write the sampled draft token into self.draft_tokens[req_idx, step].
831 - draft_token = tl.load(draft_tokens_ptr + req_idx)
832 - step = tl.load(current_draft_step_ptr)
833 - tl.store(
834 - output_draft_tokens_ptr + req_idx * output_draft_tokens_stride + step,
835 - draft_token,
836 - )
837 -
838 - if step >= num_speculative_steps - 1:
839 - # This is the final step. Skip updating draft forward inputs.
840 - return
841 -
842 - # Write the sampled draft token into the input ids tensor for the next
843 - # forward pass.
844 - tl.store(input_ids_ptr + req_idx, draft_token)
845 -
846 - # Copy hidden states into the input hidden states tensor for the next
847 - # forward pass.
848 - for i in range(0, hidden_size, BLOCK_SIZE):
849 - block = i + tl.arange(0, BLOCK_SIZE)
850 - mask = block < hidden_size
851 - hidden_states = tl.load(
852 - hidden_states_ptr + req_idx * hidden_states_stride + block,
853 - mask=mask,
854 - )
855 - tl.store(
856 - next_input_hidden_states_ptr
857 - + req_idx * next_input_hidden_states_stride
858 - + block,
859 - hidden_states,
860 - mask=mask,
861 - )
862 -
863 - # Increment position and seq_lens.
864 - # NOTE(woosuk): To prevent out-of-range access, we clamp these values
865 - # if they reach the max model length.
866 - position = tl.load(positions_ptr + req_idx)
867 - position = tl.minimum(position + 1, max_model_len - 1)
868 - tl.store(positions_ptr + req_idx, position)
869 -
870 - seq_len = tl.load(seq_lens_ptr + req_idx)
871 - seq_len = tl.minimum(seq_len + 1, max_model_len)
872 - tl.store(seq_lens_ptr + req_idx, seq_len)
873 -
874 -
875 -def update_eagle_draft_inputs(
876 - draft_tokens: torch.Tensor,
877 - current_draft_step: torch.Tensor,
878 - hidden_states: torch.Tensor,
879 - output_draft_tokens: torch.Tensor,
880 - next_input_hidden_states: torch.Tensor,
881 - input_buffers: InputBuffers,
882 - num_reqs: int,
883 - max_model_len: int,
884 - num_speculative_steps: int,
885 -):
886 - _, hidden_size = hidden_states.shape
887 - _update_eagle_draft_inputs_kernel[(num_reqs,)](
888 - output_draft_tokens,
889 - output_draft_tokens.stride(0),
890 - next_input_hidden_states,
891 - next_input_hidden_states.stride(0),
892 - input_buffers.input_ids,
893 - input_buffers.positions,
894 - input_buffers.seq_lens,
895 - draft_tokens,
896 - current_draft_step,
897 - hidden_states,
898 - hidden_states.stride(0),
899 - hidden_size,
900 - max_model_len,
901 - num_speculative_steps,
902 - BLOCK_SIZE=1024,
903 - )
15 + target_model: nn.Module,
16 + target_attn_layer_names: set[str],
17 + ) -> nn.Module:
18 + return load_eagle_model(target_model, self.vllm_config)
vllm/v1/worker/gpu/spec_decode/gemma4/__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/gemma4/speculator.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gemma4 MTP (Multi-Token Prediction) speculator for speculative decoding.
4 +
5 +The Gemma4 assistant model runs all decoder layers per draft step
6 +(producing one token), and all its attention layers share KV cache
7 +with the target model via cross-model KV sharing.
8 +"""
9 +
10 +from collections import defaultdict
11 +
12 +import torch.nn as nn
13 +
14 +from vllm.compilation.backends import set_model_tag
15 +from vllm.config import VllmConfig, replace
16 +from vllm.distributed.parallel_state import get_pp_group
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.model_loader import get_model
19 +from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
20 + AutoRegressiveSpeculator,
21 +)
22 +
23 +logger = init_logger(__name__)
24 +
25 +
26 +class Gemma4Speculator(AutoRegressiveSpeculator):
27 + @property
28 + def advance_draft_positions(self) -> bool:
29 + # Gemma4 MTP is Q-only and reads K/V from the target's existing cache.
30 + # No new KV slots are written, so positions and seq_lens stay fixed.
31 + return False
32 +
33 + @property
34 + def model_returns_tuple(self) -> bool:
35 + # forward() returns (draft_hidden_states, backbone_hidden_states).
36 + # The proposer uses draft_hidden_states for compute_logits and
37 + # backbone_hidden_states for the hidden-state feedback buffer.
38 + return True
39 +
40 + def load_draft_model(
41 + self,
42 + target_model: nn.Module,
43 + target_attn_layer_names: set[str],
44 + ) -> nn.Module:
45 + draft_vllm_config = self._create_draft_vllm_config()
46 + with set_model_tag("eagle_head"):
47 + draft_model = get_model(
48 + vllm_config=draft_vllm_config,
49 + model_config=self.speculative_config.draft_model_config,
50 + load_config=self.speculative_config.draft_load_config,
51 + )
52 + self._setup_gemma4_kv_sharing(draft_model, target_attn_layer_names)
53 + self._share_embeddings(draft_model, target_model)
54 + return draft_model
55 +
56 + def _create_draft_vllm_config(self) -> VllmConfig:
57 + """Preserve the target's forced TRITON_ATTN backend for draft layers.
58 +
59 + Gemma4 forces TRITON_ATTN due to heterogeneous head dimensions
60 + (head_dim=256 sliding, global_head_dim=512 full). The base class
61 + resets attention_config.backend to None for draft models, causing
62 + sliding layers to fall back to FLASH_ATTN which cannot handle
63 + KV-shared cache. Override to carry the target's backend through.
64 + """
65 + draft_model_config = self.speculative_config.draft_model_config
66 + draft_vllm_config = replace(
67 + self.vllm_config,
68 + model_config=draft_model_config,
69 + )
70 + target_backend = self.vllm_config.attention_config.backend
71 + if target_backend is not None:
72 + draft_vllm_config = replace(
73 + draft_vllm_config,
74 + attention_config=replace(
75 + draft_vllm_config.attention_config,
76 + backend=target_backend,
77 + ),
78 + )
79 + return draft_vllm_config
80 +
81 + def _setup_gemma4_kv_sharing(
82 + self,
83 + model: nn.Module,
84 + target_attn_layer_names: set[str],
85 + ) -> None:
86 + """Wire draft layers to share KV with the target model.
87 +
88 + Each draft decoder layer is mapped to the last non-KV-shared
89 + target layer of the same attention type (sliding or full).
90 + """
91 + draft_config = self.speculative_config.draft_model_config.hf_config
92 + draft_text_config = draft_config.get_text_config()
93 + target_config = self.vllm_config.model_config.hf_config
94 + target_text_config = target_config.get_text_config()
95 + target_layer_types = getattr(target_text_config, "layer_types", [])
96 +
97 + if not (hasattr(model, "model") and hasattr(model.model, "layers")):
98 + return
99 +
100 + target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
101 + num_non_shared = len(target_layer_types) - target_num_kv_shared
102 + type_to_target_indices: dict[str, list[int]] = defaultdict(list)
103 + for idx, lt in enumerate(target_layer_types[:num_non_shared]):
104 + type_to_target_indices[lt].append(idx)
105 +
106 + target_prefix = "model.layers"
107 + for name in target_attn_layer_names:
108 + if ".layers." in name:
109 + target_prefix = name.split(".layers.")[0] + ".layers"
110 + break
111 +
112 + draft_layer_types = getattr(draft_text_config, "layer_types", [])
113 + for draft_idx, layer in enumerate(model.model.layers):
114 + if not hasattr(layer, "self_attn"):
115 + continue
116 + attn = getattr(layer.self_attn, "attn", None)
117 + if attn is None:
118 + continue
119 +
120 + draft_layer_type = (
121 + draft_layer_types[draft_idx]
122 + if draft_idx < len(draft_layer_types)
123 + else "full_attention"
124 + )
125 + candidates = type_to_target_indices.get(draft_layer_type, [])
126 + if not candidates:
127 + logger.warning(
128 + "No target layer of type '%s' for draft layer %d",
129 + draft_layer_type,
130 + draft_idx,
131 + )
132 + continue
133 +
134 + target_idx = candidates[-1]
135 + target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
136 + attn.kv_sharing_target_layer_name = target_layer_name
137 + logger.info(
138 + "Gemma4 MTP: draft layer %d (%s) -> %s",
139 + draft_idx,
140 + draft_layer_type,
141 + target_layer_name,
142 + )
143 +
144 + def _share_embeddings(
145 + self,
146 + draft_model: nn.Module,
147 + target_model: nn.Module,
148 + ) -> None:
149 + target_language_model = (
150 + target_model.get_language_model()
151 + if hasattr(target_model, "get_language_model")
152 + else target_model
153 + )
154 + if get_pp_group().world_size == 1:
155 + target_embed = getattr(target_language_model.model, "embed_tokens", None)
156 + if target_embed is not None:
157 + del draft_model.model.embed_tokens
158 + draft_model.model.embed_tokens = target_embed
vllm/v1/worker/gpu/spec_decode/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/mtp/speculator.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch.nn as nn
5 +
6 +from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
7 + AutoRegressiveSpeculator,
8 +)
9 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
10 +
11 +
12 +class MTPSpeculator(AutoRegressiveSpeculator):
13 + @property
14 + def model_returns_tuple(self) -> bool:
15 + return False
16 +
17 + def load_draft_model(
18 + self,
19 + target_model: nn.Module,
20 + target_attn_layer_names: set[str],
21 + ) -> nn.Module:
22 + return load_eagle_model(target_model, self.vllm_config)
vllm/v1/worker/gpu/spec_decode/speculator.py new
+224
@@ -0,0 +1,224 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from typing import Any
5 +
6 +import torch
7 +import torch.nn as nn
8 +
9 +from vllm.config import VllmConfig, get_layers_from_vllm_config
10 +from vllm.config.compilation import CUDAGraphMode
11 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12 +from vllm.v1.kv_cache_interface import KVCacheConfig
13 +from vllm.v1.worker.gpu.attn_utils import (
14 + build_attn_metadata,
15 + init_attn_backend,
16 +)
17 +from vllm.v1.worker.gpu.block_table import BlockTables
18 +from vllm.v1.worker.gpu.cudagraph_utils import (
19 + AttentionStatePair,
20 + BatchExecutionDescriptor,
21 +)
22 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
23 +from vllm.v1.worker.gpu.model_states.interface import ModelState
24 +
25 +
26 +class BaseSpeculator(ABC):
27 + @abstractmethod
28 + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
29 + pass
30 +
31 + @abstractmethod
32 + def capture(
33 + self,
34 + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],
35 + ) -> None:
36 + pass
37 +
38 + @abstractmethod
39 + def propose(
40 + self,
41 + input_batch: InputBatch,
42 + attn_metadata: dict[str, Any],
43 + slot_mappings: dict[str, torch.Tensor],
44 + # [num_tokens, hidden_size]
45 + last_hidden_states: torch.Tensor,
46 + # num_layers x [num_tokens, hidden_size]
47 + aux_hidden_states: list[torch.Tensor] | None,
48 + # [num_reqs]
49 + num_sampled: torch.Tensor,
50 + # [num_reqs]
51 + num_rejected: torch.Tensor,
52 + # [max_num_reqs]
53 + last_sampled: torch.Tensor,
54 + # [max_num_reqs]
55 + next_prefill_tokens: torch.Tensor,
56 + # [max_num_reqs]
57 + temperature: torch.Tensor,
58 + # [max_num_reqs]
59 + seeds: torch.Tensor,
60 + num_tokens_across_dp: torch.Tensor | None = None,
61 + dummy_run: bool = False,
62 + skip_attn_for_dummy_run: bool = False,
63 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
64 + is_profile: bool = False,
65 + ) -> torch.Tensor:
66 + pass
67 +
68 +
69 +class DraftModelSpeculator(BaseSpeculator):
70 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
71 + self.vllm_config = vllm_config
72 + self.device = device
73 +
74 + assert vllm_config.speculative_config is not None
75 + self.speculative_config = vllm_config.speculative_config
76 + self.method = self.speculative_config.method
77 + self.num_speculative_steps = self.speculative_config.num_speculative_tokens
78 + self.draft_model_config = self.speculative_config.draft_model_config
79 +
80 + self.scheduler_config = vllm_config.scheduler_config
81 + self.max_num_reqs = self.scheduler_config.max_num_seqs
82 + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
83 + self.max_model_len = vllm_config.model_config.max_model_len
84 + self.draft_max_seq_len = self.max_model_len
85 + # We need to get the hidden size from the draft model config because
86 + # the draft model's hidden size can be different from the target model's
87 + # hidden size (e.g., Llama 3.3 70B).
88 + self.hidden_size = self.draft_model_config.get_hidden_size()
89 + # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP
90 + # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual).
91 + # Non-HC models default to hc_mult=1 and are unaffected.
92 + hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
93 + self.hidden_size = self.hidden_size * hc_mult
94 + self.vocab_size = self.draft_model_config.get_vocab_size()
95 + self.dtype = vllm_config.model_config.dtype
96 + self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
97 +
98 + # DP configuration
99 + self.dp_size = vllm_config.parallel_config.data_parallel_size
100 + self.dp_rank = vllm_config.parallel_config.data_parallel_rank
101 +
102 + self.input_buffers = InputBuffers(
103 + max_num_reqs=self.max_num_reqs,
104 + max_num_tokens=self.max_num_tokens,
105 + device=device,
106 + )
107 + self.idx_mapping = torch.zeros(
108 + self.max_num_reqs, dtype=torch.int32, device=device
109 + )
110 + self.temperature = torch.zeros(
111 + self.max_num_reqs, dtype=torch.float32, device=device
112 + )
113 + self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device)
114 + self.draft_tokens = torch.zeros(
115 + self.max_num_reqs,
116 + self.num_speculative_steps,
117 + dtype=torch.int64,
118 + device=device,
119 + )
120 + self.arange = torch.arange(
121 + self.max_num_reqs + 1, dtype=torch.int32, device="cpu"
122 + )
123 +
124 + self.draft_logits: torch.Tensor | None = None
125 + if self.speculative_config.draft_sample_method == "probabilistic":
126 + self.draft_logits = torch.zeros(
127 + self.max_num_reqs,
128 + self.num_speculative_steps,
129 + self.vocab_size,
130 + dtype=torch.float32,
131 + device=device,
132 + )
133 +
134 + @abstractmethod
135 + def load_draft_model(
136 + self,
137 + target_model: nn.Module,
138 + target_attn_layer_names: set[str],
139 + ) -> nn.Module:
140 + pass
141 +
142 + def load_model(self, target_model: nn.Module) -> None:
143 + target_attn_layer_names = set(
144 + get_layers_from_vllm_config(
145 + self.vllm_config,
146 + AttentionLayerBase, # type: ignore[type-abstract]
147 + ).keys()
148 + )
149 +
150 + self.model = self.load_draft_model(target_model, target_attn_layer_names)
151 +
152 + all_attn_layers = set[str](
153 + get_layers_from_vllm_config(
154 + self.vllm_config,
155 + AttentionLayerBase, # type: ignore[type-abstract]
156 + ).keys()
157 + )
158 + self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
159 +
160 + def set_attn(
161 + self,
162 + model_state: ModelState,
163 + kv_cache_config: KVCacheConfig,
164 + block_tables: BlockTables,
165 + ) -> None:
166 + self.model_state = model_state
167 + self.kv_cache_config = kv_cache_config
168 + self.attn_groups, _, _ = init_attn_backend(
169 + kv_cache_config,
170 + self.vllm_config,
171 + self.device,
172 + active_layer_names=self.draft_attn_layer_names,
173 + )
174 + self.block_tables = block_tables
175 +
176 + def _build_draft_attn_metadata(
177 + self,
178 + num_reqs: int,
179 + num_reqs_padded: int,
180 + num_tokens_padded: int,
181 + ) -> dict[str, Any] | None:
182 + query_start_loc_cpu = torch.clamp(
183 + self.arange[: num_reqs_padded + 1], max=num_reqs
184 + )
185 + block_tables = [
186 + x[:num_reqs_padded] for x in self.block_tables.input_block_tables
187 + ]
188 + slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
189 + attn_metadata = build_attn_metadata(
190 + attn_groups=self.attn_groups,
191 + num_reqs=num_reqs_padded,
192 + num_tokens=num_tokens_padded,
193 + query_start_loc_gpu=self.input_buffers.query_start_loc[
194 + : num_reqs_padded + 1
195 + ],
196 + query_start_loc_cpu=query_start_loc_cpu,
197 + max_query_len=1,
198 + seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
199 + max_seq_len=self.draft_max_seq_len,
200 + block_tables=block_tables,
201 + slot_mappings=slot_mappings,
202 + kv_cache_config=self.kv_cache_config,
203 + )
204 + return attn_metadata
205 +
206 + def _copy_request_inputs(
207 + self,
208 + num_reqs: int,
209 + # [num_reqs]
210 + idx_mapping: torch.Tensor,
211 + # [max_num_reqs]
212 + temperature: torch.Tensor,
213 + # [max_num_reqs]
214 + seeds: torch.Tensor,
215 + ) -> None:
216 + # Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
217 + # NOTE(woosuk): For draft sampling, we only consider the temperature
218 + # and ignore the other sampling parameters such as top_k and top_p,
219 + # for simplicity and performance.
220 + # While this may slightly degrade the acceptance rate, it does not
221 + # affect the output distribution after rejection sampling.
222 + self.temperature.copy_(temperature)
223 + self.seeds.copy_(seeds)
224 + self.idx_mapping[:num_reqs].copy_(idx_mapping)