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

[Model Runner V2][Spec Decode] Support spec decode with draft model (#43091)

Signed-off-by: Icey <1790571317@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Xiaoshuang Wang committed Sep 29, 2026 at 03:38 UTC ef42093a36b7d7cc74ecfc0533099e01f2aa8286
6 files changed +552 -4
vllm/config/vllm.py
-2
@@ -3067,8 +3067,6 @@ class VllmConfig:
3067 # https://github.com/vllm-project/vllm/pull/40704
3068 "ngram",
3069 "ngram_gpu",
3070 - # https://github.com/vllm-project/vllm/pull/43091
3071 - "draft_model",
3070 "suffix",
3071 "medusa",
3072 "mlp_speculator",
vllm/v1/worker/gpu/model_runner.py
+12 -1
@@ -799,7 +799,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
799 "skip_attn must only be True for initial memory profiling."
800 )
801
802 - # Create a dummy scheduler output.
802 + # Create a dummy scheduler output. Plain draft-model speculation adds
803 + # one correction slot per request during prefill. The scheduler
804 + # accounts for these slots, while dummy runs bypass the scheduler.
805 + # Adjust the dummy token count only when the expanded draft batch
806 + # would exceed the token budget.
807 num_reqs = min(num_tokens, self.max_num_reqs)
808 if uniform_decode:
809 # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
@@ -813,8 +817,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
817 num_reqs = min(
818 num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
819 )
820 + num_tokens = min(
821 + num_tokens,
822 + self.max_num_tokens
823 + - num_reqs * self.speculator.num_extra_query_per_req,
824 + )
825 +
826 # Distribute the remainder evenly so no dummy request exceeds
827 # ceil(num_tokens / num_reqs) <= max_model_len tokens.
828 +
829 num_tokens_per_request = [
830 num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
831 for i in range(num_reqs)
vllm/v1/worker/gpu/spec_decode/__init__.py
+6
@@ -60,5 +60,11 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
60 )
61
62 return EagleSpeculator(vllm_config, device)
63 + elif speculative_config.uses_draft_model():
64 + from vllm.v1.worker.gpu.spec_decode.draft_model.speculator import (
65 + PlainDraftModelSpeculator,
66 + )
67 +
68 + return PlainDraftModelSpeculator(vllm_config, device)
69 else:
70 raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
vllm/v1/worker/gpu/spec_decode/draft_model/__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/draft_model/speculator.py new
+527
@@ -0,0 +1,527 @@
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.compilation.backends import set_model_tag
9 +from vllm.config import VllmConfig, replace
10 +from vllm.config.compilation import CUDAGraphMode
11 +from vllm.forward_context import BatchDescriptor, set_forward_context
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.model_loader import get_model
14 +from vllm.triton_utils import tl, triton
15 +from vllm.v1.kv_cache_interface import KVCacheConfig
16 +from vllm.v1.spec_decode.utils import next_power_of_2
17 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
18 +from vllm.v1.worker.gpu.block_table import BlockTables
19 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
20 +from vllm.v1.worker.gpu.dp_utils import DPSyncState
21 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
22 +from vllm.v1.worker.gpu.model_states.interface import ModelState
23 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
24 +from vllm.v1.worker.utils import AttentionGroup
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +class PlainDraftModelSpeculator(DraftModelSpeculator):
30 + """Speculative decoding using a separate smaller draft LM.
31 +
32 + Unlike Eagle, the draft model runs fully independently of the target model.
33 + Step 0 builds an expanded buffer (accepted + correction token + rejected
34 + slots) via a Triton kernel; steps 1..k-1 are single-token decode steps.
35 + """
36 +
37 + # Plain draft model needs one extra slot per request for correction.
38 + num_extra_query_per_req = 1
39 +
40 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
41 + super().__init__(vllm_config, device)
42 +
43 + # draft_max_seq_len is read by the parent's attention metadata builder.
44 + # Plain draft model doesn't do per-batch adjustment; cap at max.
45 + self.draft_max_seq_len = self.max_model_len
46 +
47 + self.last_token_indices = torch.zeros(
48 + self.max_num_reqs, dtype=torch.int64, device=device
49 + )
50 + # GPU-side arange for decode query_start_loc initialisation.
51 + self.arange_gpu = torch.arange(
52 + self.max_num_reqs + 1, dtype=torch.int32, device=device
53 + )
54 + # Scalar step counter used as column index into draft_logits.
55 + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
56 +
57 + _expanded_max = self.max_num_tokens + self.max_num_reqs
58 + self.expanded_input_ids = torch.zeros(
59 + _expanded_max, dtype=torch.int32, device=device
60 + )
61 + self.expanded_positions = torch.zeros(
62 + _expanded_max, dtype=torch.int64, device=device
63 + )
64 + self.expanded_slot_mappings: torch.Tensor
65 + self.supports_mm_inputs = False
66 +
67 + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
68 + pass # CUDA graph not yet supported for plain draft model speculator.
69 +
70 + def capture(self) -> None:
71 + pass # CUDA graph not yet supported for plain draft model speculator.
72 +
73 + def set_attn(
74 + self,
75 + model_state: ModelState,
76 + kv_cache_config: KVCacheConfig,
77 + block_tables: BlockTables,
78 + target_input_buffers: InputBuffers,
79 + target_attn_groups: list[list[AttentionGroup]],
80 + ) -> None:
81 + super().set_attn(
82 + model_state,
83 + kv_cache_config,
84 + block_tables,
85 + target_input_buffers,
86 + target_attn_groups,
87 + )
88 + self.expanded_slot_mappings = torch.empty(
89 + block_tables.num_kv_cache_groups,
90 + self.max_num_tokens + self.max_num_reqs,
91 + dtype=torch.int64,
92 + device=self.device,
93 + )
94 +
95 + def load_draft_model(
96 + self,
97 + target_model: nn.Module,
98 + target_attn_layer_names: set[str],
99 + ) -> nn.Module:
100 + spec = self.speculative_config
101 + assert spec is not None
102 + draft_vllm_config = replace(
103 + self.vllm_config,
104 + model_config=self.draft_model_config,
105 + quant_config=None,
106 + parallel_config=replace(
107 + spec.draft_parallel_config,
108 + rank=self.vllm_config.parallel_config.rank,
109 + ),
110 + )
111 + with set_model_tag("draft_model"):
112 + return get_model(
113 + vllm_config=draft_vllm_config,
114 + prefix="draft_model",
115 + )
116 +
117 + @torch.inference_mode()
118 + def _run_model(
119 + self,
120 + input_ids: torch.Tensor,
121 + positions: torch.Tensor,
122 + attn_metadata: dict[str, Any] | None,
123 + slot_mappings: dict[str, torch.Tensor] | None,
124 + num_tokens_across_dp: torch.Tensor | None,
125 + ) -> torch.Tensor:
126 + num_tokens = input_ids.shape[0]
127 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
128 + with set_forward_context(
129 + attn_metadata,
130 + self.vllm_config,
131 + num_tokens=num_tokens,
132 + cudagraph_runtime_mode=CUDAGraphMode.NONE,
133 + num_tokens_across_dp=num_tokens_across_dp,
134 + slot_mapping=slot_mappings,
135 + batch_descriptor=batch_descriptor,
136 + ):
137 + hidden_states = self.model( # type: ignore[misc]
138 + input_ids=input_ids,
139 + positions=positions,
140 + )
141 + if isinstance(hidden_states, tuple):
142 + hidden_states = hidden_states[0]
143 + return hidden_states
144 +
145 + def _accepted_last_indices(
146 + self,
147 + input_batch: InputBatch,
148 + num_rejected: torch.Tensor,
149 + num_reqs: int,
150 + ) -> torch.Tensor:
151 + qsl = input_batch.query_start_loc
152 + adjusted_lens = qsl[1 : num_reqs + 1] - qsl[:num_reqs] - num_rejected[:num_reqs]
153 + return qsl[:num_reqs] + adjusted_lens - 1
154 +
155 + def _prefill(
156 + self,
157 + input_batch: InputBatch,
158 + num_rejected: torch.Tensor,
159 + last_sampled: torch.Tensor,
160 + num_reqs: int,
161 + skip_attn: bool,
162 + num_tokens_across_dp: torch.Tensor | None,
163 + ) -> None:
164 + if skip_attn:
165 + src_tokens = input_batch.num_tokens
166 + self.last_token_indices[:num_reqs] = self._accepted_last_indices(
167 + input_batch, num_rejected, num_reqs
168 + )
169 + hidden_states = self._run_model(
170 + input_batch.input_ids[:src_tokens],
171 + input_batch.positions[:src_tokens],
172 + None,
173 + None,
174 + num_tokens_across_dp,
175 + )
176 + positions_buf = input_batch.positions
177 + else:
178 + block_tables = self.block_tables
179 + kv_cache_config = self.kv_cache_config
180 + assert block_tables is not None
181 + assert kv_cache_config is not None
182 +
183 + total_expanded = prepare_prefill_inputs(
184 + input_buffers=self.input_buffers,
185 + input_batch=input_batch,
186 + last_sampled=last_sampled,
187 + num_rejected=num_rejected,
188 + expanded_input_ids=self.expanded_input_ids,
189 + expanded_positions=self.expanded_positions,
190 + last_token_indices=self.last_token_indices,
191 + max_num_reqs=self.max_num_reqs,
192 + max_model_len=self.max_model_len,
193 + )
194 +
195 + prefill_slot_mappings = block_tables.compute_slot_mappings(
196 + input_batch.idx_mapping,
197 + self.input_buffers.query_start_loc[: num_reqs + 1],
198 + self.expanded_positions[:total_expanded],
199 + total_expanded,
200 + out=self.expanded_slot_mappings,
201 + )
202 + qsl_np = input_batch.query_start_loc_np
203 + query_start_loc_cpu_expanded = (
204 + torch.from_numpy(qsl_np[: num_reqs + 1]).int()
205 + + torch.from_numpy(self.arange_np[: num_reqs + 1]).int()
206 + )
207 + prefill_attn_md = self._build_attn_metadata(
208 + batch_desc=BatchExecutionDescriptor(
209 + cg_mode=CUDAGraphMode.NONE,
210 + num_tokens=total_expanded,
211 + num_reqs=num_reqs,
212 + ),
213 + num_reqs=num_reqs,
214 + query_start_loc_np=query_start_loc_cpu_expanded.numpy(),
215 + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
216 + step=0,
217 + slot_mappings=prefill_slot_mappings,
218 + )
219 + prefill_slot_maps_by_layer = build_slot_mappings_by_layer(
220 + prefill_slot_mappings, kv_cache_config
221 + )
222 + hidden_states = self._run_model(
223 + self.expanded_input_ids[:total_expanded],
224 + self.expanded_positions[:total_expanded],
225 + prefill_attn_md,
226 + prefill_slot_maps_by_layer,
227 + num_tokens_across_dp,
228 + )
229 + positions_buf = self.expanded_positions
230 +
231 + last_indices = self.last_token_indices[:num_reqs]
232 + self.current_draft_step.fill_(0)
233 + self.draft_tokens[:num_reqs, 0] = self.sample_draft(
234 + hidden_states[last_indices],
235 + positions_buf[last_indices],
236 + self.idx_mapping[:num_reqs],
237 + self.temperature,
238 + self.seeds,
239 + self.current_draft_step,
240 + self.draft_logits,
241 + )
242 +
243 + def _multi_step_decode(
244 + self,
245 + input_batch: InputBatch,
246 + num_rejected: torch.Tensor,
247 + num_reqs: int,
248 + skip_attn: bool,
249 + num_tokens_across_dp: torch.Tensor | None,
250 + ) -> None:
251 + if skip_attn:
252 + last_positions = input_batch.positions[
253 + self._accepted_last_indices(input_batch, num_rejected, num_reqs)
254 + ]
255 + else:
256 + last_positions = self.expanded_positions[self.last_token_indices[:num_reqs]]
257 +
258 + self.input_buffers.positions[:num_reqs].copy_(last_positions)
259 + # The decode loop increments seq_lens BEFORE each forward.
260 + # Initial seq_lens = target_seq_lens - num_rejected + 1.
261 + # Step 1 forward uses target_seq_lens - num_rejected + 2.
262 + # Step 2 forward uses target_seq_lens - num_rejected + 3.
263 + self.input_buffers.seq_lens[:num_reqs].copy_(
264 + torch.clamp(
265 + input_batch.seq_lens[:num_reqs] - num_rejected[:num_reqs].int() + 1,
266 + max=self.max_model_len,
267 + )
268 + )
269 + self.input_buffers.query_start_loc[: num_reqs + 1].copy_(
270 + self.arange_gpu[: num_reqs + 1]
271 + )
272 +
273 + idx_mapping = input_batch.idx_mapping
274 +
275 + for step in range(1, self.num_speculative_steps):
276 + self.input_buffers.input_ids[:num_reqs].copy_(
277 + self.draft_tokens[:num_reqs, step - 1].int()
278 + )
279 + torch.clamp(
280 + self.input_buffers.positions[:num_reqs] + 1,
281 + max=self.max_model_len - 1,
282 + out=self.input_buffers.positions[:num_reqs],
283 + )
284 + # Increment seq_lens BEFORE the forward so the attention reads from
285 + # the KV slot written by the previous step.
286 + torch.clamp(
287 + self.input_buffers.seq_lens[:num_reqs] + 1,
288 + max=self.max_model_len,
289 + out=self.input_buffers.seq_lens[:num_reqs],
290 + )
291 +
292 + decode_attn_md = None
293 + step_slot_maps_by_layer = None
294 + if not skip_attn:
295 + block_tables = self.block_tables
296 + kv_cache_config = self.kv_cache_config
297 + assert block_tables is not None
298 + assert kv_cache_config is not None
299 + q_start = self.input_buffers.query_start_loc[: num_reqs + 1]
300 + positions = self.input_buffers.positions[:num_reqs]
301 + step_slot_mappings = block_tables.compute_slot_mappings(
302 + idx_mapping, q_start, positions, num_reqs
303 + )
304 + step_slot_maps_by_layer = build_slot_mappings_by_layer(
305 + step_slot_mappings, kv_cache_config
306 + )
307 + decode_attn_md = self._build_uniform_attn_metadata(
308 + batch_desc=BatchExecutionDescriptor(
309 + cg_mode=CUDAGraphMode.NONE,
310 + num_tokens=num_reqs,
311 + num_reqs=num_reqs,
312 + uniform_token_count=1,
313 + ),
314 + num_reqs=num_reqs,
315 + num_query_per_req=1,
316 + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
317 + # Include the correction token inserted during prefill.
318 + step=step + 1,
319 + )
320 + hidden_states = self._run_model(
321 + self.input_buffers.input_ids[:num_reqs],
322 + self.input_buffers.positions[:num_reqs],
323 + decode_attn_md,
324 + step_slot_maps_by_layer,
325 + num_tokens_across_dp,
326 + )
327 +
328 + self.current_draft_step.fill_(step)
329 + self.draft_tokens[:num_reqs, step] = self.sample_draft(
330 + hidden_states[:num_reqs],
331 + self.input_buffers.positions[:num_reqs],
332 + self.idx_mapping[:num_reqs],
333 + self.temperature,
334 + self.seeds,
335 + self.current_draft_step,
336 + self.draft_logits,
337 + )
338 +
339 + @torch.inference_mode()
340 + def propose(
341 + self,
342 + input_batch: InputBatch,
343 + attn_metadata: dict[str, Any],
344 + slot_mappings: dict[str, torch.Tensor],
345 + last_hidden_states: torch.Tensor, # unused
346 + aux_hidden_states: list[torch.Tensor] | None, # unused
347 + num_sampled: torch.Tensor,
348 + num_rejected: torch.Tensor,
349 + last_sampled: torch.Tensor,
350 + next_prefill_tokens: torch.Tensor,
351 + temperature: torch.Tensor,
352 + seeds: torch.Tensor,
353 + dp_sync: DPSyncState | None = None,
354 + dummy_run: bool = False,
355 + skip_attn_for_dummy_run: bool = False,
356 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
357 + is_profile: bool = False,
358 + ) -> torch.Tensor:
359 + assert self.model is not None
360 +
361 + num_reqs = input_batch.num_reqs
362 + skip_attn = dummy_run and skip_attn_for_dummy_run
363 + num_tokens_across_dp = (
364 + dp_sync.num_tokens_across_dp if dp_sync is not None else None
365 + )
366 +
367 + # Copy per-request temperature/seeds/idx_mapping into pre-allocated
368 + # buffers so sample_draft can read them without extra slicing.
369 + self._copy_request_inputs(
370 + num_reqs,
371 + input_batch.idx_mapping,
372 + temperature,
373 + seeds,
374 + )
375 +
376 + self._prefill(
377 + input_batch,
378 + num_rejected,
379 + last_sampled,
380 + num_reqs,
381 + skip_attn,
382 + num_tokens_across_dp,
383 + )
384 +
385 + if self.num_speculative_steps == 1:
386 + return self.draft_tokens[:num_reqs, :1]
387 +
388 + self._multi_step_decode(
389 + input_batch,
390 + num_rejected,
391 + num_reqs,
392 + skip_attn,
393 + num_tokens_across_dp,
394 + )
395 +
396 + return self.draft_tokens[:num_reqs]
397 +
398 +
399 +@triton.jit
400 +def _prepare_prefill_inputs_kernel(
401 + target_input_ids_ptr, # [src_tokens] int32
402 + target_positions_ptr, # [src_tokens] int64
403 + last_sampled_ptr, # [max_num_reqs] int32
404 + idx_mapping_ptr, # [num_reqs] int32
405 + out_input_ids_ptr, # [src_tokens + num_reqs] int32
406 + out_positions_ptr, # [src_tokens + num_reqs] int64
407 + last_token_indices_ptr, # [max_num_reqs] int64
408 + out_query_start_loc_ptr, # [max_num_reqs + 1] int32
409 + out_seq_lens_ptr, # [max_num_reqs] int32
410 + query_start_loc_ptr, # [num_reqs + 1] int32
411 + seq_lens_ptr, # [num_reqs] int32
412 + num_rejected_ptr, # [num_reqs] int32
413 + max_num_reqs,
414 + max_model_len,
415 + BLOCK_SIZE: tl.constexpr,
416 +):
417 + """Per-request step-0 input preparation for the plain draft-model speculator.
418 +
419 + Output layout for request i (out_start = query_start_loc[i] + i):
420 + [out_start, out_start + num_valid) accepted tokens
421 + [out_start + num_valid] correction token
422 + (out_start + num_valid, out_start + total_out) rejected slots
423 +
424 + where num_valid = query_lens[i] - num_rejected[i]
425 + total_out = query_lens[i] + 1
426 + """
427 + req_idx = tl.program_id(0)
428 + num_reqs = tl.num_programs(0)
429 +
430 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
431 +
432 + q_start = tl.load(query_start_loc_ptr + req_idx)
433 + q_next = tl.load(query_start_loc_ptr + req_idx + 1)
434 + seq_len = tl.load(seq_lens_ptr + req_idx)
435 + num_rejected = tl.load(num_rejected_ptr + req_idx)
436 +
437 + num_valid = q_next - q_start - num_rejected
438 + correction_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
439 + start_pos = tl.load(target_positions_ptr + q_start)
440 + out_start = q_start + req_idx
441 + total_out = q_next - q_start + 1
442 +
443 + for i in range(0, total_out, BLOCK_SIZE):
444 + j = i + tl.arange(0, BLOCK_SIZE)
445 + in_bounds = j < total_out
446 +
447 + is_valid = j < num_valid
448 + is_correction = j == num_valid
449 + src_idx = q_start + j
450 + token_ids = tl.load(target_input_ids_ptr + src_idx, mask=is_valid, other=0)
451 + positions = tl.minimum(start_pos + j, max_model_len - 1)
452 + token_ids = tl.where(is_correction, correction_token, token_ids)
453 +
454 + out_idx = out_start + j
455 + tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds)
456 + tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds)
457 +
458 + tl.store(last_token_indices_ptr + req_idx, out_start + num_valid)
459 + tl.store(out_query_start_loc_ptr + req_idx, out_start)
460 + # seqlen_k for the expanded prefill attention:
461 + # pre_existing = seq_len - query_len (draft KV before this round)
462 + # expanded_query_len = query_len + 1 (adds correction token)
463 + # seqlen_k = pre_existing + expanded_query_len = seq_len + 1
464 + # num_rejected does NOT change seqlen_k — the expanded query always has
465 + # query_len+1 slots. Rejected positions write temporary KV entries that
466 + # are overwritten with valid tokens in a later round.
467 + new_seq_len = tl.minimum(seq_len + 1, max_model_len)
468 + tl.store(out_seq_lens_ptr + req_idx, new_seq_len)
469 +
470 + if req_idx == num_reqs - 1:
471 + total_expanded = out_start + total_out
472 + tl.store(out_query_start_loc_ptr + num_reqs, total_expanded)
473 + for i in range(num_reqs + 1, max_num_reqs + 1, BLOCK_SIZE):
474 + block = i + tl.arange(0, BLOCK_SIZE)
475 + mask = block <= max_num_reqs
476 + tl.store(out_query_start_loc_ptr + block, total_expanded, mask=mask)
477 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
478 + block = i + tl.arange(0, BLOCK_SIZE)
479 + mask = block < max_num_reqs
480 + tl.store(out_seq_lens_ptr + block, 0, mask=mask)
481 + tl.store(last_token_indices_ptr + block, 0, mask=mask)
482 +
483 +
484 +def prepare_prefill_inputs(
485 + input_buffers: InputBuffers,
486 + input_batch: InputBatch,
487 + last_sampled: torch.Tensor, # [max_num_reqs] int64
488 + num_rejected: torch.Tensor, # [num_reqs] int64
489 + expanded_input_ids: torch.Tensor, # [max_tokens + max_reqs] int32
490 + expanded_positions: torch.Tensor, # [max_tokens + max_reqs] int64
491 + last_token_indices: torch.Tensor, # [max_num_reqs] int64
492 + max_num_reqs: int,
493 + max_model_len: int,
494 +) -> int:
495 + """Call _prepare_prefill_inputs_kernel and return total_expanded tokens.
496 +
497 + Side-effects (kernel writes):
498 + - expanded_input_ids, expanded_positions
499 + - last_token_indices
500 + - input_buffers.query_start_loc (expanded: original[i] + i)
501 + - input_buffers.seq_lens (= target_seq_lens - num_rejected + 1)
502 + """
503 + num_reqs = input_batch.num_reqs
504 + src_tokens = input_batch.num_tokens
505 + qsl_np = input_batch.query_start_loc_np
506 + query_lens = qsl_np[1 : num_reqs + 1] - qsl_np[:num_reqs]
507 + max_total_out = int(query_lens.max()) + 2 # +1 correction, +1 alignment
508 + BLOCK_SIZE = min(512, next_power_of_2(max_total_out))
509 +
510 + _prepare_prefill_inputs_kernel[(num_reqs,)](
511 + input_batch.input_ids,
512 + input_batch.positions,
513 + last_sampled.int(),
514 + input_batch.idx_mapping,
515 + expanded_input_ids,
516 + expanded_positions,
517 + last_token_indices,
518 + input_buffers.query_start_loc,
519 + input_buffers.seq_lens,
520 + input_batch.query_start_loc,
521 + input_batch.seq_lens,
522 + num_rejected.int(),
523 + max_num_reqs,
524 + max_model_len,
525 + BLOCK_SIZE=BLOCK_SIZE,
526 + )
527 + return src_tokens + num_reqs
vllm/v1/worker/gpu/spec_decode/speculator.py
+5 -1
@@ -58,6 +58,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
58
59 class BaseSpeculator(ABC):
60 num_query_per_req: int = 1
61 + # Extra query slots reserved per request outside the regular queries.
62 + num_extra_query_per_req: int = 0
63
64 @abstractmethod
65 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
@@ -302,6 +304,7 @@ class DraftModelSpeculator(BaseSpeculator):
304 step: int,
305 causal: bool | Mapping[int, bool] = True,
306 dcp_local_seq_lens: torch.Tensor | None = None,
307 + slot_mappings: torch.Tensor | None = None,
308 ) -> dict[str, Any] | None:
309 num_reqs_padded = batch_desc.num_reqs or num_reqs
310 # A FULL graph replays a captured shape whose padded requests each hold
@@ -323,7 +326,8 @@ class DraftModelSpeculator(BaseSpeculator):
326 block_tables = [
327 x[:num_reqs_padded] for x in self.block_tables.input_block_tables
328 ]
326 - slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
329 + if slot_mappings is None:
330 + slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
331 draft_seq_lens_cpu_upper_bound = torch.zeros(
332 num_reqs_padded, dtype=torch.int32, device="cpu"
333 )