From 58aa56999212b43ac18775ffcb7f774198607b4c Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Mon, 20 Apr 2026 14:57:07 +0800 Subject: [PATCH 01/34] async ngram gpu debug Signed-off-by: PatchouliTaisa --- vllm/v1/core/sched/scheduler.py | 42 ++- vllm/v1/outputs.py | 2 + vllm/v1/worker/gpu/model_runner.py | 27 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 6 + .../worker/gpu/spec_decode/ngram/__init__.py | 0 .../gpu/spec_decode/ngram/speculator.py | 332 ++++++++++++++++++ vllm/v1/worker/gpu/spec_decode/utils.py | 79 ++++- vllm/v1/worker/gpu/states.py | 7 +- 8 files changed, 474 insertions(+), 21 deletions(-) create mode 100644 vllm/v1/worker/gpu/spec_decode/ngram/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/ngram/speculator.py diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fe524ccace16..66febc371d37 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1667,9 +1667,16 @@ def _free_encoder_inputs(self, request: Request) -> None: self.encoder_cache_manager.free_encoder_input(request, input_id) def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: - for req_id, spec_token_ids in zip( - draft_token_ids.req_ids, - draft_token_ids.draft_token_ids, + # Optional per-request valid-draft-count from variable-length + # speculators (e.g. ngram_gpu). When present, we enumerate with an + # index so each request can be trimmed independently before + # further processing (structured-output grammar validation, etc.). + num_valid_list = draft_token_ids.num_valid_draft_tokens + for i, (req_id, spec_token_ids) in enumerate( + zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ) ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -1682,6 +1689,17 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: request.spec_token_ids = [] continue + # Variable-length drafters: truncate to the number of drafts + # that actually matched a real n-gram. Anything beyond that + # count was a safe-fallback substitution (see + # NgramGPUSpeculator) and must not be surfaced as a draft. + if num_valid_list is not None: + num_valid = num_valid_list[i] + if num_valid < len(spec_token_ids): + # Slice defensively; spec_token_ids is a Python list + # produced by ``ndarray.tolist()`` so slicing is O(k). + spec_token_ids = spec_token_ids[:num_valid] + # Add newly generated spec token ids to the request. if self.structured_output_manager.should_advance(request): metadata = request.structured_output_request @@ -1692,11 +1710,14 @@ def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput ) -> None: num_invalid_spec_tokens: dict[str, int] = {} + num_valid_list = draft_token_ids.num_valid_draft_tokens sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens - for req_id, spec_token_ids in zip( - draft_token_ids.req_ids, - draft_token_ids.draft_token_ids, + for i, (req_id, spec_token_ids) in enumerate( + zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ) ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -1710,7 +1731,14 @@ def update_draft_token_ids_in_output( orig_num_spec_tokens = len(placeholder_spec_tokens) # Trim drafts to scheduled number of spec tokens # (needed for chunked prefill case for example). - del spec_token_ids[orig_num_spec_tokens:] + effective_num_spec_tokens = orig_num_spec_tokens + if num_valid_list is not None: + effective_num_spec_tokens = max( + 0, + min(num_valid_list[i], orig_num_spec_tokens), + ) + + del spec_token_ids[effective_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. if self.structured_output_manager.should_advance(request): metadata = request.structured_output_request diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 1f102ec61783..7bb0f588f907 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -221,6 +221,8 @@ class DraftTokenIds: req_ids: list[str] # num_reqs x num_draft_tokens draft_token_ids: list[list[int]] + # [num_reqs] + num_valid_draft_tokens: list[int] | None = None def make_empty_encoder_model_runner_output( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a2f83c52e951..86da8ff32b87 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -183,6 +183,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.pooling_runner: PoolingRunner | None = None # General request states. + use_dense_all_token_ids = ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ) self.req_states = RequestState( max_num_reqs=self.max_num_reqs, max_model_len=self.max_model_len, @@ -190,6 +194,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): num_speculative_steps=self.num_speculative_steps, vocab_size=self.vocab_size, device=self.device, + use_dense_all_token_ids=use_dense_all_token_ids, ) self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, @@ -197,6 +202,13 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=self.device, ) + # Inject RequestState into speculators that consume the persistent + # token store directly (e.g. NgramGPUSpeculator). Attribute-based + # injection keeps speculators that don't need it (EagleSpeculator) + # untouched — they simply don't declare the `req_states` attribute. + if self.speculator is not None and hasattr(self.speculator, "req_states"): + self.speculator.req_states = self.req_states + self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None self.prompt_logprobs_worker: PromptLogprobsWorker | None = None @@ -296,7 +308,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: if not load_dummy_weights: prepare_communication_buffer_for_model(self.model) - if self.speculator is not None: + if self.speculator is not None and hasattr(self.speculator, "model"): prepare_communication_buffer_for_model(self.speculator.model) # Initialize the components that require the model. @@ -1219,7 +1231,18 @@ def sample_tokens( mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) + # Speculators that can produce fewer than `num_speculative_steps` + # real drafts (e.g. ngram_gpu on a no-match row) expose a + # `get_num_valid_draft_tokens` accessor. When present, the + # companion tensor is shipped to the scheduler via the same + # async D2H copy stream as the draft tokens themselves. + num_valid_draft_tokens: torch.Tensor | None = None + get_num_valid = getattr(self.speculator, "get_num_valid_draft_tokens", None) + if get_num_valid is not None: + num_valid_draft_tokens = get_num_valid(input_batch.num_reqs) + self.draft_tokens_handler.set_draft_tokens( + input_batch, draft_tokens, num_valid_draft_tokens + ) if self.use_async_scheduling: return async_output diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 536b7526bddd..eca2eb718c8a 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -12,4 +12,10 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator return EagleSpeculator(vllm_config, device) + if speculative_config.use_ngram_gpu(): + from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( + NgramGPUSpeculator, + ) + + return NgramGPUSpeculator(vllm_config, device) raise NotImplementedError(f"{speculative_config.method} is not supported yet.") diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py b/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py new file mode 100644 index 000000000000..468794a69c8d --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""V2 GPU-native N-gram speculator. + +Unlike V1's ``NgramProposerGPU`` (in ``vllm/v1/spec_decode/ngram_proposer_gpu.py``), +this speculator is designed to slot directly into the V2 runner's speculator +interface. It reuses the runner-owned ``RequestState.all_token_ids`` / +``RequestState.total_len`` tensors as the persistent token store and thus +does not maintain any shadow GPU state. + +Key properties: + * Fully vectorized ``unfold → broadcast-compare → argmax → gather`` algorithm + — no CPU-GPU synchronization inside ``propose``. + * ``-1`` no-match positions are rewritten on GPU to ``last_sampled_tokens`` + so that the downstream ``combine_sampled_and_draft_tokens`` kernel never + writes an invalid token id into ``input_ids``. + * A companion ``num_valid_draft_tokens`` tensor is emitted every step and + propagated to the scheduler so that ``request.spec_token_ids`` can be + correctly truncated before the next scheduling round. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.v1.worker.gpu.input_batch import InputBatch + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.states import RequestState + + +@support_torch_compile() +class _NgramKernel(nn.Module): + """Pure, stateless n-gram match kernel. + + Exposed as an ``nn.Module`` so that ``@support_torch_compile`` can hoist + the whole body into a single Inductor-compiled region. The module has + no parameters; it is compiled purely for operator fusion. + """ + + def __init__(self, min_n: int, max_n: int, k: int): + super().__init__() + assert 1 <= min_n <= max_n, ( + f"min_n must be in [1, max_n]; got min_n={min_n}, max_n={max_n}" + ) + assert k >= 1 + self.min_n = min_n + self.max_n = max_n + self.k = k + self.num_sizes = max_n - min_n + 1 + + def forward( + self, + token_ids: torch.Tensor, # [B, L] int32 + seq_lens: torch.Tensor, # [B] int32 (current total_len per req) + valid_mask: torch.Tensor, # [B] bool (row eligible for n-gram lookup) + last_sampled: torch.Tensor, # [B] int64 (fallback for -1 positions) + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(draft_tokens[B, k] int64, num_valid_draft_tokens[B] int32)``. + + The ``draft_tokens`` tensor is guaranteed to contain only valid + vocabulary IDs. Positions that correspond to no-match / out-of-range + slots are filled with ``last_sampled`` so that writing them to + ``input_ids`` later is safe. The companion ``num_valid_draft_tokens`` + encodes how many *leading contiguous* draft positions were actually + derived from a real n-gram match. + """ + B, L = token_ids.shape + device = token_ids.device + + # --- (1) For each n-gram size in [min_n, max_n], find the earliest + # prior occurrence of the trailing suffix. The outer `for` loop runs + # at most `max_n - min_n + 1` times (typically <= 4), and is required + # because `Tensor.unfold` needs a Python int window size. Each + # iteration executes as one fused GPU pass. + first_match_pos = torch.full( + (B, self.num_sizes), -1, dtype=torch.long, device=device + ) + batch_idx = torch.arange(B, device=device) + + for i, n in enumerate(range(self.min_n, self.max_n + 1)): + # [B, L-n+1, n] view — O(1) cost (contiguous stride trick). + windows = token_ids.unfold(1, n, 1) + num_windows = windows.shape[1] + + # Gather the trailing n tokens from each sequence as the query. + # `suffix_start` is clamped to 0 to keep the gather in-bounds; the + # corresponding row is also filtered by `valid_mask` upstream. + suffix_start = (seq_lens.long() - n).clamp(min=0) + offsets = torch.arange(n, device=device) + suffix_idx = suffix_start.unsqueeze(1) + offsets # [B, n] + suffix = torch.gather(token_ids, 1, suffix_idx) # [B, n] + + # Element-wise equality → all(-1) → per-window boolean match. + matches = (windows == suffix.unsqueeze(1)).all(dim=-1) # [B, L-n+1] + + # A match is only actionable if at least one draft token follows. + max_valid_pos = seq_lens.long() - n - 1 # may be negative + window_pos = torch.arange(num_windows, device=device) + matches = matches & (window_pos.unsqueeze(0) <= max_valid_pos.unsqueeze(1)) + + # `argmax(int)` returns 0 on all-false rows; verify with a lookup. + idx = matches.int().argmax(dim=1) # [B] + has_match = matches[batch_idx, idx] # [B] + first_match_pos[:, i] = torch.where(has_match, idx.long(), -1) + + # --- (2) Pick the LONGEST n-gram with a match (flip+argmax, no sync). + # Equivalent to "last True from the right". + best_i = (first_match_pos >= 0).int().flip(dims=[1]).argmax(dim=1) + best_i = self.num_sizes - 1 - best_i # flip index back + best_pos = first_match_pos[batch_idx, best_i] # [B] + ngram_lens_table = torch.arange( + self.min_n, self.max_n + 1, device=device, dtype=torch.long + ) + best_n = ngram_lens_table[best_i] # [B] + has_any = best_pos >= 0 # [B] bool + + # --- (3) Gather k tokens starting right after the matched suffix. + draft_start = torch.where( + has_any, + best_pos + best_n, + torch.zeros_like(best_pos), + ) # [B] + k_range = torch.arange(self.k, device=device) # [k] + draft_idx = (draft_start.unsqueeze(1) + k_range).clamp_(0, L - 1) + drafts = torch.gather(token_ids, 1, draft_idx).to(torch.int64) # [B, k] + + # --- (4) Compute the leading-valid mask per row. + tokens_available = (seq_lens.long() - draft_start).clamp_(min=0) # [B] + valid_positions = k_range.unsqueeze(0) < tokens_available.unsqueeze(1) + row_valid = has_any & valid_mask # [B] + leading_valid_mask = valid_positions & row_valid.unsqueeze(1) # [B, k] + + # --- (5) num_valid = length of the leading contiguous valid run. + # cumsum counts valid positions so far; match against [1, 2, ..., k] + # tells us whether each position is still in the leading run. + cum_valid = leading_valid_mask.int().cumsum(dim=1) # [B, k] + positions = torch.arange(1, self.k + 1, device=device) # [k] + num_valid = (cum_valid == positions.unsqueeze(0)).int().sum(dim=1) # [B] + + # --- (6) Safe-substitute invalid slots with last_sampled so that the + # next step's input_ids are always in-vocabulary. `last_sampled` is + # int64 already (matches RequestState.last_sampled_tokens dtype). + safe_drafts = torch.where( + leading_valid_mask, + drafts, + last_sampled.view(-1, 1).expand(B, self.k), + ) + return safe_drafts, num_valid.to(torch.int32) + + +class NgramGPUSpeculator: + """V2-compatible GPU n-gram speculator. + + Public surface mirrors ``EagleSpeculator`` (``load_model``, ``set_attn``, + ``init_cudagraph_manager``, ``capture_model``, ``propose``). Unlike the + Eagle path, this speculator: + + * has no neural model, no attention layers, no CUDA graph capture; + * reuses the runner-owned ``RequestState`` persistent token store; + * emits ``num_valid_draft_tokens`` as a secondary output so that the + scheduler can trim per-request drafts that were effectively a miss. + """ + + # Consumed by the runner's spec-decode sample path — see + # ``vllm/v1/worker/gpu/model_runner.py::sample``. + supports_mm_inputs = False + draft_logits = None # probabilistic rejection sampling is not supported + + def __init__(self, vllm_config: VllmConfig, device: torch.device): + spec = vllm_config.speculative_config + assert spec is not None + assert spec.prompt_lookup_min is not None, ( + "prompt_lookup_min must be configured for ngram_gpu" + ) + assert spec.prompt_lookup_max is not None, ( + "prompt_lookup_max must be configured for ngram_gpu" + ) + + self.vllm_config = vllm_config + self.device = device + self.speculative_config = spec + self.num_speculative_steps: int = spec.num_speculative_tokens + + self.min_n: int = spec.prompt_lookup_min + self.max_n: int = spec.prompt_lookup_max + + self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs + self.max_model_len: int = vllm_config.model_config.max_model_len + + self.kernel = ( + _NgramKernel( + min_n=self.min_n, + max_n=self.max_n, + k=self.num_speculative_steps, + ) + .to(device) + .eval() + ) + + # Persistent scratch buffer — shaped [max_num_reqs] on device — used to + # stash the last ``num_valid_draft_tokens`` tensor so that the runner + # can forward it to ``DraftTokensHandler`` after ``propose`` returns. + self.num_valid_draft_tokens: torch.Tensor = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + + # Injected by the model runner once ``req_states`` is constructed. + # Accessed inside ``propose``. + self.req_states: RequestState | None = None + + # ------------------------------------------------------------------ + # V2 Speculator interface: intentional no-ops. + # ------------------------------------------------------------------ + def load_model(self, target_model: nn.Module) -> None: + """No weights to load — ngram is a data-only proposer.""" + pass + + def set_attn(self, *args: Any, **kwargs: Any) -> None: + """No attention layers owned by this speculator.""" + pass + + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + """N-gram kernel is torch.compile-managed; no explicit CG capture.""" + pass + + def capture_model(self) -> None: + """No graph capture phase required.""" + pass + + # ------------------------------------------------------------------ + # V2 Speculator interface: main entry point. + # ------------------------------------------------------------------ + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: Any, # unused + slot_mappings: Any, # unused + # [num_tokens, hidden_size] — unused by ngram, kept for signature parity + last_hidden_states: torch.Tensor, + # num_layers x [num_tokens, hidden_size] — unused + aux_hidden_states: list[torch.Tensor] | None, + # [num_reqs] int32 + num_sampled: torch.Tensor, + # [num_reqs] int32 — unused by ngram + num_rejected: torch.Tensor, + # [max_num_reqs, 1] int64 + last_sampled_tokens: torch.Tensor, + # [max_num_reqs] int32 — unused + next_prefill_tokens: torch.Tensor, + # [max_num_reqs] — unused + temperature: torch.Tensor, + # [max_num_reqs] — unused + seeds: torch.Tensor, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + """Propose up to ``num_speculative_steps`` draft tokens per request. + + Returns a ``[num_reqs, num_speculative_steps]`` int64 tensor. Invalid + or no-match positions are backfilled with ``last_sampled`` rather + than ``-1`` so that downstream ``combine_sampled_and_draft_tokens`` + never writes out-of-vocab ids into ``input_ids``. Truth about how + many of those drafts are "real" lives in + ``self.num_valid_draft_tokens[:num_reqs]``. + """ + assert self.req_states is not None, ( + "NgramGPUSpeculator.req_states was not injected by the model " + "runner. Ensure model_runner sets `speculator.req_states = " + "self.req_states` after RequestState is constructed." + ) + + num_reqs = input_batch.num_reqs + # `idx_mapping` is [num_reqs] int32 on device; advanced indexing + # requires int64 → one-time cast (negligible cost vs. the + # O(B * L * num_sizes) compare). + idx_mapping_long = input_batch.idx_mapping.long() + + # Persistent [max_num_reqs, max_model_len] store, UVA-backed. + # Advanced indexing materialises a contiguous [num_reqs, L] view on + # device — identical access pattern to V1's shadow tensor. No extra + # H2D copies are issued. + active_tokens: torch.Tensor = self.req_states.all_token_ids.gpu[ + idx_mapping_long + ] + active_seq_lens: torch.Tensor = self.req_states.total_len.gpu[idx_mapping_long] + active_last_sampled: torch.Tensor = last_sampled_tokens.view(-1)[ + idx_mapping_long + ] + + # A request can draft iff (a) at least one real token was just + # sampled for it (otherwise we cannot trust its suffix) AND + # (b) the sequence already contains min_n tokens for the lookup. + valid_mask = (num_sampled > 0) & (active_seq_lens >= self.min_n) + + with set_forward_context(None, self.vllm_config): + drafts, num_valid = self.kernel( + active_tokens, + active_seq_lens, + valid_mask, + active_last_sampled, + ) + + # Stash num_valid so the runner can forward it to DraftTokensHandler + # after scatter. Note: this is a device tensor; D2H happens on a + # side stream in the handler. + self.num_valid_draft_tokens[:num_reqs].copy_(num_valid) + # Zero out the tail slots to avoid leaking stale values from prior + # steps in case the runner ever peeks beyond num_reqs. + if num_reqs < self.max_num_reqs: + self.num_valid_draft_tokens[num_reqs:].zero_() + + return drafts # [num_reqs, num_speculative_steps] int64, no -1 + + def get_num_valid_draft_tokens(self, num_reqs: int) -> torch.Tensor: + """Return the last step's per-request valid draft counts. + + Sliced view of the internal buffer; safe to pass to an async D2H + copy on a side stream. + """ + return self.num_valid_draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index e1fa21aeb8ae..3a2cbaf8dbca 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -16,32 +16,91 @@ def __init__(self, device: torch.device | None = None): self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None + # Per-request valid-draft-count, populated by speculators that + # return variable-length drafts (e.g. NgramGPUSpeculator). ``None`` + # preserves legacy fixed-length semantics (Eagle / Medusa). + self.num_valid_draft_tokens_np: np.ndarray | None = None self.num_draft_tokens: int = 0 def set_draft_tokens( - self, input_batch: InputBatch, draft_tokens: torch.Tensor + self, + input_batch: InputBatch, + draft_tokens: torch.Tensor, + num_valid_draft_tokens: torch.Tensor | None = None, ) -> None: + """Stage the per-step draft tokens for async D2H copy. + + Args: + input_batch: The just-executed step's input batch (used for + ``req_ids`` and the structured-outputs flag). + draft_tokens: ``[num_reqs, num_speculative_steps]`` int64 on + device. Must be a final tensor (no further GPU mutation + after this call) so the async copy is safe. + num_valid_draft_tokens: Optional ``[num_reqs]`` int32 on device. + Required for variable-length speculators so the scheduler + can trim per-request drafts before the next scheduling + round. When ``None`` (Eagle-style fixed-length drafts), + all draft positions are treated as valid. + """ self.req_ids = input_batch.req_ids self.num_draft_tokens = draft_tokens.shape[1] - if not input_batch.has_structured_output_reqs: - # No draft token validation needs to be performed by - # the scheduler for this batch. + + needs_draft_copy = input_batch.has_structured_output_reqs + needs_valid_copy = num_valid_draft_tokens is not None + + if not needs_draft_copy and not needs_valid_copy: + # No downstream consumer for this step's drafts. self.draft_tokens_np = None + self.num_valid_draft_tokens_np = None return - # For spec decoding + structured outputs, we must transfer the - # draft tokens back to the scheduler for grammar validation. + # Kick off all D2H copies on a dedicated side stream so that they + # overlap with draft-scatter and any subsequent main-stream work. current_stream = torch.cuda.current_stream(self.device) self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): - self.draft_tokens_np = async_copy_to_np(draft_tokens) + if needs_draft_copy: + self.draft_tokens_np = async_copy_to_np(draft_tokens) + else: + self.draft_tokens_np = None + if needs_valid_copy: + # ``async_copy_to_np`` already issues a non-blocking copy + # and returns a NumPy view over the resulting pinned host + # buffer. We synchronise on ``copy_event`` in + # ``get_draft_tokens`` before materialising the list. + self.num_valid_draft_tokens_np = async_copy_to_np( + num_valid_draft_tokens + ) + else: + self.num_valid_draft_tokens_np = None self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: - if self.draft_tokens_np is not None: + # If either side-band was copied, we must sync the event exactly + # once before reading the pinned buffers. + if ( + self.draft_tokens_np is not None + or self.num_valid_draft_tokens_np is not None + ): self.copy_event.synchronize() + + if self.draft_tokens_np is not None: draft_token_ids = self.draft_tokens_np.tolist() else: - # This case only happens when async scheduling is disabled. + # This branch is taken when async scheduling is disabled AND + # there are no structured-output requests AND the speculator + # did not emit per-request valid counts. The scheduler treats + # an empty/``-1`` placeholder as "drafts were not echoed + # back" — compatible with legacy Eagle behaviour. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] - return DraftTokenIds(self.req_ids, draft_token_ids) + + num_valid_list: list[int] | None = None + if self.num_valid_draft_tokens_np is not None: + # ``.tolist()`` on an int32 ndarray returns ``list[int]``. + num_valid_list = self.num_valid_draft_tokens_np.tolist() + + return DraftTokenIds( + req_ids=self.req_ids, + draft_token_ids=draft_token_ids, + num_valid_draft_tokens=num_valid_list, + ) diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 24d225886106..ad702a2b564a 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -15,6 +15,7 @@ def __init__( num_speculative_steps: int, vocab_size: int, device: torch.device, + use_dense_all_token_ids: bool = False, ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -29,12 +30,14 @@ def __init__( # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs) # depending on the configured max_num_reqs and max_model_len. - # To save GPU memory, we use UVA instead of GPU for this tensor. + # To save GPU memory, we use UVA instead of GPU by default, but + # ngram_gpu benefits from dense device residency because it scans + # active rows repeatedly during proposal. self.all_token_ids = StagedWriteTensor( (self.max_num_reqs, self.max_model_len), dtype=torch.int32, device=device, - uva_instead_of_gpu=True, + uva_instead_of_gpu=not use_dense_all_token_ids, ) # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len: # - prompt_len: Number of tokens in the user-provided prompt. From 4bfb745af0ac6bd5bfd28c580eff7c324c8333da Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 21 Apr 2026 12:00:42 +0800 Subject: [PATCH 02/34] complete async ngram v2 Signed-off-by: PatchouliTaisa --- vllm/envs.py | 9 +++++ vllm/v1/engine/core.py | 36 +++++++++++++++++-- vllm/v1/worker/gpu/model_runner.py | 8 +++-- vllm/v1/worker/gpu/sample/gumbel.py | 33 +++++++++++++---- .../gpu/spec_decode/ngram/speculator.py | 21 +++++++---- 5 files changed, 88 insertions(+), 19 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index c2f8ca8c5808..de8d0a2f890b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -193,6 +193,7 @@ VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False + VLLM_SAMPLER_FP64_GUMBEL: bool = False VLLM_USE_NVFP4_CT_EMULATIONS: bool = False VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ "FP", "INT8", "INT6", "INT4", "NONE" @@ -1423,6 +1424,14 @@ def _get_or_set_default() -> str: "VLLM_COMPUTE_NANS_IN_LOGITS": lambda: bool( int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0")) ), + # If set to 1, the V2 sampler's Gumbel-max kernel will promote logits + # to fp64 and draw a full 64-bit uniform before the double-log transform. + # Default 0 (fp32) matches V1's behaviour and is 16-32x faster on H100 + # without any observable accuracy regression. Intended for statistical + # validation only. + "VLLM_SAMPLER_FP64_GUMBEL": lambda: bool( + int(os.getenv("VLLM_SAMPLER_FP64_GUMBEL", "0")) + ), # Controls whether or not emulations are used for NVFP4 # generations on machines < 100 for compressed-tensors # models diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 0fa59579ee76..0c1dc53a6254 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -70,7 +70,7 @@ from vllm.v1.executor import Executor from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.metrics.stats import SchedulerStats -from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder from vllm.v1.structured_output import StructuredOutputManager @@ -406,7 +406,30 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: scheduler_output, model_output ) - return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 + model_executed = scheduler_output.total_num_scheduled_tokens > 0 + self._maybe_update_async_draft_token_ids(model_executed) + + return engine_core_outputs, model_executed + + def _maybe_update_async_draft_token_ids( + self, model_executed: bool + ) -> "DraftTokenIds | None": + """Consume variable-length draft metadata from the just-completed + batch and apply it to scheduler request state. + + Returns the fetched ``DraftTokenIds`` payload (or ``None``) so that + callers with a deferred structured-output path can reuse it instead + of fetching twice. + """ + if not (self.async_scheduling and self.use_spec_decode and model_executed): + return None + draft_token_ids = self.model_executor.take_draft_token_ids() + if draft_token_ids is None: + return None + if draft_token_ids.num_valid_draft_tokens is None: + return draft_token_ids + self.scheduler.update_draft_token_ids(draft_token_ids) + return draft_token_ids def post_step(self, model_executed: bool) -> None: # When using async scheduling we can't get draft token ids in advance, @@ -510,6 +533,11 @@ def step_with_batch_queue( scheduler_output, model_output ) + popped_batch_executed = scheduler_output.total_num_scheduled_tokens > 0 + async_draft_token_ids = self._maybe_update_async_draft_token_ids( + popped_batch_executed + ) + # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. @@ -518,7 +546,9 @@ def step_with_batch_queue( # we need to get the draft token ids from the prior step before # we can compute the grammar bitmask for the deferred request. if self.use_spec_decode: - draft_token_ids = self.model_executor.take_draft_token_ids() + draft_token_ids = async_draft_token_ids + if draft_token_ids is None: + draft_token_ids = self.model_executor.take_draft_token_ids() assert draft_token_ids is not None # Update the draft token ids in the scheduler output to # filter out the invalid spec tokens, which will be padded diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 86da8ff32b87..3855bb7b13ff 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -45,6 +45,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput from vllm.v1.worker.gpu.attn_utils import ( @@ -1147,9 +1148,10 @@ def sample_tokens( return None # Last rank: sample tokens - sampler_output, num_sampled, num_rejected = self.sample( - hidden_states, input_batch, grammar_output - ) + with record_function_or_nullcontext("gpu_model_runner: sample"): + sampler_output, num_sampled, num_rejected = self.sample( + hidden_states, input_batch, grammar_output + ) if self.use_pp: # Broadcast to non-last PP ranks (handles spec decode multi-token). diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 0d08ceb83bc0..67a0339535ba 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,8 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +import vllm.envs as envs from vllm.triton_utils import tl, triton +# Smallest positive normal fp32 value. Used to clamp the uniform draw so that +# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). fp32 +# normal min is 2**-126 ≈ 1.1754944e-38; `-log(-log(2**-126)) ≈ -4.47`. +_FP32_TINY = tl.constexpr(1.1754943508222875e-38) + @triton.jit def _temperature_kernel( @@ -81,6 +87,7 @@ def _gumbel_sample_kernel( vocab_size, BLOCK_SIZE: tl.constexpr, APPLY_TEMPERATURE: tl.constexpr, + USE_FP64: tl.constexpr, ): token_idx = tl.program_id(0) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) @@ -110,19 +117,26 @@ def _gumbel_sample_kernel( mask=mask, ) - logits = logits.to(tl.float64) + # Promote to the reduction dtype. fp32 is the default — on H100/Ada/Blackwell + # fp64 has 1/32x-1/64x throughput of fp32, and the Gumbel-max result does not + # benefit from the extra precision in any measurable way. The fp64 branch is + # retained behind VLLM_SAMPLER_FP64_GUMBEL=1 for statistical validation. + if USE_FP64: + logits = logits.to(tl.float64) if temp != 0.0: # Calculate the seed for gumbel noise. seed = tl.load(seeds_ptr + req_state_idx) pos = tl.load(pos_ptr + token_idx) gumbel_seed = tl.randint(seed, pos) - # tl.rand returns fp32, so build a true fp64 uniform from 64 random - # bits before applying the double-log transform. - u = tl_rand64(gumbel_seed, block, includes_zero=False) - gumbel_noise = -tl.log(-tl.log(u)) + if USE_FP64: + u = tl_rand64(gumbel_seed, block, includes_zero=False) + gumbel_noise = -tl.log(-tl.log(u)) + else: + u = tl.rand(gumbel_seed, block) + u = tl.maximum(u, _FP32_TINY) + gumbel_noise = -tl.log(-tl.log(u)) - # Apply gumbel noise. logits = tl.where(mask, logits + gumbel_noise, float("-inf")) value, idx = tl.max(logits, axis=0, return_indices=True) @@ -139,12 +153,16 @@ def gumbel_sample( pos: torch.Tensor, # [num_tokens] apply_temperature: bool, processed_logits_out: torch.Tensor | None = None, # [num_reqs, vocab_size] + use_fp64: bool | None = None, ) -> torch.Tensor: + if use_fp64 is None: + use_fp64 = bool(envs.VLLM_SAMPLER_FP64_GUMBEL) num_tokens, vocab_size = logits.shape BLOCK_SIZE = 1024 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE) local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64) - local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float64) + local_max_dtype = torch.float64 if use_fp64 else torch.float32 + local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype) _gumbel_sample_kernel[(num_tokens, num_blocks)]( local_argmax, local_argmax.stride(0), @@ -161,6 +179,7 @@ def gumbel_sample( vocab_size, BLOCK_SIZE=BLOCK_SIZE, APPLY_TEMPERATURE=apply_temperature, + USE_FP64=use_fp64, ) # NOTE(woosuk): Use int64 for later indexing. max_block_idx = local_max.argmax(dim=-1, keepdim=True) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 468794a69c8d..7efc556f9db4 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -36,7 +36,19 @@ from vllm.v1.worker.gpu.states import RequestState -@support_torch_compile() +# NOTE: `from __future__ import annotations` above turns every forward-method +# annotation into a *string*, which defeats `@support_torch_compile`'s +# auto-inference of dynamic dims (it does `v.annotation in [torch.Tensor, ...]` +# with class objects, not strings). We therefore pass `dynamic_arg_dims` +# explicitly. All four tensor inputs have the batch size on dim 0. +@support_torch_compile( + dynamic_arg_dims={ + "token_ids": 0, + "seq_lens": 0, + "valid_mask": 0, + "last_sampled": 0, + } +) class _NgramKernel(nn.Module): """Pure, stateless n-gram match kernel. @@ -253,7 +265,7 @@ def propose( # [num_reqs] int32 — unused by ngram num_rejected: torch.Tensor, # [max_num_reqs, 1] int64 - last_sampled_tokens: torch.Tensor, + last_sampled: torch.Tensor, # [max_num_reqs] int32 — unused next_prefill_tokens: torch.Tensor, # [max_num_reqs] — unused @@ -264,7 +276,6 @@ def propose( dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, - is_profile: bool = False, ) -> torch.Tensor: """Propose up to ``num_speculative_steps`` draft tokens per request. @@ -295,9 +306,7 @@ def propose( idx_mapping_long ] active_seq_lens: torch.Tensor = self.req_states.total_len.gpu[idx_mapping_long] - active_last_sampled: torch.Tensor = last_sampled_tokens.view(-1)[ - idx_mapping_long - ] + active_last_sampled: torch.Tensor = last_sampled.view(-1)[idx_mapping_long] # A request can draft iff (a) at least one real token was just # sampled for it (otherwise we cannot trust its suffix) AND From 02185acd6d2abd262c025d513fe5e91e25edbe9d Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 21 Apr 2026 14:45:54 +0800 Subject: [PATCH 03/34] remove overcomments Signed-off-by: PatchouliTaisa --- vllm/v1/core/sched/scheduler.py | 9 ------- vllm/v1/engine/core.py | 7 ++--- vllm/v1/worker/gpu/model_runner.py | 5 ---- vllm/v1/worker/gpu/spec_decode/utils.py | 35 +++---------------------- 4 files changed, 5 insertions(+), 51 deletions(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 66febc371d37..f3a149be37e4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1667,10 +1667,6 @@ def _free_encoder_inputs(self, request: Request) -> None: self.encoder_cache_manager.free_encoder_input(request, input_id) def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: - # Optional per-request valid-draft-count from variable-length - # speculators (e.g. ngram_gpu). When present, we enumerate with an - # index so each request can be trimmed independently before - # further processing (structured-output grammar validation, etc.). num_valid_list = draft_token_ids.num_valid_draft_tokens for i, (req_id, spec_token_ids) in enumerate( zip( @@ -1690,14 +1686,9 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: continue # Variable-length drafters: truncate to the number of drafts - # that actually matched a real n-gram. Anything beyond that - # count was a safe-fallback substitution (see - # NgramGPUSpeculator) and must not be surfaced as a draft. if num_valid_list is not None: num_valid = num_valid_list[i] if num_valid < len(spec_token_ids): - # Slice defensively; spec_token_ids is a Python list - # produced by ``ndarray.tolist()`` so slicing is O(k). spec_token_ids = spec_token_ids[:num_valid] # Add newly generated spec token ids to the request. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 0c1dc53a6254..89111b2ae39e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -414,12 +414,9 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: def _maybe_update_async_draft_token_ids( self, model_executed: bool ) -> "DraftTokenIds | None": - """Consume variable-length draft metadata from the just-completed + """ + Consume variable-length draft metadata from the just-completed batch and apply it to scheduler request state. - - Returns the fetched ``DraftTokenIds`` payload (or ``None``) so that - callers with a deferred structured-output path can reuse it instead - of fetching twice. """ if not (self.async_scheduling and self.use_spec_decode and model_executed): return None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3855bb7b13ff..177bb4ebb807 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1233,11 +1233,6 @@ def sample_tokens( mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - # Speculators that can produce fewer than `num_speculative_steps` - # real drafts (e.g. ngram_gpu on a no-match row) expose a - # `get_num_valid_draft_tokens` accessor. When present, the - # companion tensor is shipped to the scheduler via the same - # async D2H copy stream as the draft tokens themselves. num_valid_draft_tokens: torch.Tensor | None = None get_num_valid = getattr(self.speculator, "get_num_valid_draft_tokens", None) if get_num_valid is not None: diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 3a2cbaf8dbca..3203e5b6a6a6 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -16,9 +16,6 @@ def __init__(self, device: torch.device | None = None): self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None - # Per-request valid-draft-count, populated by speculators that - # return variable-length drafts (e.g. NgramGPUSpeculator). ``None`` - # preserves legacy fixed-length semantics (Eagle / Medusa). self.num_valid_draft_tokens_np: np.ndarray | None = None self.num_draft_tokens: int = 0 @@ -28,20 +25,6 @@ def set_draft_tokens( draft_tokens: torch.Tensor, num_valid_draft_tokens: torch.Tensor | None = None, ) -> None: - """Stage the per-step draft tokens for async D2H copy. - - Args: - input_batch: The just-executed step's input batch (used for - ``req_ids`` and the structured-outputs flag). - draft_tokens: ``[num_reqs, num_speculative_steps]`` int64 on - device. Must be a final tensor (no further GPU mutation - after this call) so the async copy is safe. - num_valid_draft_tokens: Optional ``[num_reqs]`` int32 on device. - Required for variable-length speculators so the scheduler - can trim per-request drafts before the next scheduling - round. When ``None`` (Eagle-style fixed-length drafts), - all draft positions are treated as valid. - """ self.req_ids = input_batch.req_ids self.num_draft_tokens = draft_tokens.shape[1] @@ -49,13 +32,12 @@ def set_draft_tokens( needs_valid_copy = num_valid_draft_tokens is not None if not needs_draft_copy and not needs_valid_copy: - # No downstream consumer for this step's drafts. self.draft_tokens_np = None self.num_valid_draft_tokens_np = None return - # Kick off all D2H copies on a dedicated side stream so that they - # overlap with draft-scatter and any subsequent main-stream work. + # For spec decoding + structured outputs, we must transfer the + # draft tokens back to the scheduler for grammar validation. current_stream = torch.cuda.current_stream(self.device) self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): @@ -64,10 +46,6 @@ def set_draft_tokens( else: self.draft_tokens_np = None if needs_valid_copy: - # ``async_copy_to_np`` already issues a non-blocking copy - # and returns a NumPy view over the resulting pinned host - # buffer. We synchronise on ``copy_event`` in - # ``get_draft_tokens`` before materialising the list. self.num_valid_draft_tokens_np = async_copy_to_np( num_valid_draft_tokens ) @@ -76,8 +54,6 @@ def set_draft_tokens( self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: - # If either side-band was copied, we must sync the event exactly - # once before reading the pinned buffers. if ( self.draft_tokens_np is not None or self.num_valid_draft_tokens_np is not None @@ -87,16 +63,11 @@ def get_draft_tokens(self) -> DraftTokenIds | None: if self.draft_tokens_np is not None: draft_token_ids = self.draft_tokens_np.tolist() else: - # This branch is taken when async scheduling is disabled AND - # there are no structured-output requests AND the speculator - # did not emit per-request valid counts. The scheduler treats - # an empty/``-1`` placeholder as "drafts were not echoed - # back" — compatible with legacy Eagle behaviour. + # This case only happens when async scheduling is disabled. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] num_valid_list: list[int] | None = None if self.num_valid_draft_tokens_np is not None: - # ``.tolist()`` on an int32 ndarray returns ``list[int]``. num_valid_list = self.num_valid_draft_tokens_np.tolist() return DraftTokenIds( From 242fb8275e4fd3dfa4f038b0d8a694df52744a89 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 23 Apr 2026 16:30:05 +0800 Subject: [PATCH 04/34] set minimal values in FP32 Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/sample/gumbel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 67a0339535ba..9c9c7c60fee6 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -8,7 +8,7 @@ # Smallest positive normal fp32 value. Used to clamp the uniform draw so that # `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). fp32 # normal min is 2**-126 ≈ 1.1754944e-38; `-log(-log(2**-126)) ≈ -4.47`. -_FP32_TINY = tl.constexpr(1.1754943508222875e-38) +_FP32_TINY = tl.constexpr(float.fromhex("0x1p-126")) @triton.jit From b4f1db99816a378fb54c8368073e04e46c2c6383 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 23 Apr 2026 20:37:56 +0800 Subject: [PATCH 05/34] remove debug Signed-off-by: PatchouliTaisa --- vllm/envs.py | 6 +- vllm/v1/worker/gpu/model_runner.py | 12 +- vllm/v1/worker/gpu/sample/gumbel.py | 3 +- .../gpu/spec_decode/ngram/speculator.py | 175 +++--------------- 4 files changed, 36 insertions(+), 160 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index e58c912abece..663f232c9b34 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1418,11 +1418,7 @@ def _get_or_set_default() -> str: "VLLM_COMPUTE_NANS_IN_LOGITS": lambda: bool( int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0")) ), - # If set to 1, the V2 sampler's Gumbel-max kernel will promote logits - # to fp64 and draw a full 64-bit uniform before the double-log transform. - # Default 0 (fp32) matches V1's behaviour and is 16-32x faster on H100 - # without any observable accuracy regression. Intended for statistical - # validation only. + # Optimization for Gumbel sampler. "VLLM_SAMPLER_FP64_GUMBEL": lambda: bool( int(os.getenv("VLLM_SAMPLER_FP64_GUMBEL", "0")) ), diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a39b8943f37c..bf8aa9266267 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -48,7 +48,6 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput -from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput from vllm.v1.worker.gpu.attn_utils import ( @@ -207,9 +206,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) # Inject RequestState into speculators that consume the persistent - # token store directly (e.g. NgramGPUSpeculator). Attribute-based - # injection keeps speculators that don't need it (EagleSpeculator) - # untouched — they simply don't declare the `req_states` attribute. + # token store directly (e.g. NgramGPUSpeculator). if self.speculator is not None and hasattr(self.speculator, "req_states"): self.speculator.req_states = self.req_states @@ -1164,10 +1161,9 @@ def sample_tokens( return None # Last rank: sample tokens - with record_function_or_nullcontext("gpu_model_runner: sample"): - sampler_output, num_sampled, num_rejected = self.sample( - hidden_states, input_batch, grammar_output - ) + sampler_output, num_sampled, num_rejected = self.sample( + hidden_states, input_batch, grammar_output + ) if self.use_pp: # Broadcast to non-last PP ranks (handles spec decode multi-token). diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index d266b90d2a6a..f4cb7b97813c 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -6,8 +6,7 @@ from vllm.triton_utils import tl, triton # Smallest positive normal fp32 value. Used to clamp the uniform draw so that -# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). fp32 -# normal min is 2**-126 ≈ 1.1754944e-38; `-log(-log(2**-126)) ≈ -4.47`. +# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). _FP32_TINY = tl.constexpr(float.fromhex("0x1p-126")) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 7efc556f9db4..c1228b33231e 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -1,24 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""V2 GPU-native N-gram speculator. - -Unlike V1's ``NgramProposerGPU`` (in ``vllm/v1/spec_decode/ngram_proposer_gpu.py``), -this speculator is designed to slot directly into the V2 runner's speculator -interface. It reuses the runner-owned ``RequestState.all_token_ids`` / -``RequestState.total_len`` tensors as the persistent token store and thus -does not maintain any shadow GPU state. - -Key properties: - * Fully vectorized ``unfold → broadcast-compare → argmax → gather`` algorithm - — no CPU-GPU synchronization inside ``propose``. - * ``-1`` no-match positions are rewritten on GPU to ``last_sampled_tokens`` - so that the downstream ``combine_sampled_and_draft_tokens`` kernel never - writes an invalid token id into ``input_ids``. - * A companion ``num_valid_draft_tokens`` tensor is emitted every step and - propagated to the scheduler so that ``request.spec_token_ids`` can be - correctly truncated before the next scheduling round. -""" - from __future__ import annotations from typing import TYPE_CHECKING, Any @@ -36,11 +17,6 @@ from vllm.v1.worker.gpu.states import RequestState -# NOTE: `from __future__ import annotations` above turns every forward-method -# annotation into a *string*, which defeats `@support_torch_compile`'s -# auto-inference of dynamic dims (it does `v.annotation in [torch.Tensor, ...]` -# with class objects, not strings). We therefore pass `dynamic_arg_dims` -# explicitly. All four tensor inputs have the batch size on dim 0. @support_torch_compile( dynamic_arg_dims={ "token_ids": 0, @@ -50,12 +26,7 @@ } ) class _NgramKernel(nn.Module): - """Pure, stateless n-gram match kernel. - - Exposed as an ``nn.Module`` so that ``@support_torch_compile`` can hoist - the whole body into a single Inductor-compiled region. The module has - no parameters; it is compiled purely for operator fusion. - """ + """GPU-accelerated N-gram proposer using fully async tensor operations.""" def __init__(self, min_n: int, max_n: int, k: int): super().__init__() @@ -75,91 +46,60 @@ def forward( valid_mask: torch.Tensor, # [B] bool (row eligible for n-gram lookup) last_sampled: torch.Tensor, # [B] int64 (fallback for -1 positions) ) -> tuple[torch.Tensor, torch.Tensor]: - """Return ``(draft_tokens[B, k] int64, num_valid_draft_tokens[B] int32)``. - - The ``draft_tokens`` tensor is guaranteed to contain only valid - vocabulary IDs. Positions that correspond to no-match / out-of-range - slots are filled with ``last_sampled`` so that writing them to - ``input_ids`` later is safe. The companion ``num_valid_draft_tokens`` - encodes how many *leading contiguous* draft positions were actually - derived from a real n-gram match. - """ B, L = token_ids.shape device = token_ids.device - # --- (1) For each n-gram size in [min_n, max_n], find the earliest - # prior occurrence of the trailing suffix. The outer `for` loop runs - # at most `max_n - min_n + 1` times (typically <= 4), and is required - # because `Tensor.unfold` needs a Python int window size. Each - # iteration executes as one fused GPU pass. first_match_pos = torch.full( (B, self.num_sizes), -1, dtype=torch.long, device=device ) batch_idx = torch.arange(B, device=device) for i, n in enumerate(range(self.min_n, self.max_n + 1)): - # [B, L-n+1, n] view — O(1) cost (contiguous stride trick). windows = token_ids.unfold(1, n, 1) num_windows = windows.shape[1] - # Gather the trailing n tokens from each sequence as the query. - # `suffix_start` is clamped to 0 to keep the gather in-bounds; the - # corresponding row is also filtered by `valid_mask` upstream. suffix_start = (seq_lens.long() - n).clamp(min=0) offsets = torch.arange(n, device=device) - suffix_idx = suffix_start.unsqueeze(1) + offsets # [B, n] - suffix = torch.gather(token_ids, 1, suffix_idx) # [B, n] + suffix_idx = suffix_start.unsqueeze(1) + offsets + suffix = torch.gather(token_ids, 1, suffix_idx) - # Element-wise equality → all(-1) → per-window boolean match. - matches = (windows == suffix.unsqueeze(1)).all(dim=-1) # [B, L-n+1] + matches = (windows == suffix.unsqueeze(1)).all(dim=-1) - # A match is only actionable if at least one draft token follows. - max_valid_pos = seq_lens.long() - n - 1 # may be negative + max_valid_pos = seq_lens.long() - n - 1 window_pos = torch.arange(num_windows, device=device) matches = matches & (window_pos.unsqueeze(0) <= max_valid_pos.unsqueeze(1)) - # `argmax(int)` returns 0 on all-false rows; verify with a lookup. - idx = matches.int().argmax(dim=1) # [B] - has_match = matches[batch_idx, idx] # [B] + idx = matches.int().argmax(dim=1) + has_match = matches[batch_idx, idx] first_match_pos[:, i] = torch.where(has_match, idx.long(), -1) - # --- (2) Pick the LONGEST n-gram with a match (flip+argmax, no sync). - # Equivalent to "last True from the right". best_i = (first_match_pos >= 0).int().flip(dims=[1]).argmax(dim=1) - best_i = self.num_sizes - 1 - best_i # flip index back - best_pos = first_match_pos[batch_idx, best_i] # [B] + best_i = self.num_sizes - 1 - best_i + best_pos = first_match_pos[batch_idx, best_i] ngram_lens_table = torch.arange( self.min_n, self.max_n + 1, device=device, dtype=torch.long ) - best_n = ngram_lens_table[best_i] # [B] - has_any = best_pos >= 0 # [B] bool + best_n = ngram_lens_table[best_i] + has_any = best_pos >= 0 - # --- (3) Gather k tokens starting right after the matched suffix. draft_start = torch.where( has_any, best_pos + best_n, torch.zeros_like(best_pos), - ) # [B] - k_range = torch.arange(self.k, device=device) # [k] + ) + k_range = torch.arange(self.k, device=device) draft_idx = (draft_start.unsqueeze(1) + k_range).clamp_(0, L - 1) - drafts = torch.gather(token_ids, 1, draft_idx).to(torch.int64) # [B, k] + drafts = torch.gather(token_ids, 1, draft_idx).to(torch.int64) - # --- (4) Compute the leading-valid mask per row. - tokens_available = (seq_lens.long() - draft_start).clamp_(min=0) # [B] + tokens_available = (seq_lens.long() - draft_start).clamp_(min=0) valid_positions = k_range.unsqueeze(0) < tokens_available.unsqueeze(1) - row_valid = has_any & valid_mask # [B] - leading_valid_mask = valid_positions & row_valid.unsqueeze(1) # [B, k] - - # --- (5) num_valid = length of the leading contiguous valid run. - # cumsum counts valid positions so far; match against [1, 2, ..., k] - # tells us whether each position is still in the leading run. - cum_valid = leading_valid_mask.int().cumsum(dim=1) # [B, k] - positions = torch.arange(1, self.k + 1, device=device) # [k] - num_valid = (cum_valid == positions.unsqueeze(0)).int().sum(dim=1) # [B] - - # --- (6) Safe-substitute invalid slots with last_sampled so that the - # next step's input_ids are always in-vocabulary. `last_sampled` is - # int64 already (matches RequestState.last_sampled_tokens dtype). + row_valid = has_any & valid_mask + leading_valid_mask = valid_positions & row_valid.unsqueeze(1) + + cum_valid = leading_valid_mask.int().cumsum(dim=1) + positions = torch.arange(1, self.k + 1, device=device) + num_valid = (cum_valid == positions.unsqueeze(0)).int().sum(dim=1) + safe_drafts = torch.where( leading_valid_mask, drafts, @@ -169,22 +109,12 @@ def forward( class NgramGPUSpeculator: - """V2-compatible GPU n-gram speculator. - - Public surface mirrors ``EagleSpeculator`` (``load_model``, ``set_attn``, - ``init_cudagraph_manager``, ``capture_model``, ``propose``). Unlike the - Eagle path, this speculator: - - * has no neural model, no attention layers, no CUDA graph capture; - * reuses the runner-owned ``RequestState`` persistent token store; - * emits ``num_valid_draft_tokens`` as a secondary output so that the - scheduler can trim per-request drafts that were effectively a miss. + """ + V2-compatible GPU n-gram speculator. """ - # Consumed by the runner's spec-decode sample path — see - # ``vllm/v1/worker/gpu/model_runner.py::sample``. supports_mm_inputs = False - draft_logits = None # probabilistic rejection sampling is not supported + draft_logits = None def __init__(self, vllm_config: VllmConfig, device: torch.device): spec = vllm_config.speculative_config @@ -217,20 +147,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): .eval() ) - # Persistent scratch buffer — shaped [max_num_reqs] on device — used to - # stash the last ``num_valid_draft_tokens`` tensor so that the runner - # can forward it to ``DraftTokensHandler`` after ``propose`` returns. self.num_valid_draft_tokens: torch.Tensor = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=device ) - # Injected by the model runner once ``req_states`` is constructed. - # Accessed inside ``propose``. self.req_states: RequestState | None = None - # ------------------------------------------------------------------ - # V2 Speculator interface: intentional no-ops. - # ------------------------------------------------------------------ def load_model(self, target_model: nn.Module) -> None: """No weights to load — ngram is a data-only proposer.""" pass @@ -247,45 +169,25 @@ def capture_model(self) -> None: """No graph capture phase required.""" pass - # ------------------------------------------------------------------ - # V2 Speculator interface: main entry point. - # ------------------------------------------------------------------ @torch.inference_mode() def propose( self, input_batch: InputBatch, - attn_metadata: Any, # unused - slot_mappings: Any, # unused - # [num_tokens, hidden_size] — unused by ngram, kept for signature parity + attn_metadata: Any, + slot_mappings: Any, last_hidden_states: torch.Tensor, - # num_layers x [num_tokens, hidden_size] — unused aux_hidden_states: list[torch.Tensor] | None, - # [num_reqs] int32 num_sampled: torch.Tensor, - # [num_reqs] int32 — unused by ngram num_rejected: torch.Tensor, - # [max_num_reqs, 1] int64 last_sampled: torch.Tensor, - # [max_num_reqs] int32 — unused next_prefill_tokens: torch.Tensor, - # [max_num_reqs] — unused temperature: torch.Tensor, - # [max_num_reqs] — unused seeds: torch.Tensor, num_tokens_across_dp: torch.Tensor | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, ) -> torch.Tensor: - """Propose up to ``num_speculative_steps`` draft tokens per request. - - Returns a ``[num_reqs, num_speculative_steps]`` int64 tensor. Invalid - or no-match positions are backfilled with ``last_sampled`` rather - than ``-1`` so that downstream ``combine_sampled_and_draft_tokens`` - never writes out-of-vocab ids into ``input_ids``. Truth about how - many of those drafts are "real" lives in - ``self.num_valid_draft_tokens[:num_reqs]``. - """ assert self.req_states is not None, ( "NgramGPUSpeculator.req_states was not injected by the model " "runner. Ensure model_runner sets `speculator.req_states = " @@ -293,24 +195,14 @@ def propose( ) num_reqs = input_batch.num_reqs - # `idx_mapping` is [num_reqs] int32 on device; advanced indexing - # requires int64 → one-time cast (negligible cost vs. the - # O(B * L * num_sizes) compare). idx_mapping_long = input_batch.idx_mapping.long() - # Persistent [max_num_reqs, max_model_len] store, UVA-backed. - # Advanced indexing materialises a contiguous [num_reqs, L] view on - # device — identical access pattern to V1's shadow tensor. No extra - # H2D copies are issued. active_tokens: torch.Tensor = self.req_states.all_token_ids.gpu[ idx_mapping_long ] active_seq_lens: torch.Tensor = self.req_states.total_len.gpu[idx_mapping_long] active_last_sampled: torch.Tensor = last_sampled.view(-1)[idx_mapping_long] - # A request can draft iff (a) at least one real token was just - # sampled for it (otherwise we cannot trust its suffix) AND - # (b) the sequence already contains min_n tokens for the lookup. valid_mask = (num_sampled > 0) & (active_seq_lens >= self.min_n) with set_forward_context(None, self.vllm_config): @@ -321,21 +213,14 @@ def propose( active_last_sampled, ) - # Stash num_valid so the runner can forward it to DraftTokensHandler - # after scatter. Note: this is a device tensor; D2H happens on a - # side stream in the handler. self.num_valid_draft_tokens[:num_reqs].copy_(num_valid) - # Zero out the tail slots to avoid leaking stale values from prior - # steps in case the runner ever peeks beyond num_reqs. if num_reqs < self.max_num_reqs: self.num_valid_draft_tokens[num_reqs:].zero_() - return drafts # [num_reqs, num_speculative_steps] int64, no -1 + return drafts def get_num_valid_draft_tokens(self, num_reqs: int) -> torch.Tensor: - """Return the last step's per-request valid draft counts. - - Sliced view of the internal buffer; safe to pass to an async D2H - copy on a side stream. + """ + Return the last step's per-request valid draft counts. """ return self.num_valid_draft_tokens[:num_reqs] From f22224a93c6569b8ce2d635313f367cfcf93d8dd Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Fri, 24 Apr 2026 11:35:57 +0800 Subject: [PATCH 06/34] reorg codes Signed-off-by: PatchouliTaisa --- vllm/v1/engine/core.py | 13 +++++++------ vllm/v1/worker/gpu/sample/gumbel.py | 11 +++-------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index eb1c77e2a0c1..1cfb674f1489 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -444,9 +444,8 @@ def _maybe_update_async_draft_token_ids( draft_token_ids = self.model_executor.take_draft_token_ids() if draft_token_ids is None: return None - if draft_token_ids.num_valid_draft_tokens is None: - return draft_token_ids - self.scheduler.update_draft_token_ids(draft_token_ids) + if draft_token_ids.num_valid_draft_tokens is not None: + self.scheduler.update_draft_token_ids(draft_token_ids) return draft_token_ids def post_step(self, model_executed: bool) -> None: @@ -564,9 +563,11 @@ def step_with_batch_queue( # we need to get the draft token ids from the prior step before # we can compute the grammar bitmask for the deferred request. if self.use_spec_decode: - draft_token_ids = async_draft_token_ids - if draft_token_ids is None: - draft_token_ids = self.model_executor.take_draft_token_ids() + draft_token_ids = ( + async_draft_token_ids + if async_draft_token_ids is not None + else self.model_executor.take_draft_token_ids() + ) assert draft_token_ids is not None # Update the draft token ids in the scheduler output to # filter out the invalid spec tokens, which will be padded diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index f4cb7b97813c..da8a16287807 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -23,7 +23,7 @@ def _temperature_kernel( req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32) if temperature == 0.0 or temperature == 1.0: - # Early return to avoid loading logits. + # Greedy or no-op rescale: avoid loading logits at all. return block_idx = tl.program_id(1) @@ -87,27 +87,22 @@ def gumbel_block_argmax( req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) if temp != 0.0 and APPLY_TEMPERATURE: - # Apply temperature. # NOTE(woosuk): Match the behavior of _temperature_kernel. # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too. logits = logits / temp if processed_logits_ptr is not None: - # Store the temperature-applied logits. tl.store( processed_logits_ptr + req_state_idx * processed_logits_stride + block, logits, mask=mask, ) - # Promote to the reduction dtype. fp32 is the default — on H100/Ada/Blackwell - # fp64 has 1/32x-1/64x throughput of fp32, and the Gumbel-max result does not - # benefit from the extra precision in any measurable way. The fp64 branch is - # retained behind VLLM_SAMPLER_FP64_GUMBEL=1 for statistical validation. + # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput + # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max. if USE_FP64: logits = logits.to(tl.float64) if temp != 0.0: - # Calculate the seed for gumbel noise. seed = tl.load(seeds_ptr + req_state_idx) pos = tl.load(pos_ptr + token_idx) gumbel_seed = tl.randint(seed, pos) From 3af6547920ca16da0889b498c6bb4489c499e902 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 28 Apr 2026 16:18:24 +0800 Subject: [PATCH 07/34] fix argmax Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/spec_decode/ngram/speculator.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index c1228b33231e..f4f4602f3704 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -69,9 +69,10 @@ def forward( window_pos = torch.arange(num_windows, device=device) matches = matches & (window_pos.unsqueeze(0) <= max_valid_pos.unsqueeze(1)) - idx = matches.int().argmax(dim=1) + matched_indices = torch.where(matches, window_pos.unsqueeze(0), -1) + idx = matched_indices.argmax(dim=1) has_match = matches[batch_idx, idx] - first_match_pos[:, i] = torch.where(has_match, idx.long(), -1) + first_match_pos[:, i] = torch.where(has_match, idx, -1) best_i = (first_match_pos >= 0).int().flip(dims=[1]).argmax(dim=1) best_i = self.num_sizes - 1 - best_i From a285ef63e69c7049c761c1c24eacea2c1b247e34 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Mon, 11 May 2026 19:26:33 +0800 Subject: [PATCH 08/34] test added for ngram gpu Signed-off-by: PatchouliTaisa --- tests/v1/spec_decode/test_max_len.py | 23 + tests/v1/spec_decode/test_ngram_gpu.py | 573 +++++++++++++++++++++++++ 2 files changed, 596 insertions(+) create mode 100644 tests/v1/spec_decode/test_ngram_gpu.py diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 77c041d84a94..4c808654abee 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -34,6 +34,29 @@ def test_ngram_max_len(num_speculative_tokens: int): llm.generate(_PROMPTS, sampling_params) +@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) +def test_ngram_gpu_max_len(num_speculative_tokens: int): + """V2 GPU n-gram counterpart of ``test_ngram_max_len``. + + Verifies that the V2 model runner with ``method="ngram_gpu"`` correctly + handles the ``max_model_len`` boundary across various speculative-token + counts. + """ + llm = LLM( + model="facebook/opt-125m", + max_model_len=100, + enforce_eager=True, # For faster initialization. + speculative_config={ + "method": "ngram_gpu", + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": num_speculative_tokens, + }, + ) + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + llm.generate(_PROMPTS, sampling_params) + + @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) def test_eagle_max_len( diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py new file mode 100644 index 000000000000..5da7064e4f35 --- /dev/null +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -0,0 +1,573 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU-accelerated n-gram speculator. + +These tests target the kernel and proposer logic in +``vllm.v1.worker.gpu.spec_decode.ngram.speculator`` and complement the CPU +``NgramProposer`` tests in ``test_ngram.py``. The GPU speculator follows a +slightly different policy than the CPU one: when multiple n-gram matches of the +same length exist, the GPU kernel picks the right-most (most recent) match +inside the active context, whereas the CPU implementation returns the +left-most. The expectations below reflect the GPU behavior. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import ( + ModelConfig, + SchedulerConfig, + SpeculativeConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( + NgramGPUSpeculator, + _NgramKernel, +) + +# The kernel uses ``@support_torch_compile`` and ``set_forward_context``, both +# of which expect a real CUDA device. Skip these tests gracefully on +# CPU-only or non-CUDA platforms. +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for NgramGPUSpeculator tests", + allow_module_level=True, + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_vllm_config( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 8, + max_model_len: int = 64, +) -> VllmConfig: + """Build a minimal VllmConfig configured for ngram_gpu speculative decoding. + + ``enforce_eager=True`` is used so the kernel runs in eager mode (no + torch.compile) — sufficient for unit-testing the math. + """ + model_config = ModelConfig( + model="facebook/opt-125m", + max_model_len=max_model_len, + enforce_eager=True, + ) + scheduler_config = SchedulerConfig(max_num_seqs=max_num_seqs) + speculative_config = SpeculativeConfig( + method="ngram_gpu", + prompt_lookup_min=min_n, + prompt_lookup_max=max_n, + num_speculative_tokens=k, + ) + return VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + speculative_config=speculative_config, + ) + + +def _make_kernel(min_n: int, max_n: int, k: int) -> _NgramKernel: + """Construct a ``_NgramKernel`` bound to a minimal ``VllmConfig``. + + The kernel is decorated by ``@support_torch_compile`` whose generated + ``__init__`` accepts ``vllm_config`` as a kwarg. + """ + vllm_config = _make_vllm_config(min_n=min_n, max_n=max_n, k=k) + with set_current_vllm_config(vllm_config): + kernel = _NgramKernel( + min_n=min_n, + max_n=max_n, + k=k, + vllm_config=vllm_config, + ) + return kernel.to("cuda").eval() + + +def _pad_tokens(rows: list[list[int]], pad_to: int) -> torch.Tensor: + """Right-pad ragged ``rows`` to ``pad_to`` length with zeros (int32).""" + out = torch.zeros((len(rows), pad_to), dtype=torch.int32) + for i, row in enumerate(rows): + if row: + out[i, : len(row)] = torch.tensor(row, dtype=torch.int32) + return out + + +def _run_kernel( + kernel: _NgramKernel, + rows: list[list[int]], + seq_lens: list[int], + valid_mask: list[bool] | None = None, + last_sampled: list[int] | None = None, + pad_to: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build inputs from python lists, run the kernel, return (drafts, num_valid).""" + B = len(rows) + if pad_to is None: + pad_to = max((len(r) for r in rows), default=1) + # The kernel needs at least 1 column to satisfy unfold(1, n, 1). + pad_to = max(pad_to, kernel.max_n) + + if valid_mask is None: + valid_mask = [True] * B + if last_sampled is None: + last_sampled = [0] * B + + token_ids = _pad_tokens(rows, pad_to).cuda() + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + valid_mask_t = torch.tensor(valid_mask, dtype=torch.bool, device="cuda") + last_sampled_t = torch.tensor(last_sampled, dtype=torch.int64, device="cuda") + with torch.inference_mode(): + drafts, num_valid = kernel( + token_ids, + seq_lens_t, + valid_mask_t, + last_sampled_t, + ) + return drafts.cpu(), num_valid.cpu() + + +# --------------------------------------------------------------------------- +# _NgramKernel tests (mirror & extend tests/v1/spec_decode/test_ngram.py) +# --------------------------------------------------------------------------- + + +def test_kernel_no_match_returns_zero_valid(): + """No 2-gram match in [1,2,3,4,5] → kernel reports 0 valid drafts.""" + kernel = _make_kernel(min_n=2, max_n=2, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 4, 5]], + seq_lens=[5], + last_sampled=[42], + ) + assert num_valid.tolist() == [0] + # Padding positions must fall back to ``last_sampled`` for invalid rows. + assert drafts.tolist() == [[42, 42]] + + +def test_kernel_no_4gram_match_only(): + """No 4-gram match in [1,2,3,4,1,2,3] → 0 valid drafts.""" + kernel = _make_kernel(min_n=4, max_n=4, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 4, 1, 2, 3]], + seq_lens=[7], + last_sampled=[7], + ) + assert num_valid.tolist() == [0] + assert drafts.tolist() == [[7, 7]] + + +def test_kernel_falls_back_to_3gram_when_4gram_missing(): + """No 4-gram match but a 3-gram match exists → propose [4, 1].""" + kernel = _make_kernel(min_n=3, max_n=4, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 4, 1, 2, 3]], + seq_lens=[7], + ) + assert num_valid.tolist() == [2] + assert drafts.tolist() == [[4, 1]] + + +def test_kernel_prefers_longer_ngram(): + """Both a 4-gram (1,2,3,4) and a 3-gram match are present. + + The kernel must prefer the longer match, returning [1, 2] instead of [5, 1]. + """ + kernel = _make_kernel(min_n=3, max_n=4, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]], + seq_lens=[12], + ) + assert num_valid.tolist() == [2] + assert drafts.tolist() == [[1, 2]] + + +def test_kernel_picks_longest_match_among_2_3_4_grams(): + """2-gram and 3-gram match, 4-gram does not → propose 3-gram match [1, 2].""" + kernel = _make_kernel(min_n=2, max_n=4, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]], + seq_lens=[10], + ) + assert num_valid.tolist() == [2] + assert drafts.tolist() == [[1, 2]] + + +def test_kernel_picks_rightmost_when_multiple_matches(): + """Multiple 3-gram matches exist for suffix (1,2,3). + + The GPU kernel picks the RIGHT-most (most recent) match, unlike the CPU + proposer which picks the left-most. Tokens after the right-most match + starting at index 8 are [300, 1]. + """ + kernel = _make_kernel(min_n=3, max_n=3, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]], + seq_lens=[15], + ) + assert num_valid.tolist() == [2] + assert drafts.tolist() == [[300, 1]] + + +def test_kernel_short_context_yields_zero_valid(): + """Context length < min_n → kernel returns 0 valid drafts. + + Note: callers (NgramGPUSpeculator.propose) gate this via ``valid_mask``. + Here we exercise the kernel directly with a row whose seq_len < min_n. + The only window of length 2 cannot match a length-2 suffix because it + overlaps the suffix itself, so no valid match is produced. + """ + kernel = _make_kernel(min_n=2, max_n=2, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[5, 6]], + seq_lens=[2], + last_sampled=[99], + pad_to=4, + ) + assert num_valid.tolist() == [0] + assert drafts.tolist() == [[99, 99]] + + +def test_kernel_valid_mask_disables_row(): + """``valid_mask=False`` disables proposals for that row regardless of match.""" + kernel = _make_kernel(min_n=2, max_n=2, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 1, 2]], # has a 2-gram match for suffix (1,2) → [3, 1] + seq_lens=[5], + valid_mask=[False], + last_sampled=[77], + ) + assert num_valid.tolist() == [0] + # With the row disabled, fallback ``last_sampled`` fills the draft slots. + assert drafts.tolist() == [[77, 77]] + + +def test_kernel_truncates_num_valid_when_few_tokens_after_match(): + """When fewer than k tokens are available after the match, num_valid < k. + + Tokens: [1, 2, 1, 2] (seq_len=4). Suffix (1, 2) matches at position 0 + (the right-most match at position 2 is the suffix itself and is excluded + by ``max_valid_pos = seq_len - n - 1``). With k=3, only the first 2 slots + fit within the active context; the third slot must fall back to + ``last_sampled``. + """ + kernel = _make_kernel(min_n=2, max_n=2, k=3) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 1, 2]], + seq_lens=[4], + last_sampled=[55], + pad_to=4, + ) + # Only the first 2 speculative slots map to tokens inside the context. + assert num_valid.tolist() == [2] + drafts_row = drafts.tolist()[0] + # The first 2 drafts come from the tokens following the match. + assert drafts_row[:2] == [1, 2] + # The third slot must fall back to ``last_sampled`` since it is invalid. + assert drafts_row[2] == 55 + + +def test_kernel_multibatch_mixed(): + """Mixed batch: row 0 matches, row 1 has no match.""" + kernel = _make_kernel(min_n=2, max_n=2, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 1, 2], [4, 5, 6, 0, 0]], + seq_lens=[5, 3], + last_sampled=[10, 20], + ) + assert num_valid.tolist() == [2, 0] + assert drafts.tolist()[0] == [3, 1] + # Row 1: no match → all slots filled with last_sampled[1] = 20. + assert drafts.tolist()[1] == [20, 20] + + +def test_kernel_multibatch_independent_choice_of_n(): + """Row 0 matches as 3-gram, row 1 only matches as 2-gram. + + Verifies that each row independently picks its longest matched n. + """ + kernel = _make_kernel(min_n=2, max_n=3, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[ + [9, 1, 2, 3, 8, 1, 2, 3], # 3-gram (1,2,3) matches at idx 1 → [8, 1] + [7, 1, 2, 9, 1, 2, 0, 0], # 2-gram (1,2) matches at idx 1 → [9, 1] + ], + seq_lens=[8, 6], + ) + assert num_valid.tolist() == [2, 2] + assert drafts.tolist()[0] == [8, 1] + assert drafts.tolist()[1] == [9, 1] + + +def test_kernel_min_n_eq_1(): + """min_n=max_n=1 — single-token n-grams always match if context > 1.""" + kernel = _make_kernel(min_n=1, max_n=1, k=2) + drafts, num_valid = _run_kernel( + kernel, + rows=[[1, 2, 3, 4, 1]], # suffix (1,) matches at idx 0 → tokens after = [2, 3] + seq_lens=[5], + ) + assert num_valid.tolist() == [2] + assert drafts.tolist() == [[2, 3]] + + +# --------------------------------------------------------------------------- +# NgramGPUSpeculator.propose tests +# --------------------------------------------------------------------------- + + +class _FakeStaged: + """Lightweight stand-in for ``StagedWriteTensor`` used by ``RequestState``. + + Only ``.gpu`` is consulted by ``NgramGPUSpeculator.propose``. + """ + + def __init__(self, tensor: torch.Tensor): + self.gpu = tensor + + +class _FakeRequestState: + """Minimal duck-typed ``RequestState`` exposing only the fields propose() reads.""" + + def __init__(self, all_token_ids: torch.Tensor, total_len: torch.Tensor): + self.all_token_ids = _FakeStaged(all_token_ids) + self.total_len = _FakeStaged(total_len) + + +def _make_speculator( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 4, + max_model_len: int = 32, +) -> NgramGPUSpeculator: + cfg = _make_vllm_config( + min_n=min_n, + max_n=max_n, + k=k, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) + with set_current_vllm_config(cfg): + spec = NgramGPUSpeculator(vllm_config=cfg, device=torch.device("cuda")) + return spec + + +def test_speculator_propose_basic_single_request(): + """propose() should populate drafts and num_valid_draft_tokens correctly.""" + spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) + + # State for a single active request with tokens [1, 2, 3, 1, 2]; suffix (1,2) + # matches at index 0 → expected draft tokens [3, 1]. + all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") + all_token_ids[2, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) + total_len = torch.zeros(4, dtype=torch.int32, device="cuda") + total_len[2] = 5 + + spec.req_states = _FakeRequestState(all_token_ids, total_len) + + # The active batch contains a single request mapping to slot index 2. + input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping=torch.tensor([2], dtype=torch.int32, device="cuda"), + ) + last_sampled = torch.full((4, 1), 99, dtype=torch.int64, device="cuda") + num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") + + drafts = spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device="cuda"), + aux_hidden_states=None, + num_sampled=num_sampled, + num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), + last_sampled=last_sampled, + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), + temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), + seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + + assert drafts.shape == (1, 2) + assert drafts.cpu().tolist() == [[3, 1]] + assert spec.get_num_valid_draft_tokens(1).cpu().tolist() == [2] + + +def test_speculator_propose_zero_sampled_disables_proposal(): + """When num_sampled==0 for a request, valid_mask is False → 0 valid drafts.""" + spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) + + all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") + all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) + total_len = torch.zeros(4, dtype=torch.int32, device="cuda") + total_len[0] = 5 + + spec.req_states = _FakeRequestState(all_token_ids, total_len) + + input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + ) + last_sampled = torch.full((4, 1), 7, dtype=torch.int64, device="cuda") + # num_sampled=0 must disable speculation for this request. + num_sampled = torch.tensor([0], dtype=torch.int32, device="cuda") + + spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device="cuda"), + aux_hidden_states=None, + num_sampled=num_sampled, + num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), + last_sampled=last_sampled, + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), + temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), + seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + + assert spec.get_num_valid_draft_tokens(1).cpu().tolist() == [0] + + +def test_speculator_propose_multibatch_noncontiguous_idx_mapping(): + """Verify that propose() correctly reads via idx_mapping (non-contiguous).""" + spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) + + all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") + # Request A at slot 0: tokens [1,2,3,1,2] → drafts [3,1] + all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) + # Request B at slot 3: tokens [7,8,9,7,8] → drafts [9,7] + all_token_ids[3, :5] = torch.tensor([7, 8, 9, 7, 8], dtype=torch.int32) + + total_len = torch.zeros(4, dtype=torch.int32, device="cuda") + total_len[0] = 5 + total_len[3] = 5 + + spec.req_states = _FakeRequestState(all_token_ids, total_len) + + # Batch order is [slot 3, slot 0] — verify the idx_mapping is honored. + input_batch = SimpleNamespace( + num_reqs=2, + idx_mapping=torch.tensor([3, 0], dtype=torch.int32, device="cuda"), + ) + last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") + num_sampled = torch.tensor([1, 1], dtype=torch.int32, device="cuda") + + drafts = spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device="cuda"), + aux_hidden_states=None, + num_sampled=num_sampled, + num_rejected=torch.zeros(2, dtype=torch.int32, device="cuda"), + last_sampled=last_sampled, + next_prefill_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + temperature=torch.zeros(2, dtype=torch.float32, device="cuda"), + seeds=torch.zeros(2, dtype=torch.int64, device="cuda"), + ) + + assert drafts.cpu().tolist() == [[9, 7], [3, 1]] + assert spec.get_num_valid_draft_tokens(2).cpu().tolist() == [2, 2] + + +def test_speculator_propose_resets_unused_slots(): + """num_valid_draft_tokens must be zeroed beyond ``num_reqs``.""" + spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) + + all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") + all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) + total_len = torch.zeros(4, dtype=torch.int32, device="cuda") + total_len[0] = 5 + + spec.req_states = _FakeRequestState(all_token_ids, total_len) + + # Pre-populate with stale data to verify the slots are zeroed. + spec.num_valid_draft_tokens.fill_(123) + + input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + ) + last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") + num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") + + spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device="cuda"), + aux_hidden_states=None, + num_sampled=num_sampled, + num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), + last_sampled=last_sampled, + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), + temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), + seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + + # Slot 0 should report 2 valid drafts; slots 1..3 must be zeroed. + assert spec.num_valid_draft_tokens.cpu().tolist() == [2, 0, 0, 0] + + +def test_speculator_propose_requires_req_states(): + """propose() must assert that req_states has been injected by the model runner.""" + spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=2, max_model_len=8) + # req_states is None by default and should trigger an AssertionError. + assert spec.req_states is None + + input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + ) + last_sampled = torch.zeros((2, 1), dtype=torch.int64, device="cuda") + num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") + + with pytest.raises(AssertionError, match="req_states"): + spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device="cuda"), + aux_hidden_states=None, + num_sampled=num_sampled, + num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), + last_sampled=last_sampled, + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), + temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), + seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + + +def test_speculator_construction_validates_speculative_config(): + """NgramGPUSpeculator requires prompt_lookup_min and prompt_lookup_max to be set.""" + cfg = _make_vllm_config(min_n=2, max_n=3, k=2) + with set_current_vllm_config(cfg): + spec = NgramGPUSpeculator(vllm_config=cfg, device=torch.device("cuda")) + assert spec.min_n == 2 + assert spec.max_n == 3 + assert spec.num_speculative_steps == 2 + # No-op hooks must not raise. + spec.load_model(target_model=None) + spec.set_attn() + spec.capture_model() From fcdab55f85f629143cd6f6b308e864510886323d Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 12 May 2026 14:19:59 +0800 Subject: [PATCH 09/34] fix return value Signed-off-by: PatchouliTaisa --- tests/v1/spec_decode/test_ngram_gpu.py | 33 ++++++++++--------- vllm/v1/worker/gpu/model_runner.py | 6 +--- .../gpu/spec_decode/eagle/speculator.py | 6 ++-- .../gpu/spec_decode/ngram/speculator.py | 19 ++--------- 4 files changed, 24 insertions(+), 40 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 5da7064e4f35..273860c40657 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -374,7 +374,7 @@ def _make_speculator( def test_speculator_propose_basic_single_request(): - """propose() should populate drafts and num_valid_draft_tokens correctly.""" + """propose() should return drafts plus per-request valid counts.""" spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) # State for a single active request with tokens [1, 2, 3, 1, 2]; suffix (1,2) @@ -394,7 +394,7 @@ def test_speculator_propose_basic_single_request(): last_sampled = torch.full((4, 1), 99, dtype=torch.int64, device="cuda") num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - drafts = spec.propose( + drafts, num_valid = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, @@ -410,7 +410,8 @@ def test_speculator_propose_basic_single_request(): assert drafts.shape == (1, 2) assert drafts.cpu().tolist() == [[3, 1]] - assert spec.get_num_valid_draft_tokens(1).cpu().tolist() == [2] + assert num_valid is not None + assert num_valid.cpu().tolist() == [2] def test_speculator_propose_zero_sampled_disables_proposal(): @@ -432,7 +433,7 @@ def test_speculator_propose_zero_sampled_disables_proposal(): # num_sampled=0 must disable speculation for this request. num_sampled = torch.tensor([0], dtype=torch.int32, device="cuda") - spec.propose( + _drafts, num_valid = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, @@ -446,7 +447,8 @@ def test_speculator_propose_zero_sampled_disables_proposal(): seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), ) - assert spec.get_num_valid_draft_tokens(1).cpu().tolist() == [0] + assert num_valid is not None + assert num_valid.cpu().tolist() == [0] def test_speculator_propose_multibatch_noncontiguous_idx_mapping(): @@ -473,7 +475,7 @@ def test_speculator_propose_multibatch_noncontiguous_idx_mapping(): last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") num_sampled = torch.tensor([1, 1], dtype=torch.int32, device="cuda") - drafts = spec.propose( + drafts, num_valid = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, @@ -488,11 +490,12 @@ def test_speculator_propose_multibatch_noncontiguous_idx_mapping(): ) assert drafts.cpu().tolist() == [[9, 7], [3, 1]] - assert spec.get_num_valid_draft_tokens(2).cpu().tolist() == [2, 2] + assert num_valid is not None + assert num_valid.cpu().tolist() == [2, 2] -def test_speculator_propose_resets_unused_slots(): - """num_valid_draft_tokens must be zeroed beyond ``num_reqs``.""" +def test_speculator_propose_returns_num_valid_matching_batch_size(): + """num_valid should be shaped [num_reqs], independent of max_num_seqs.""" spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") @@ -502,9 +505,6 @@ def test_speculator_propose_resets_unused_slots(): spec.req_states = _FakeRequestState(all_token_ids, total_len) - # Pre-populate with stale data to verify the slots are zeroed. - spec.num_valid_draft_tokens.fill_(123) - input_batch = SimpleNamespace( num_reqs=1, idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), @@ -512,7 +512,7 @@ def test_speculator_propose_resets_unused_slots(): last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - spec.propose( + drafts, num_valid = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, @@ -526,8 +526,11 @@ def test_speculator_propose_resets_unused_slots(): seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), ) - # Slot 0 should report 2 valid drafts; slots 1..3 must be zeroed. - assert spec.num_valid_draft_tokens.cpu().tolist() == [2, 0, 0, 0] + # Returned tensors match the active batch size, not max_num_seqs. + assert drafts.shape == (1, 2) + assert num_valid is not None + assert num_valid.shape == (1,) + assert num_valid.cpu().tolist() == [2] def test_speculator_propose_requires_req_states(): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1c9441900611..6ad22ae92507 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1295,7 +1295,7 @@ def sample_tokens( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] - draft_tokens = self.speculator.propose( + draft_tokens, num_valid_draft_tokens = self.speculator.propose( input_batch, attn_metadata, slot_mappings_by_layer, @@ -1310,10 +1310,6 @@ def sample_tokens( mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - num_valid_draft_tokens: torch.Tensor | None = None - get_num_valid = getattr(self.speculator, "get_num_valid_draft_tokens", None) - if get_num_valid is not None: - num_valid_draft_tokens = get_num_valid(input_batch.num_reqs) self.draft_tokens_handler.set_draft_tokens( input_batch, draft_tokens, num_valid_draft_tokens ) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index efe510f16e22..7043af7d90f4 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -479,7 +479,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() @@ -561,7 +561,7 @@ def propose( if self.num_speculative_steps == 1: # Early exit. - return self.draft_tokens[:num_reqs, :1] + return self.draft_tokens[:num_reqs, :1], None # Prepare the inputs for the decode steps. prepare_eagle_decode( @@ -593,7 +593,7 @@ def propose( num_tokens_across_dp, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs], None @triton.jit diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index f4f4602f3704..603bdee6eada 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -148,10 +148,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): .eval() ) - self.num_valid_draft_tokens: torch.Tensor = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=device - ) - self.req_states: RequestState | None = None def load_model(self, target_model: nn.Module) -> None: @@ -188,14 +184,13 @@ def propose( dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: assert self.req_states is not None, ( "NgramGPUSpeculator.req_states was not injected by the model " "runner. Ensure model_runner sets `speculator.req_states = " "self.req_states` after RequestState is constructed." ) - num_reqs = input_batch.num_reqs idx_mapping_long = input_batch.idx_mapping.long() active_tokens: torch.Tensor = self.req_states.all_token_ids.gpu[ @@ -214,14 +209,4 @@ def propose( active_last_sampled, ) - self.num_valid_draft_tokens[:num_reqs].copy_(num_valid) - if num_reqs < self.max_num_reqs: - self.num_valid_draft_tokens[num_reqs:].zero_() - - return drafts - - def get_num_valid_draft_tokens(self, num_reqs: int) -> torch.Tensor: - """ - Return the last step's per-request valid draft counts. - """ - return self.num_valid_draft_tokens[:num_reqs] + return drafts, num_valid From 3374d6fb3f4d37230a8b66fc075a22043460b68c Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Wed, 13 May 2026 10:31:36 +0800 Subject: [PATCH 10/34] add tests in yaml and change return values into tuple Signed-off-by: PatchouliTaisa --- .../gpu/spec_decode/ngram/speculator.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 603bdee6eada..3d9b674df3f7 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -46,34 +46,45 @@ def forward( valid_mask: torch.Tensor, # [B] bool (row eligible for n-gram lookup) last_sampled: torch.Tensor, # [B] int64 (fallback for -1 positions) ) -> tuple[torch.Tensor, torch.Tensor]: + """For each row, find the longest n-gram suffix match in its context + and propose the next k tokens. Fully vectorized; no data-dependent + control flow so the kernel stays torch.compile / CUDA-graph friendly. + """ B, L = token_ids.shape device = token_ids.device + # Phase 1: For each n in [min_n, max_n], find the right-most position + # of the length-n suffix inside each row (excluding the suffix itself). + # first_match_pos[b, i] = match position for n=min_n+i, or -1. first_match_pos = torch.full( (B, self.num_sizes), -1, dtype=torch.long, device=device ) batch_idx = torch.arange(B, device=device) for i, n in enumerate(range(self.min_n, self.max_n + 1)): + # Sliding length-n windows; needle = length-n suffix of each row. windows = token_ids.unfold(1, n, 1) num_windows = windows.shape[1] - suffix_start = (seq_lens.long() - n).clamp(min=0) offsets = torch.arange(n, device=device) suffix_idx = suffix_start.unsqueeze(1) + offsets suffix = torch.gather(token_ids, 1, suffix_idx) - matches = (windows == suffix.unsqueeze(1)).all(dim=-1) + # Mask out windows that overlap (or live past) the suffix. max_valid_pos = seq_lens.long() - n - 1 window_pos = torch.arange(num_windows, device=device) matches = matches & (window_pos.unsqueeze(0) <= max_valid_pos.unsqueeze(1)) + # Right-most match via argmax on (pos if match else -1); re-check + # the chosen index to distinguish a real match from the fallback. matched_indices = torch.where(matches, window_pos.unsqueeze(0), -1) idx = matched_indices.argmax(dim=1) has_match = matches[batch_idx, idx] first_match_pos[:, i] = torch.where(has_match, idx, -1) + # Phase 2: Pick the largest n that produced a match (right-most True + # along the n-axis). best_i = (first_match_pos >= 0).int().flip(dims=[1]).argmax(dim=1) best_i = self.num_sizes - 1 - best_i best_pos = first_match_pos[batch_idx, best_i] @@ -83,6 +94,8 @@ def forward( best_n = ngram_lens_table[best_i] has_any = best_pos >= 0 + # Phase 3: Gather the next k tokens after the match. No-match rows + # use draft_start=0 to keep indices in-bounds; they're masked below. draft_start = torch.where( has_any, best_pos + best_n, @@ -92,15 +105,18 @@ def forward( draft_idx = (draft_start.unsqueeze(1) + k_range).clamp_(0, L - 1) drafts = torch.gather(token_ids, 1, draft_idx).to(torch.int64) + # Phase 4: A slot j is valid iff the row has a match, valid_mask is + # True, and j < seq_len - draft_start (so we don't read past context). + # num_valid = length of the leading run of True values per row. tokens_available = (seq_lens.long() - draft_start).clamp_(min=0) valid_positions = k_range.unsqueeze(0) < tokens_available.unsqueeze(1) row_valid = has_any & valid_mask leading_valid_mask = valid_positions & row_valid.unsqueeze(1) - cum_valid = leading_valid_mask.int().cumsum(dim=1) positions = torch.arange(1, self.k + 1, device=device) num_valid = (cum_valid == positions.unsqueeze(0)).int().sum(dim=1) + # Phase 5: Replace invalid slots with last_sampled. safe_drafts = torch.where( leading_valid_mask, drafts, From 34fa2871f1cce308bcaee62c3a9c49c3e19e806c Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Wed, 13 May 2026 10:31:54 +0800 Subject: [PATCH 11/34] add tests in yaml and change return values into tuple Signed-off-by: PatchouliTaisa --- .buildkite/test_areas/model_runner_v2.yaml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 9dfd046289e8..d271b27fb2c7 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -19,7 +19,7 @@ steps: - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" # This requires eager until we sort out CG correctness issues. # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged. - - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" + - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py - pytest -v -s v1/e2e/general/test_context_length.py - pytest -v -s v1/e2e/general/test_min_tokens.py # Temporary hack filter to exclude ngram spec decoding based tests. @@ -106,13 +106,15 @@ steps: - vllm/v1/worker/gpu/ - vllm/v1/worker/gpu_worker.py - tests/v1/spec_decode/test_max_len.py + - tests/v1/spec_decode/test_ngram_gpu.py - tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py - tests/v1/e2e/spec_decode/test_spec_decode.py commands: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp" + - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp or ngram_gpu" - pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp" + - pytest -v -s v1/spec_decode/test_ngram_gpu.py + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp or ngram_gpu" From a9101c7103df39082e2071a719de2218897c8a07 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Wed, 13 May 2026 19:43:02 +0800 Subject: [PATCH 12/34] fix test error Signed-off-by: PatchouliTaisa --- vllm/config/vllm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index d220aa65035d..1629a67252a9 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1877,7 +1877,8 @@ def _validate_v2_model_runner(self) -> None: if ( self.speculative_config is not None - and self.speculative_config.method not in ("eagle", "eagle3", "mtp") + and self.speculative_config.method not in get_args(EagleModelTypes) + and self.speculative_config.method not in get_args(NgramGPUTypes) ): unsupported.append(f"speculative method '{self.speculative_config.method}'") From 5986d2b21a4cc35905ca4490deb39fa2c0584147 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 14 May 2026 11:59:20 +0800 Subject: [PATCH 13/34] inputs params fixed Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/spec_decode/ngram/speculator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 3d9b674df3f7..a80174bfdd77 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -200,6 +200,7 @@ def propose( dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: assert self.req_states is not None, ( "NgramGPUSpeculator.req_states was not injected by the model " From 07641bc8678c037a86d26d40fa400127812f361f Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 14 May 2026 15:38:42 +0800 Subject: [PATCH 14/34] fix(test): Add required InitVar fields to SchedulerConfig in ngram_gpu test helper SchedulerConfig requires max_model_len and is_encoder_decoder as InitVar fields. Use SchedulerConfig.default_factory() to properly handle these fields with appropriate defaults (max_model_len=8192, is_encoder_decoder=False). Co-authored-by: GitHub Copilot Signed-off-by: PatchouliTaisa --- tests/v1/spec_decode/test_ngram_gpu.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 273860c40657..fe9146243395 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -62,7 +62,10 @@ def _make_vllm_config( max_model_len=max_model_len, enforce_eager=True, ) - scheduler_config = SchedulerConfig(max_num_seqs=max_num_seqs) + scheduler_config = SchedulerConfig.default_factory( + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) speculative_config = SpeculativeConfig( method="ngram_gpu", prompt_lookup_min=min_n, From 29cb65e7e1866bbecc94b1747e5a0fd3fca8d9f1 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Fri, 15 May 2026 14:17:24 +0800 Subject: [PATCH 15/34] modify vllm config to pass ngram_gpu Signed-off-by: PatchouliTaisa --- vllm/config/vllm.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e5d3653656fa..330c2e095386 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1975,10 +1975,9 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported.append("sequence parallelism") if speculative_config is not None: - # TODO: ngram are not supported by the v2 model runner yet - if speculative_config.method in ("ngram",): - unsupported.append("ngram cpu speculative decoding") - elif speculative_config.method not in ("eagle", "eagle3", "mtp"): + if speculative_config.method not in get_args( + EagleModelTypes + ) and speculative_config.method not in get_args(NgramGPUTypes): unsupported.append(f"speculative method '{speculative_config.method}'") if ( From 98cc7dae57e9a197ff87ed1076809b8e30b1d7fa Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Fri, 15 May 2026 19:20:33 +0800 Subject: [PATCH 16/34] fix method signature Signed-off-by: PatchouliTaisa --- tests/v1/spec_decode/test_ngram_gpu.py | 2 +- vllm/v1/worker/gpu/spec_decode/ngram/speculator.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index fe9146243395..903ceade82b4 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -576,4 +576,4 @@ def test_speculator_construction_validates_speculative_config(): # No-op hooks must not raise. spec.load_model(target_model=None) spec.set_attn() - spec.capture_model() + spec.capture() diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 0ebd1c408339..0398460040e4 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -178,7 +178,7 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: """N-gram kernel is torch.compile-managed; no explicit CG capture.""" pass - def capture(self) -> None: + def capture(self, *args: Any, **kwargs: Any) -> None: """No graph capture phase required.""" pass From 445322dc8b7421424d6fb8f4eefd3c60b996b290 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 19 May 2026 10:35:55 +0800 Subject: [PATCH 17/34] format fixed Signed-off-by: PatchouliTaisa --- vllm/v1/core/sched/scheduler.py | 13 +++---------- vllm/v1/engine/core.py | 11 +++++------ 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 77710ac6eda4..60ea5833d1e4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1691,10 +1691,7 @@ def _free_encoder_inputs(self, request: Request) -> None: def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: num_valid_list = draft_token_ids.num_valid_draft_tokens for i, (req_id, spec_token_ids) in enumerate( - zip( - draft_token_ids.req_ids, - draft_token_ids.draft_token_ids, - ) + zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids) ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -1727,10 +1724,7 @@ def update_draft_token_ids_in_output( sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens for i, (req_id, spec_token_ids) in enumerate( - zip( - draft_token_ids.req_ids, - draft_token_ids.draft_token_ids, - ) + zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids) ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -1747,8 +1741,7 @@ def update_draft_token_ids_in_output( effective_num_spec_tokens = orig_num_spec_tokens if num_valid_list is not None: effective_num_spec_tokens = max( - 0, - min(num_valid_list[i], orig_num_spec_tokens), + 0, min(num_valid_list[i], orig_num_spec_tokens) ) del spec_token_ids[effective_num_spec_tokens:] diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 66830327d65e..48dc5682e909 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -466,9 +466,10 @@ def _maybe_update_async_draft_token_ids( if not (self.async_scheduling and self.use_spec_decode and model_executed): return None draft_token_ids = self.model_executor.take_draft_token_ids() - if draft_token_ids is None: - return None - if draft_token_ids.num_valid_draft_tokens is not None: + if ( + draft_token_ids is not None + and draft_token_ids.num_valid_draft_tokens is not None + ): self.scheduler.update_draft_token_ids(draft_token_ids) return draft_token_ids @@ -588,9 +589,7 @@ def step_with_batch_queue( # we can compute the grammar bitmask for the deferred request. if self.use_spec_decode: draft_token_ids = ( - async_draft_token_ids - if async_draft_token_ids is not None - else self.model_executor.take_draft_token_ids() + async_draft_token_ids or self.model_executor.take_draft_token_ids() ) assert draft_token_ids is not None # Update the draft token ids in the scheduler output to From 98a08091e3e89167eb6a8f162abc2f85faccb837 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Tue, 26 May 2026 15:23:39 +0800 Subject: [PATCH 18/34] fix gpu utils drop bug Signed-off-by: PatchouliTaisa --- vllm/v1/engine/core.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 5cac1419ce73..01ec7f8acdcb 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -155,6 +155,13 @@ def __init__( hash_block_size=hash_block_size, ) self.use_spec_decode = vllm_config.speculative_config is not None + # Variable-length drafters (e.g., ngram_gpu) need to retrieve + # num_valid_draft_tokens in async scheduling mode to truncate drafts. + # Fixed-length drafters (Eagle/Eagle3/MTP) don't need this extra RPC. + self.use_variable_length_drafter = ( + vllm_config.speculative_config is not None + and vllm_config.speculative_config.use_ngram_gpu() + ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore @@ -466,7 +473,11 @@ def _maybe_update_async_draft_token_ids( Consume variable-length draft metadata from the just-completed batch and apply it to scheduler request state. """ - if not (self.async_scheduling and self.use_spec_decode and model_executed): + if not ( + self.async_scheduling + and self.use_variable_length_drafter + and model_executed + ): return None draft_token_ids = self.model_executor.take_draft_token_ids() if ( From 5228bd698df5823a4c5dcb794e38321d8c622dc7 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Wed, 3 Jun 2026 10:48:29 +0800 Subject: [PATCH 19/34] duplicated codes removed Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/model_runner.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e8299e9440a2..bcb402b8f8ab 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1452,10 +1452,6 @@ def sample_tokens( kv_connector_output = self.kv_connector.post_forward(finished_req_ids) model_runner_output.kv_connector_output = kv_connector_output - # Post-step KV connector related operations. - kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - model_runner_output.kv_connector_output = kv_connector_output - if self.use_async_scheduling: return async_output return async_output.get_output() From af498eaccda575a1fa089da358876a8f6803797d Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 4 Jun 2026 20:23:20 +0800 Subject: [PATCH 20/34] triton kernel for ngram gpu Signed-off-by: PatchouliTaisa --- .../gpu/spec_decode/ngram/speculator.py | 260 +++++++++++++++++- 1 file changed, 249 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 0398460040e4..ada9a3c81fe0 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -11,6 +11,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import set_forward_context +from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.worker.gpu.input_batch import InputBatch if TYPE_CHECKING: @@ -125,6 +126,226 @@ def forward( return safe_drafts, num_valid.to(torch.int32) +if HAS_TRITON: + + @triton.jit + def _ngram_scan_kernel( + token_ids_ptr, # *int32 [B, L] + seq_lens_ptr, # *int32 [B] + valid_mask_ptr, # *int8 [B] + scratch_ptr, # *int64 [B, N_BLOCKS] (output) + L, # int64 scalar + L_PLUS_1, # int64 scalar (= L + 1, used for packing) + N_BLOCKS, # int64 scalar (stride of scratch's second dim) + MIN_N: tl.constexpr, + MAX_N: tl.constexpr, + MAX_N_PO2: tl.constexpr, + BLOCK_L: tl.constexpr, + ): + b = tl.program_id(0).to(tl.int64) + blk = tl.program_id(1).to(tl.int64) + L_ = tl.cast(L, tl.int64) + Lp1 = tl.cast(L_PLUS_1, tl.int64) + NB = tl.cast(N_BLOCKS, tl.int64) + + seq_len = tl.load(seq_lens_ptr + b).to(tl.int64) + row_valid = tl.load(valid_mask_ptr + b).to(tl.int1) + eligible_row = row_valid & (seq_len >= MIN_N) + + scratch_off = b * NB + blk + + # Ineligible rows (or blocks past the last valid pos) write 0. + if not eligible_row: + tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64)) + return + + row_off = b * L_ + + # Load the length-MAX_N suffix once into registers. + suf_iota = tl.arange(0, MAX_N_PO2).to(tl.int64) + suf_pos = seq_len - MAX_N + suf_iota + suf_in_range = (suf_iota < MAX_N) & (suf_pos >= 0) & (suf_pos < seq_len) + suffix = tl.load( + token_ids_ptr + row_off + suf_pos, + mask=suf_in_range, + other=-1, + ).to(tl.int32) + + pos_iota = tl.arange(0, BLOCK_L).to(tl.int64) + pos = blk * BLOCK_L + pos_iota # ascending + + best_score = tl.zeros([BLOCK_L], dtype=tl.int64) + + for n_iter in tl.static_range(MIN_N, MAX_N + 1): + max_pos_n = seq_len - n_iter - 1 + match = (pos >= 0) & (pos <= max_pos_n) + for j in tl.static_range(0, n_iter): + tok = tl.load( + token_ids_ptr + row_off + (pos + j), + mask=match, + other=0, + ).to(tl.int32) + suf_idx = (MAX_N - n_iter) + j + suf_val = tl.sum(tl.where(suf_iota == suf_idx, suffix, 0)) + match = match & (tok == suf_val) + + cand = n_iter * Lp1 + pos + 1 + best_score = tl.where(match, cand, best_score) + + block_best = tl.max(best_score, axis=0) + tl.store(scratch_ptr + scratch_off, block_best) + + @triton.jit + def _ngram_finalize_kernel( + token_ids_ptr, # *int32 [B, L] + seq_lens_ptr, # *int32 [B] + valid_mask_ptr, # *int8 [B] + last_sampled_ptr, # *int64 [B] + scratch_ptr, # *int64 [B, N_BLOCKS] + drafts_ptr, # *int64 [B, K] (output) + num_valid_ptr, # *int32 [B] (output) + L, + L_PLUS_1, + N_BLOCKS, + K: tl.constexpr, + K_PO2: tl.constexpr, + N_BLOCKS_PO2: tl.constexpr, + ): + b = tl.program_id(0).to(tl.int64) + L_ = tl.cast(L, tl.int64) + Lp1 = tl.cast(L_PLUS_1, tl.int64) + NB = tl.cast(N_BLOCKS, tl.int64) + + nb_iota = tl.arange(0, N_BLOCKS_PO2).to(tl.int64) + nb_in_range = nb_iota < NB + block_scores = tl.load( + scratch_ptr + b * NB + nb_iota, + mask=nb_in_range, + other=0, + ) + score = tl.max(block_scores, axis=0) + + seq_len = tl.load(seq_lens_ptr + b).to(tl.int64) + row_valid = tl.load(valid_mask_ptr + b).to(tl.int1) + last_tok = tl.load(last_sampled_ptr + b) + + has_match = score > 0 + s1 = score - 1 + best_n = tl.where(has_match, s1 // Lp1, tl.zeros_like(s1)) + best_pos = tl.where(has_match, s1 - best_n * Lp1, tl.zeros_like(s1)) + draft_start = tl.where(has_match, best_pos + best_n, tl.zeros_like(s1)) + + tokens_avail = tl.maximum(seq_len - draft_start, 0) + write_ok = row_valid & has_match + nv = tl.where(write_ok, tl.minimum(tl.cast(K, tl.int64), tokens_avail), 0) + tl.store(num_valid_ptr + b, nv.to(tl.int32)) + + row_off = b * L_ + k_iota = tl.arange(0, K_PO2).to(tl.int64) + k_in_range = k_iota < K + gather_idx = tl.minimum(draft_start + k_iota, L_ - 1) + slot_valid = (k_iota < tokens_avail) & write_ok & k_in_range + gathered = tl.load( + token_ids_ptr + row_off + gather_idx, + mask=slot_valid, + other=0, + ).to(tl.int64) + out = tl.where(slot_valid, gathered, last_tok) + tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) + + +_NGRAM_SCRATCH: dict[tuple, torch.Tensor] = {} + + +def _get_ngram_scratch(B: int, n_blocks: int, device: torch.device) -> torch.Tensor: + key = (device, B, n_blocks) + buf = _NGRAM_SCRATCH.get(key) + if buf is None: + buf = torch.empty((B, n_blocks), dtype=torch.int64, device=device) + _NGRAM_SCRATCH[key] = buf + return buf + + +def _ngram_propose_triton( + token_ids: torch.Tensor, + seq_lens: torch.Tensor, + valid_mask: torch.Tensor, + last_sampled: torch.Tensor, + min_n: int, + max_n: int, + k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + For each row, find the longest n-gram suffix match in its context + and propose the next k tokens. + """ + B, L = token_ids.shape + device = token_ids.device + + drafts = torch.empty((B, k), dtype=torch.int64, device=device) + num_valid = torch.empty((B,), dtype=torch.int32, device=device) + + if B == 0: + return drafts, num_valid + + tok = token_ids.contiguous().to(torch.int32) + seq = seq_lens.contiguous().to(torch.int32) + vmask = valid_mask.contiguous().to(torch.bool).to(torch.int8) + last = last_sampled.contiguous().to(torch.int64).view(-1) + + if L >= 1024: + BLOCK_L = 256 + elif L >= 256: + BLOCK_L = 128 + elif L >= 64: + BLOCK_L = 64 + else: + BLOCK_L = max(16, triton.next_power_of_2(max(L, 1))) + + K_PO2 = max(1, triton.next_power_of_2(k)) + MAX_N_PO2 = max(1, triton.next_power_of_2(max_n)) + n_blocks = (L + BLOCK_L - 1) // BLOCK_L + n_blocks_po2 = max(1, triton.next_power_of_2(n_blocks)) + + scratch = _get_ngram_scratch(B, n_blocks, device) + + L_plus_1 = L + 1 + _ngram_scan_kernel[(B, n_blocks)]( + tok, + seq, + vmask, + scratch, + L, + L_plus_1, + n_blocks, + min_n, + max_n, + MAX_N_PO2, + BLOCK_L, + num_warps=4, + num_stages=2, + ) + + _ngram_finalize_kernel[(B,)]( + tok, + seq, + vmask, + last, + scratch, + drafts, + num_valid, + L, + L_plus_1, + n_blocks, + k, + K_PO2, + n_blocks_po2, + num_warps=2, + num_stages=1, + ) + return drafts, num_valid + + class NgramGPUSpeculator: """ V2-compatible GPU n-gram speculator. @@ -154,15 +375,21 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs self.max_model_len: int = vllm_config.model_config.max_model_len - self.kernel = ( - _NgramKernel( - min_n=self.min_n, - max_n=self.max_n, - k=self.num_speculative_steps, + # Triton is the default fast path; the torch.compile kernel is + # only constructed (and used) when Triton is unavailable. + self.use_triton: bool = HAS_TRITON + if self.use_triton: + self.kernel: _NgramKernel + else: + self.kernel = ( + _NgramKernel( + min_n=self.min_n, + max_n=self.max_n, + k=self.num_speculative_steps, + ) + .to(device) + .eval() ) - .to(device) - .eval() - ) self.req_states: RequestState | None = None @@ -175,7 +402,7 @@ def set_attn(self, *args: Any, **kwargs: Any) -> None: pass def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - """N-gram kernel is torch.compile-managed; no explicit CG capture.""" + """N-gram kernels are launched directly; no explicit CG capture.""" pass def capture(self, *args: Any, **kwargs: Any) -> None: @@ -218,12 +445,23 @@ def propose( valid_mask = (num_sampled > 0) & (active_seq_lens >= self.min_n) - with set_forward_context(None, self.vllm_config): - drafts, num_valid = self.kernel( + if self.use_triton: + drafts, num_valid = _ngram_propose_triton( active_tokens, active_seq_lens, valid_mask, active_last_sampled, + self.min_n, + self.max_n, + self.num_speculative_steps, ) + else: + with set_forward_context(None, self.vllm_config): + drafts, num_valid = self.kernel( + active_tokens, + active_seq_lens, + valid_mask, + active_last_sampled, + ) return drafts, num_valid From 6dd58ad2a01c1e302eae28fc5f365f57ae3cbf84 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 11 Jun 2026 10:11:06 +0800 Subject: [PATCH 21/34] return values fixed Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 868540437b23..862c5cf9a215 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -161,7 +161,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() @@ -244,7 +244,7 @@ def propose( if self.num_speculative_steps == 1: # Early exit. - return self.draft_tokens[:num_reqs, :1] + return self.draft_tokens[:num_reqs, :1], None # Prepare the inputs for the decode steps. prepare_decode_inputs( @@ -277,7 +277,7 @@ def propose( num_tokens_across_dp, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs], None def sample_draft( self, From 3e20869fb6f0690c1aca89ccb9544ad5992166e6 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Thu, 11 Jun 2026 19:59:21 +0800 Subject: [PATCH 22/34] config bug fixed Signed-off-by: PatchouliTaisa --- vllm/config/vllm.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a91b81c1074c..0fc91c9f1808 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2005,10 +2005,16 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported.append("sequence parallelism") if speculative_config is not None: - # TODO: ngram are not supported by the v2 model runner yet + # CPU ngram is not supported by the V2 model runner yet. if speculative_config.method == "ngram": unsupported.append("ngram speculative decoding") - elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"): + elif speculative_config.method not in ( + "eagle", + "eagle3", + "mtp", + "dflash", + "ngram_gpu", + ): unsupported.append(f"speculative method '{speculative_config.method}'") # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle) From 4033877c16b2692f876fea86855ad78f4db7e3c9 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Fri, 12 Jun 2026 09:59:38 +0800 Subject: [PATCH 23/34] return type fixed Signed-off-by: PatchouliTaisa --- vllm/v1/worker/gpu/spec_decode/dflash/speculator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 1bd130838a1a..6237a009f62f 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -286,7 +286,7 @@ def propose( num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=CUDAGraphMode.NONE, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs], None # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. @@ -367,7 +367,7 @@ def propose( cudagraph_runtime_mode=batch_desc.cg_mode, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs], None @triton.jit From 66e1adec6c6451e22d49888cb40749e9245aabb2 Mon Sep 17 00:00:00 2001 From: PatchouliTaisa Date: Sat, 1 Aug 2026 15:14:10 +0800 Subject: [PATCH 24/34] test vllm_runner updated Signed-off-by: PatchouliTaisa --- tests/v1/spec_decode/test_max_len.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index a669e37dc324..91d53b5b57e3 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -37,16 +37,18 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) -def test_ngram_gpu_max_len(num_speculative_tokens: int): +def test_ngram_gpu_max_len(num_speculative_tokens: int, vllm_runner): """V2 GPU n-gram counterpart of ``test_ngram_max_len``. Verifies that the V2 model runner with ``method="ngram_gpu"`` correctly handles the ``max_model_len`` boundary across various speculative-token counts. """ - llm = LLM( - model="facebook/opt-125m", + with vllm_runner( + "facebook/opt-125m", + trust_remote_code=False, max_model_len=100, + enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ "method": "ngram_gpu", @@ -54,9 +56,9 @@ def test_ngram_gpu_max_len(num_speculative_tokens: int): "prompt_lookup_min": 3, "num_speculative_tokens": num_speculative_tokens, }, - ) - sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) - llm.generate(_PROMPTS, sampling_params) + ) as runner: + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + runner.llm.generate(_PROMPTS, sampling_params) @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) From 9ff65bcfcf5d20dd99bdb7dc0a2474a206c8fa10 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 14 Aug 2026 17:47:00 -0700 Subject: [PATCH 25/34] [Spec Decode] ngram_gpu on V2: GPU-side draft trimming via shared verification layout Replace the scheduler-notification design (per-step take_draft_token_ids RPC + D2H event sync) with GPU-side verification trimming reusing the DSpark adaptive-verification machinery: - The drafter records per-request valid draft counts in a persistent GPU tensor; the scheduler always schedules the full num_speculative_tokens. - At the next step, VariableDraftTrimmer clamps scheduled draft slots to min(num_valid, scheduled) and rebuilds cu_num_logits/query_start_loc on device via build_verification_layout (factored out of AdaptiveVerificationManager.reallocate_drafts). CPU totals remain upper bounds; the trimmed gap behaves as cudagraph padding downstream, and trimmed slots reconcile through the existing num_rejected accounting. - Falls back to pad-and-verify (still correct) when the attention backend or config cannot support device-side varlen trimming. - Ngram kernels now index all_token_ids rows in place via idx_mapping (no per-step [B, max_model_len] materialization), use persistent scratch, early-exit scan blocks past seq_len, and drop the torch.compile fallback. No CPU<->GPU syncs are introduced on any path. Co-authored-by: Claude Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_ngram_gpu.py | 702 +++++++----------- vllm/v1/core/sched/scheduler.py | 26 +- vllm/v1/engine/core.py | 46 +- vllm/v1/outputs.py | 2 - vllm/v1/worker/gpu/input_batch.py | 23 +- vllm/v1/worker/gpu/model_runner.py | 69 +- .../gpu/spec_decode/adaptive_verification.py | 178 ++++- .../spec_decode/autoregressive/speculator.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 4 +- .../gpu/spec_decode/ngram/speculator.py | 598 ++++++--------- vllm/v1/worker/gpu/spec_decode/utils.py | 53 +- 11 files changed, 761 insertions(+), 946 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 903ceade82b4..470caf9618f9 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -2,19 +2,23 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the V2 GPU-accelerated n-gram speculator. -These tests target the kernel and proposer logic in +These tests target the Triton proposer in ``vllm.v1.worker.gpu.spec_decode.ngram.speculator`` and complement the CPU ``NgramProposer`` tests in ``test_ngram.py``. The GPU speculator follows a -slightly different policy than the CPU one: when multiple n-gram matches of the -same length exist, the GPU kernel picks the right-most (most recent) match -inside the active context, whereas the CPU implementation returns the +slightly different policy than the CPU one: when multiple n-gram matches of +the same length exist, the GPU kernel picks the right-most (most recent) +match inside the active context, whereas the CPU implementation returns the left-most. The expectations below reflect the GPU behavior. + +Also covers the GPU draft-trimming layout helpers in +``adaptive_verification`` that ngram_gpu shares with DSpark. """ from __future__ import annotations from types import SimpleNamespace +import numpy as np import pytest import torch @@ -23,26 +27,20 @@ SchedulerConfig, SpeculativeConfig, VllmConfig, - set_current_vllm_config, ) -from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( - NgramGPUSpeculator, - _NgramKernel, +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( + VariableDraftTrimmer, + build_verification_layout, ) +from vllm.v1.worker.gpu.spec_decode.ngram.speculator import NgramGPUSpeculator -# The kernel uses ``@support_torch_compile`` and ``set_forward_context``, both -# of which expect a real CUDA device. Skip these tests gracefully on -# CPU-only or non-CUDA platforms. if not torch.cuda.is_available(): pytest.skip( "CUDA required for NgramGPUSpeculator tests", allow_module_level=True, ) - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- +DEVICE = torch.device("cuda") def _make_vllm_config( @@ -52,11 +50,6 @@ def _make_vllm_config( max_num_seqs: int = 8, max_model_len: int = 64, ) -> VllmConfig: - """Build a minimal VllmConfig configured for ngram_gpu speculative decoding. - - ``enforce_eager=True`` is used so the kernel runs in eager mode (no - torch.compile) — sufficient for unit-testing the math. - """ model_config = ModelConfig( model="facebook/opt-125m", max_model_len=max_model_len, @@ -79,278 +72,15 @@ def _make_vllm_config( ) -def _make_kernel(min_n: int, max_n: int, k: int) -> _NgramKernel: - """Construct a ``_NgramKernel`` bound to a minimal ``VllmConfig``. - - The kernel is decorated by ``@support_torch_compile`` whose generated - ``__init__`` accepts ``vllm_config`` as a kwarg. - """ - vllm_config = _make_vllm_config(min_n=min_n, max_n=max_n, k=k) - with set_current_vllm_config(vllm_config): - kernel = _NgramKernel( - min_n=min_n, - max_n=max_n, - k=k, - vllm_config=vllm_config, - ) - return kernel.to("cuda").eval() - - -def _pad_tokens(rows: list[list[int]], pad_to: int) -> torch.Tensor: - """Right-pad ragged ``rows`` to ``pad_to`` length with zeros (int32).""" - out = torch.zeros((len(rows), pad_to), dtype=torch.int32) - for i, row in enumerate(rows): - if row: - out[i, : len(row)] = torch.tensor(row, dtype=torch.int32) - return out - - -def _run_kernel( - kernel: _NgramKernel, - rows: list[list[int]], - seq_lens: list[int], - valid_mask: list[bool] | None = None, - last_sampled: list[int] | None = None, - pad_to: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Build inputs from python lists, run the kernel, return (drafts, num_valid).""" - B = len(rows) - if pad_to is None: - pad_to = max((len(r) for r in rows), default=1) - # The kernel needs at least 1 column to satisfy unfold(1, n, 1). - pad_to = max(pad_to, kernel.max_n) - - if valid_mask is None: - valid_mask = [True] * B - if last_sampled is None: - last_sampled = [0] * B - - token_ids = _pad_tokens(rows, pad_to).cuda() - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") - valid_mask_t = torch.tensor(valid_mask, dtype=torch.bool, device="cuda") - last_sampled_t = torch.tensor(last_sampled, dtype=torch.int64, device="cuda") - with torch.inference_mode(): - drafts, num_valid = kernel( - token_ids, - seq_lens_t, - valid_mask_t, - last_sampled_t, - ) - return drafts.cpu(), num_valid.cpu() - - -# --------------------------------------------------------------------------- -# _NgramKernel tests (mirror & extend tests/v1/spec_decode/test_ngram.py) -# --------------------------------------------------------------------------- - - -def test_kernel_no_match_returns_zero_valid(): - """No 2-gram match in [1,2,3,4,5] → kernel reports 0 valid drafts.""" - kernel = _make_kernel(min_n=2, max_n=2, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 4, 5]], - seq_lens=[5], - last_sampled=[42], - ) - assert num_valid.tolist() == [0] - # Padding positions must fall back to ``last_sampled`` for invalid rows. - assert drafts.tolist() == [[42, 42]] - - -def test_kernel_no_4gram_match_only(): - """No 4-gram match in [1,2,3,4,1,2,3] → 0 valid drafts.""" - kernel = _make_kernel(min_n=4, max_n=4, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 4, 1, 2, 3]], - seq_lens=[7], - last_sampled=[7], - ) - assert num_valid.tolist() == [0] - assert drafts.tolist() == [[7, 7]] - - -def test_kernel_falls_back_to_3gram_when_4gram_missing(): - """No 4-gram match but a 3-gram match exists → propose [4, 1].""" - kernel = _make_kernel(min_n=3, max_n=4, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 4, 1, 2, 3]], - seq_lens=[7], - ) - assert num_valid.tolist() == [2] - assert drafts.tolist() == [[4, 1]] - - -def test_kernel_prefers_longer_ngram(): - """Both a 4-gram (1,2,3,4) and a 3-gram match are present. - - The kernel must prefer the longer match, returning [1, 2] instead of [5, 1]. - """ - kernel = _make_kernel(min_n=3, max_n=4, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]], - seq_lens=[12], - ) - assert num_valid.tolist() == [2] - assert drafts.tolist() == [[1, 2]] - - -def test_kernel_picks_longest_match_among_2_3_4_grams(): - """2-gram and 3-gram match, 4-gram does not → propose 3-gram match [1, 2].""" - kernel = _make_kernel(min_n=2, max_n=4, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]], - seq_lens=[10], - ) - assert num_valid.tolist() == [2] - assert drafts.tolist() == [[1, 2]] - - -def test_kernel_picks_rightmost_when_multiple_matches(): - """Multiple 3-gram matches exist for suffix (1,2,3). - - The GPU kernel picks the RIGHT-most (most recent) match, unlike the CPU - proposer which picks the left-most. Tokens after the right-most match - starting at index 8 are [300, 1]. - """ - kernel = _make_kernel(min_n=3, max_n=3, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]], - seq_lens=[15], - ) - assert num_valid.tolist() == [2] - assert drafts.tolist() == [[300, 1]] - - -def test_kernel_short_context_yields_zero_valid(): - """Context length < min_n → kernel returns 0 valid drafts. - - Note: callers (NgramGPUSpeculator.propose) gate this via ``valid_mask``. - Here we exercise the kernel directly with a row whose seq_len < min_n. - The only window of length 2 cannot match a length-2 suffix because it - overlaps the suffix itself, so no valid match is produced. - """ - kernel = _make_kernel(min_n=2, max_n=2, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[5, 6]], - seq_lens=[2], - last_sampled=[99], - pad_to=4, - ) - assert num_valid.tolist() == [0] - assert drafts.tolist() == [[99, 99]] - - -def test_kernel_valid_mask_disables_row(): - """``valid_mask=False`` disables proposals for that row regardless of match.""" - kernel = _make_kernel(min_n=2, max_n=2, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 1, 2]], # has a 2-gram match for suffix (1,2) → [3, 1] - seq_lens=[5], - valid_mask=[False], - last_sampled=[77], - ) - assert num_valid.tolist() == [0] - # With the row disabled, fallback ``last_sampled`` fills the draft slots. - assert drafts.tolist() == [[77, 77]] - - -def test_kernel_truncates_num_valid_when_few_tokens_after_match(): - """When fewer than k tokens are available after the match, num_valid < k. - - Tokens: [1, 2, 1, 2] (seq_len=4). Suffix (1, 2) matches at position 0 - (the right-most match at position 2 is the suffix itself and is excluded - by ``max_valid_pos = seq_len - n - 1``). With k=3, only the first 2 slots - fit within the active context; the third slot must fall back to - ``last_sampled``. - """ - kernel = _make_kernel(min_n=2, max_n=2, k=3) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 1, 2]], - seq_lens=[4], - last_sampled=[55], - pad_to=4, - ) - # Only the first 2 speculative slots map to tokens inside the context. - assert num_valid.tolist() == [2] - drafts_row = drafts.tolist()[0] - # The first 2 drafts come from the tokens following the match. - assert drafts_row[:2] == [1, 2] - # The third slot must fall back to ``last_sampled`` since it is invalid. - assert drafts_row[2] == 55 - - -def test_kernel_multibatch_mixed(): - """Mixed batch: row 0 matches, row 1 has no match.""" - kernel = _make_kernel(min_n=2, max_n=2, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 1, 2], [4, 5, 6, 0, 0]], - seq_lens=[5, 3], - last_sampled=[10, 20], - ) - assert num_valid.tolist() == [2, 0] - assert drafts.tolist()[0] == [3, 1] - # Row 1: no match → all slots filled with last_sampled[1] = 20. - assert drafts.tolist()[1] == [20, 20] - - -def test_kernel_multibatch_independent_choice_of_n(): - """Row 0 matches as 3-gram, row 1 only matches as 2-gram. - - Verifies that each row independently picks its longest matched n. - """ - kernel = _make_kernel(min_n=2, max_n=3, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[ - [9, 1, 2, 3, 8, 1, 2, 3], # 3-gram (1,2,3) matches at idx 1 → [8, 1] - [7, 1, 2, 9, 1, 2, 0, 0], # 2-gram (1,2) matches at idx 1 → [9, 1] - ], - seq_lens=[8, 6], - ) - assert num_valid.tolist() == [2, 2] - assert drafts.tolist()[0] == [8, 1] - assert drafts.tolist()[1] == [9, 1] - - -def test_kernel_min_n_eq_1(): - """min_n=max_n=1 — single-token n-grams always match if context > 1.""" - kernel = _make_kernel(min_n=1, max_n=1, k=2) - drafts, num_valid = _run_kernel( - kernel, - rows=[[1, 2, 3, 4, 1]], # suffix (1,) matches at idx 0 → tokens after = [2, 3] - seq_lens=[5], - ) - assert num_valid.tolist() == [2] - assert drafts.tolist() == [[2, 3]] - - -# --------------------------------------------------------------------------- -# NgramGPUSpeculator.propose tests -# --------------------------------------------------------------------------- - - class _FakeStaged: - """Lightweight stand-in for ``StagedWriteTensor`` used by ``RequestState``. - - Only ``.gpu`` is consulted by ``NgramGPUSpeculator.propose``. - """ + """Stand-in for ``StagedWriteTensor``; only ``.gpu`` is read by propose().""" def __init__(self, tensor: torch.Tensor): self.gpu = tensor class _FakeRequestState: - """Minimal duck-typed ``RequestState`` exposing only the fields propose() reads.""" + """Minimal duck-typed ``RequestState`` exposing what propose() reads.""" def __init__(self, all_token_ids: torch.Tensor, total_len: torch.Tensor): self.all_token_ids = _FakeStaged(all_token_ids) @@ -361,7 +91,7 @@ def _make_speculator( min_n: int, max_n: int, k: int, - max_num_seqs: int = 4, + max_num_seqs: int = 8, max_model_len: int = 32, ) -> NgramGPUSpeculator: cfg = _make_vllm_config( @@ -371,205 +101,267 @@ def _make_speculator( max_num_seqs=max_num_seqs, max_model_len=max_model_len, ) - with set_current_vllm_config(cfg): - spec = NgramGPUSpeculator(vllm_config=cfg, device=torch.device("cuda")) - return spec + return NgramGPUSpeculator(vllm_config=cfg, device=DEVICE) -def test_speculator_propose_basic_single_request(): - """propose() should return drafts plus per-request valid counts.""" - spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) +def _propose( + spec: NgramGPUSpeculator, + rows: list[list[int]], + seq_lens: list[int] | None = None, + num_sampled: list[int] | None = None, + last_sampled: list[int] | None = None, + slots: list[int] | None = None, +) -> tuple[list[list[int]], list[int]]: + """Place each batch row at a request slot and run propose(). - # State for a single active request with tokens [1, 2, 3, 1, 2]; suffix (1,2) - # matches at index 0 → expected draft tokens [3, 1]. - all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") - all_token_ids[2, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) - total_len = torch.zeros(4, dtype=torch.int32, device="cuda") - total_len[2] = 5 + Returns (drafts, num_valid) as python lists in batch order. + """ + B = len(rows) + if seq_lens is None: + seq_lens = [len(r) for r in rows] + if num_sampled is None: + num_sampled = [1] * B + if last_sampled is None: + last_sampled = [0] * B + if slots is None: + slots = list(range(B)) + + max_num_reqs = spec.max_num_reqs + L = spec.max_model_len + all_token_ids = torch.zeros((max_num_reqs, L), dtype=torch.int32, device=DEVICE) + total_len = torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + last_sampled_t = torch.zeros((max_num_reqs, 1), dtype=torch.int64, device=DEVICE) + for row, slot, seq_len, last in zip(rows, slots, seq_lens, last_sampled): + if row: + all_token_ids[slot, : len(row)] = torch.tensor( + row, dtype=torch.int32, device=DEVICE + ) + total_len[slot] = seq_len + last_sampled_t[slot, 0] = last spec.req_states = _FakeRequestState(all_token_ids, total_len) + idx_mapping = torch.tensor(slots, dtype=torch.int64, device=DEVICE) + input_batch = SimpleNamespace(num_reqs=B, idx_mapping=idx_mapping) - # The active batch contains a single request mapping to slot index 2. - input_batch = SimpleNamespace( - num_reqs=1, - idx_mapping=torch.tensor([2], dtype=torch.int32, device="cuda"), - ) - last_sampled = torch.full((4, 1), 99, dtype=torch.int64, device="cuda") - num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - - drafts, num_valid = spec.propose( + drafts = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, - last_hidden_states=torch.empty(0, device="cuda"), + last_hidden_states=torch.empty(0, device=DEVICE), aux_hidden_states=None, - num_sampled=num_sampled, - num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), - last_sampled=last_sampled, - next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), - temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), - seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + num_sampled=torch.tensor(num_sampled, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(B, dtype=torch.int32, device=DEVICE), + last_sampled=last_sampled_t, + next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), ) + num_valid = spec.num_valid_drafts[idx_mapping] + return drafts.cpu().tolist(), num_valid.cpu().tolist() - assert drafts.shape == (1, 2) - assert drafts.cpu().tolist() == [[3, 1]] - assert num_valid is not None - assert num_valid.cpu().tolist() == [2] +# --------------------------------------------------------------------------- +# Proposal behavior +# --------------------------------------------------------------------------- -def test_speculator_propose_zero_sampled_disables_proposal(): - """When num_sampled==0 for a request, valid_mask is False → 0 valid drafts.""" - spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) - all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") - all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) - total_len = torch.zeros(4, dtype=torch.int32, device="cuda") - total_len[0] = 5 +def test_no_match_returns_zero_valid(): + """No 2-gram match in [1,2,3,4,5] → 0 valid drafts, last_sampled fill.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 5]], last_sampled=[42]) + assert num_valid == [0] + assert drafts == [[42, 42]] - spec.req_states = _FakeRequestState(all_token_ids, total_len) - input_batch = SimpleNamespace( - num_reqs=1, - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - ) - last_sampled = torch.full((4, 1), 7, dtype=torch.int64, device="cuda") - # num_sampled=0 must disable speculation for this request. - num_sampled = torch.tensor([0], dtype=torch.int32, device="cuda") +def test_no_4gram_match_only(): + """No 4-gram match in [1,2,3,4,1,2,3] → 0 valid drafts.""" + spec = _make_speculator(min_n=4, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]], last_sampled=[7]) + assert num_valid == [0] + assert drafts == [[7, 7]] - _drafts, num_valid = spec.propose( - input_batch=input_batch, - attn_metadata=None, - slot_mappings=None, - last_hidden_states=torch.empty(0, device="cuda"), - aux_hidden_states=None, - num_sampled=num_sampled, - num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), - last_sampled=last_sampled, - next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), - temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), - seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), - ) - assert num_valid is not None - assert num_valid.cpu().tolist() == [0] +def test_falls_back_to_3gram_when_4gram_missing(): + """No 4-gram match but a 3-gram match exists → propose [4, 1].""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]]) + assert num_valid == [2] + assert drafts == [[4, 1]] -def test_speculator_propose_multibatch_noncontiguous_idx_mapping(): - """Verify that propose() correctly reads via idx_mapping (non-contiguous).""" - spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) +def test_prefers_longer_ngram(): + """Both a 4-gram and a 3-gram match exist → prefer the 4-gram match.""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]]) + assert num_valid == [2] + assert drafts == [[1, 2]] - all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") - # Request A at slot 0: tokens [1,2,3,1,2] → drafts [3,1] - all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) - # Request B at slot 3: tokens [7,8,9,7,8] → drafts [9,7] - all_token_ids[3, :5] = torch.tensor([7, 8, 9, 7, 8], dtype=torch.int32) - total_len = torch.zeros(4, dtype=torch.int32, device="cuda") - total_len[0] = 5 - total_len[3] = 5 +def test_picks_longest_match_among_2_3_4_grams(): + """2-gram and 3-gram match, 4-gram does not → propose 3-gram match [1, 2].""" + spec = _make_speculator(min_n=2, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]]) + assert num_valid == [2] + assert drafts == [[1, 2]] - spec.req_states = _FakeRequestState(all_token_ids, total_len) - # Batch order is [slot 3, slot 0] — verify the idx_mapping is honored. - input_batch = SimpleNamespace( - num_reqs=2, - idx_mapping=torch.tensor([3, 0], dtype=torch.int32, device="cuda"), +def test_picks_rightmost_when_multiple_matches(): + """Multiple 3-gram matches for suffix (1,2,3) → pick the right-most.""" + spec = _make_speculator(min_n=3, max_n=3, k=2) + drafts, num_valid = _propose( + spec, [[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]] ) - last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") - num_sampled = torch.tensor([1, 1], dtype=torch.int32, device="cuda") + assert num_valid == [2] + assert drafts == [[300, 1]] - drafts, num_valid = spec.propose( - input_batch=input_batch, - attn_metadata=None, - slot_mappings=None, - last_hidden_states=torch.empty(0, device="cuda"), - aux_hidden_states=None, - num_sampled=num_sampled, - num_rejected=torch.zeros(2, dtype=torch.int32, device="cuda"), - last_sampled=last_sampled, - next_prefill_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - temperature=torch.zeros(2, dtype=torch.float32, device="cuda"), - seeds=torch.zeros(2, dtype=torch.int64, device="cuda"), + +def test_short_context_yields_zero_valid(): + """The only length-2 window overlaps the suffix itself → no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose(spec, [[5, 6]], last_sampled=[99]) + assert num_valid == [0] + assert drafts == [[99, 99]] + + +def test_zero_sampled_disables_proposal(): + """num_sampled==0 disables proposals for that request regardless of match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, [[1, 2, 3, 1, 2]], num_sampled=[0], last_sampled=[77] ) + assert num_valid == [0] + assert drafts == [[77, 77]] - assert drafts.cpu().tolist() == [[9, 7], [3, 1]] - assert num_valid is not None - assert num_valid.cpu().tolist() == [2, 2] +def test_truncates_num_valid_when_few_tokens_after_match(): + """Fewer than k tokens after the match → num_valid < k, tail falls back. -def test_speculator_propose_returns_num_valid_matching_batch_size(): - """num_valid should be shaped [num_reqs], independent of max_num_seqs.""" - spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=4, max_model_len=16) + Tokens: [1, 2, 1, 2] (seq_len=4). Suffix (1, 2) matches at position 0 + (the match at position 2 is the suffix itself and is excluded). With + k=3, only 2 slots map to tokens inside the context. + """ + spec = _make_speculator(min_n=2, max_n=2, k=3) + drafts, num_valid = _propose(spec, [[1, 2, 1, 2]], last_sampled=[55]) + assert num_valid == [2] + assert drafts[0][:2] == [1, 2] + assert drafts[0][2] == 55 - all_token_ids = torch.zeros((4, 16), dtype=torch.int32, device="cuda") - all_token_ids[0, :5] = torch.tensor([1, 2, 3, 1, 2], dtype=torch.int32) - total_len = torch.zeros(4, dtype=torch.int32, device="cuda") - total_len[0] = 5 - spec.req_states = _FakeRequestState(all_token_ids, total_len) +def test_multibatch_mixed(): + """Mixed batch: row 0 matches, row 1 has no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, + [[1, 2, 3, 1, 2], [4, 5, 6]], + last_sampled=[10, 20], + ) + assert num_valid == [2, 0] + assert drafts[0] == [3, 1] + assert drafts[1] == [20, 20] + + +def test_multibatch_independent_choice_of_n(): + """Each row independently picks its longest matched n.""" + spec = _make_speculator(min_n=2, max_n=3, k=2) + drafts, num_valid = _propose( + spec, + [ + [9, 1, 2, 3, 8, 1, 2, 3], # 3-gram (1,2,3) at idx 1 → [8, 1] + [7, 1, 2, 9, 1, 2], # 2-gram (1,2) at idx 1 → [9, 1] + ], + ) + assert num_valid == [2, 2] + assert drafts[0] == [8, 1] + assert drafts[1] == [9, 1] + + +def test_min_n_eq_1(): + """min_n=max_n=1 — single-token n-grams always match if context > 1.""" + spec = _make_speculator(min_n=1, max_n=1, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1]]) + assert num_valid == [2] + assert drafts == [[2, 3]] + + +def test_noncontiguous_idx_mapping(): + """propose() reads token rows in place via idx_mapping (non-contiguous).""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, + [[7, 8, 9, 7, 8], [1, 2, 3, 1, 2]], + slots=[3, 0], + ) + assert drafts == [[9, 7], [3, 1]] + assert num_valid == [2, 2] + + +def test_num_valid_written_to_request_slots(): + """num_valid_drafts is req-slot indexed for the GPU draft trimmer.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + _propose( + spec, + [[7, 8, 9, 7, 8], [1, 2, 3, 4, 5]], + slots=[5, 2], + ) + nv = spec.num_valid_drafts.cpu() + assert nv[5].item() == 2 # match + assert nv[2].item() == 0 # no match + + +def test_dummy_run_does_not_touch_state(): + """Dummy runs must not mutate persistent request or drafter state.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + _propose(spec, [[1, 2, 3, 1, 2]], slots=[1]) + before = spec.num_valid_drafts.clone() input_batch = SimpleNamespace( num_reqs=1, - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + idx_mapping=torch.tensor([1], dtype=torch.int64, device=DEVICE), ) - last_sampled = torch.zeros((4, 1), dtype=torch.int64, device="cuda") - num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - - drafts, num_valid = spec.propose( + drafts = spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, - last_hidden_states=torch.empty(0, device="cuda"), + last_hidden_states=torch.empty(0, device=DEVICE), aux_hidden_states=None, - num_sampled=num_sampled, - num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), - last_sampled=last_sampled, - next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), - temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), - seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + num_sampled=torch.ones(1, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(1, dtype=torch.int32, device=DEVICE), + last_sampled=torch.zeros((8, 1), dtype=torch.int64, device=DEVICE), + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(1, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(1, dtype=torch.int64, device=DEVICE), + dummy_run=True, ) - - # Returned tensors match the active batch size, not max_num_seqs. assert drafts.shape == (1, 2) - assert num_valid is not None - assert num_valid.shape == (1,) - assert num_valid.cpu().tolist() == [2] + assert torch.equal(spec.num_valid_drafts.cpu(), before.cpu()) -def test_speculator_propose_requires_req_states(): - """propose() must assert that req_states has been injected by the model runner.""" - spec = _make_speculator(min_n=2, max_n=2, k=2, max_num_seqs=2, max_model_len=8) - # req_states is None by default and should trigger an AssertionError. +def test_propose_requires_req_states(): + """propose() must assert that req_states was injected by the model runner.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) assert spec.req_states is None - input_batch = SimpleNamespace( num_reqs=1, - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + idx_mapping=torch.tensor([0], dtype=torch.int64, device=DEVICE), ) - last_sampled = torch.zeros((2, 1), dtype=torch.int64, device="cuda") - num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda") - with pytest.raises(AssertionError, match="req_states"): spec.propose( input_batch=input_batch, attn_metadata=None, slot_mappings=None, - last_hidden_states=torch.empty(0, device="cuda"), + last_hidden_states=torch.empty(0, device=DEVICE), aux_hidden_states=None, - num_sampled=num_sampled, - num_rejected=torch.zeros(1, dtype=torch.int32, device="cuda"), - last_sampled=last_sampled, - next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device="cuda"), - temperature=torch.zeros(1, dtype=torch.float32, device="cuda"), - seeds=torch.zeros(1, dtype=torch.int64, device="cuda"), + num_sampled=torch.ones(1, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(1, dtype=torch.int32, device=DEVICE), + last_sampled=torch.zeros((8, 1), dtype=torch.int64, device=DEVICE), + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(1, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(1, dtype=torch.int64, device=DEVICE), ) -def test_speculator_construction_validates_speculative_config(): - """NgramGPUSpeculator requires prompt_lookup_min and prompt_lookup_max to be set.""" - cfg = _make_vllm_config(min_n=2, max_n=3, k=2) - with set_current_vllm_config(cfg): - spec = NgramGPUSpeculator(vllm_config=cfg, device=torch.device("cuda")) +def test_construction_validates_speculative_config(): + spec = _make_speculator(min_n=2, max_n=3, k=2) assert spec.min_n == 2 assert spec.max_n == 3 assert spec.num_speculative_steps == 2 @@ -577,3 +369,59 @@ def test_speculator_construction_validates_speculative_config(): spec.load_model(target_model=None) spec.set_attn() spec.capture() + + +# --------------------------------------------------------------------------- +# GPU draft trimming (shared verification-layout machinery) +# --------------------------------------------------------------------------- + + +def test_build_verification_layout_exact_and_gpu_tail(): + """Layout cumsums match a numpy reference; padding tail equals the total.""" + capacities = torch.tensor([2, 0, 1], dtype=torch.int32, device=DEVICE) + non_draft = torch.tensor([1, 5, 1], dtype=torch.int32, device=DEVICE) + num_bonus = 1 + max_num_reqs = 6 + cu_num_logits = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + qsl = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + + for num_tokens in (10, None): # exact CPU total vs GPU cumsum tail + cnl, out_qsl = build_verification_layout( + capacities, non_draft, num_bonus, cu_num_logits, qsl, num_tokens + ) + assert cnl.cpu().tolist() == [0, 3, 4, 6] + assert out_qsl.cpu().tolist()[:4] == [0, 3, 8, 10] + # Trailing (padding) entries hold the batch total. + assert out_qsl.cpu().tolist()[4:] == [10, 10, 10] + + +def test_variable_draft_trimmer_clamps_to_num_valid(): + """Scheduled draft slots are clamped per request to the drafter's counts.""" + max_num_reqs = 8 + num_valid_drafts = torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + num_valid_drafts[4] = 1 # drafter produced 1 valid draft for slot 4 + num_valid_drafts[2] = 3 # more than scheduled for slot 2 + qsl_buf = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + + trimmer = VariableDraftTrimmer( + num_valid_drafts, + qsl_buf, + num_bonus_tokens=1, + max_num_reqs=max_num_reqs, + max_total_logits=1024, + device=DEVICE, + ) + # Batch: [slot 4 (2 drafts scheduled), slot 2 (2 drafts), slot 0 (prefill)]. + idx_mapping = torch.tensor([4, 2, 0], dtype=torch.int64, device=DEVICE) + num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) + num_scheduled_tokens = np.array([3, 3, 7], dtype=np.int32) + + cu_num_logits, qsl = trimmer.trim( + idx_mapping, num_draft_tokens_per_req, num_scheduled_tokens + ) + # capacities = min(scheduled, num_valid) = [1, 2, 0] + assert cu_num_logits.cpu().tolist() == [0, 2, 5, 6] + # query lens = non-draft + capacities = [1+1, 1+2, 7+0] + assert qsl.cpu().tolist()[:4] == [0, 2, 5, 12] + # Padding tail equals the (GPU) batch total. + assert qsl.cpu().tolist()[4:] == [12] * (max_num_reqs - 3) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index be269550d3c4..e4a21328660a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2232,9 +2232,9 @@ def _free_encoder_inputs(self, request: Request) -> None: self.encoder_cache_manager.free_encoder_input(request, input_id) def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: - num_valid_list = draft_token_ids.num_valid_draft_tokens - for i, (req_id, spec_token_ids) in enumerate( - zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids) + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -2247,12 +2247,6 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: request.spec_token_ids = [] continue - # Variable-length drafters: truncate to the number of drafts - if num_valid_list is not None: - num_valid = num_valid_list[i] - if num_valid < len(spec_token_ids): - spec_token_ids = spec_token_ids[:num_valid] - # Add newly generated spec token ids to the request. if self.structured_output_manager.should_advance(request): metadata = request.structured_output_request @@ -2263,11 +2257,11 @@ def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput ) -> None: num_invalid_spec_tokens: dict[str, int] = {} - num_valid_list = draft_token_ids.num_valid_draft_tokens sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens - for i, (req_id, spec_token_ids) in enumerate( - zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids) + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, ): request = self.requests.get(req_id) if request is None or request.is_finished(): @@ -2281,13 +2275,7 @@ def update_draft_token_ids_in_output( orig_num_spec_tokens = len(placeholder_spec_tokens) # Trim drafts to scheduled number of spec tokens # (needed for chunked prefill case for example). - effective_num_spec_tokens = orig_num_spec_tokens - if num_valid_list is not None: - effective_num_spec_tokens = max( - 0, min(num_valid_list[i], orig_num_spec_tokens) - ) - - del spec_token_ids[effective_num_spec_tokens:] + del spec_token_ids[orig_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. if self.structured_output_manager.should_advance(request): metadata = request.structured_output_request diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 488bf2be5a40..55fe5b90c587 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -86,7 +86,7 @@ ) from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats -from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput +from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr from vllm.v1.structured_output import StructuredOutputManager @@ -168,13 +168,6 @@ def __init__( hash_block_size=hash_block_size, ) self.use_spec_decode = vllm_config.speculative_config is not None - # Variable-length drafters (e.g., ngram_gpu) need to retrieve - # num_valid_draft_tokens in async scheduling mode to truncate drafts. - # Fixed-length drafters (Eagle/Eagle3/MTP) don't need this extra RPC. - self.use_variable_length_drafter = ( - vllm_config.speculative_config is not None - and vllm_config.speculative_config.use_ngram_gpu() - ) self.check_for_draft_tokens = ( self.use_spec_decode or vllm_config.model_config.is_diffusion ) @@ -615,31 +608,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: ) self._attach_iteration_details(engine_core_outputs, iteration_details) - model_executed = scheduler_output.total_num_scheduled_tokens > 0 - self._maybe_update_async_draft_token_ids(model_executed) - - return engine_core_outputs, model_executed - - def _maybe_update_async_draft_token_ids( - self, model_executed: bool - ) -> "DraftTokenIds | None": - """ - Consume variable-length draft metadata from the just-completed - batch and apply it to scheduler request state. - """ - if not ( - self.async_scheduling - and self.use_variable_length_drafter - and model_executed - ): - return None - draft_token_ids = self.model_executor.take_draft_token_ids() - if ( - draft_token_ids is not None - and draft_token_ids.num_valid_draft_tokens is not None - ): - self.scheduler.update_draft_token_ids(draft_token_ids) - return draft_token_ids + return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 def post_step(self, model_executed: bool) -> None: # When using async scheduling we can't get draft token ids in advance, @@ -741,23 +710,14 @@ def step_with_batch_queue( ) self._attach_iteration_details(engine_core_outputs, iteration_details) - popped_batch_executed = scheduler_output.total_num_scheduled_tokens > 0 - async_draft_token_ids = self._maybe_update_async_draft_token_ids( - popped_batch_executed - ) - # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. if deferred_scheduler_output: # When draft tokens are used with structured output, validate them # before computing the grammar bitmask for the deferred request. - # In async scheduling the variable-length draft metadata is fetched - # eagerly above; otherwise pull it from the worker here. if self.check_for_draft_tokens: - draft_token_ids = ( - async_draft_token_ids or self.model_executor.take_draft_token_ids() - ) + draft_token_ids = self.model_executor.take_draft_token_ids() if draft_token_ids is not None: # Update the draft token ids in the scheduler output to # filter out the invalid spec tokens, which will be padded diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 29052d6ab136..0bbee7667527 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -382,8 +382,6 @@ class DraftTokenIds: req_ids: list[str] # num_reqs x num_draft_tokens draft_token_ids: list[list[int]] - # [num_reqs] - num_valid_draft_tokens: list[int] | None = None def make_empty_encoder_model_runner_output( diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 9a574a4faac5..858f19adcf1b 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -457,6 +457,9 @@ def combine_sampled_and_draft_tokens( cu_num_logits: torch.Tensor, num_logits: int, num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens + # Set when num_logits is only an upper bound (GPU draft trimming), so + # unwritten trailing entries hold benign in-bounds indices. + zero_init_logits_indices: bool = False, ) -> torch.Tensor: assert num_new_sampled_tokens in (0, 1), ( f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" @@ -465,7 +468,8 @@ def combine_sampled_and_draft_tokens( num_reqs = idx_mapping.shape[0] num_speculative_steps = draft_tokens.shape[-1] - logits_indices = torch.empty( + alloc = torch.zeros if zero_init_logits_indices else torch.empty + logits_indices = alloc( num_logits, dtype=torch.int64, device=input_ids.device, @@ -699,12 +703,21 @@ def expand_idx_mapping( total_num_logits: int, cu_num_logits: torch.Tensor, max_expand_len: int, + # Set when total_num_logits is only an upper bound (GPU draft trimming), + # so unwritten trailing entries hold benign in-bounds values. + zero_init: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: num_reqs = idx_mapping.shape[0] - expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) - expanded_local_pos = torch.empty( - total_num_logits, dtype=torch.int32, device=idx_mapping.device - ) + if zero_init: + expanded_idx_mapping = idx_mapping.new_zeros(total_num_logits) + expanded_local_pos = torch.zeros( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) + else: + expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) _expand_idx_mapping_kernel[(num_reqs,)]( idx_mapping, expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 21ea66f31abc..25ce7b5aba62 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -130,7 +130,9 @@ from vllm.v1.worker.gpu.spec_decode import init_speculator from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, + VariableDraftTrimmer, maybe_create_adaptive_verification_manager, + maybe_create_draft_trimmer, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, @@ -283,6 +285,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): use_dense_all_token_ids=use_dense_all_token_ids, ) self.adaptive_verification: AdaptiveVerificationManager | None = None + self.draft_trimmer: VariableDraftTrimmer | None = None self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -551,6 +554,26 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, max_total_logits=get_max_chunk_logits(self.vocab_size), ) + # Variable-length drafters (ngram_gpu) trim scheduled draft slots to + # the drafter's valid counts on GPU, when supported. + self.draft_trimmer = None + if self.speculator is not None and self.adaptive_verification is None: + self.draft_trimmer = maybe_create_draft_trimmer( + speculator=self.speculator, + attn_groups=self.attn_groups, + attn_cg_support=attn_cg_support, + uses_full_cudagraphs=self.compilation_config.cudagraph_mode is not None + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE, + has_lora=self.lora_config is not None, + uses_pipeline_parallel=self.use_pp, + uses_context_parallel=self.dcp_size > 1 + or self.parallel_config.prefill_context_parallel_size > 1, + query_start_loc=self.input_buffers.query_start_loc, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, + max_total_logits=get_max_chunk_logits(self.vocab_size), + max_num_reqs=self.max_num_reqs, + device=self.device, + ) self.block_tables = BlockTables( block_sizes=block_sizes, @@ -576,6 +599,14 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + elif ( + self.draft_trimmer is not None + and self.compilation_config.cudagraph_mode is not None + and self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ): + # Trimmed decode batches have per-request varlen queries, which + # uniform-decode full graphs cannot replay. + self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, @@ -591,7 +622,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cudagraph_mode, decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, - varlen_decode=self.adaptive_verification is not None, + varlen_decode=self.adaptive_verification is not None + or self.draft_trimmer is not None, ) check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): @@ -1155,6 +1187,16 @@ def prepare_inputs( adaptive_verification = ( self.adaptive_verification if num_draft_tokens_per_req is not None else None ) + draft_trimmer = None + if ( + adaptive_verification is None + and num_draft_tokens_per_req is not None + and self.draft_trimmer is not None + # The chunked logits path indexes by the CPU (untrimmed) offsets, + # which cannot address the trimmed layout. + and total_num_logits <= self.draft_trimmer.max_total_logits + ): + draft_trimmer = self.draft_trimmer num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across @@ -1184,9 +1226,22 @@ def prepare_inputs( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + elif draft_trimmer is not None: + # Clamp scheduled draft slots to the drafter's valid counts on + # GPU. CPU-side totals remain upper bounds; the trimmed gap is + # treated as padding downstream. + cu_num_logits, query_start_loc = draft_trimmer.trim( + idx_mapping, num_draft_tokens_per_req, num_scheduled_tokens_np + ) if draft_tokens: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( - idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len + idx_mapping, + total_num_logits, + cu_num_logits, + self.decode_query_len, + # With GPU trimming, total_num_logits is an upper bound; the + # gap must hold benign (in-bounds) values. + zero_init=draft_trimmer is not None, ) query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] query_start_loc = query_start_loc[: num_reqs_padded + 1] @@ -1239,6 +1294,7 @@ def prepare_inputs( cu_num_logits, total_num_logits, self.model_state.num_new_sampled_tokens_per_step, + zero_init_logits_indices=draft_trimmer is not None, ) # CPU upper bound on seq_lens; padded entries left at zero. @@ -1295,7 +1351,7 @@ def prepare_inputs( prompt_lens=prompt_lens, max_query_len=( int(num_scheduled_tokens_upper_bound.max()) - if adaptive_verification is not None + if adaptive_verification is not None or draft_trimmer is not None else None ), ) @@ -1809,7 +1865,7 @@ def sample_tokens( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] - draft_tokens, num_valid_draft_tokens = self.speculator.propose( + draft_tokens = self.speculator.propose( input_batch, attn_metadata, slot_mappings_by_layer, @@ -1828,11 +1884,6 @@ def sample_tokens( self.adaptive_verification.record_confidences( self.speculator.draft_token_confidence_probs, input_batch ) - # Pass num_valid_draft_tokens so variable-length drafters (ngram_gpu) - # can truncate drafts in async scheduling mode. - self.draft_tokens_handler.set_draft_tokens( - input_batch, draft_tokens, num_valid_draft_tokens - ) if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 12dc8c7c656e..7bbcf049bb00 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -65,6 +65,161 @@ def _assign_draft_token_budget( ) +def build_verification_layout( + capacities: torch.Tensor, + num_non_draft_tokens: torch.Tensor, + num_bonus_tokens: int, + cu_num_logits: torch.Tensor, + query_start_loc: torch.Tensor, + num_tokens: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build GPU cu_num_logits / query_start_loc from per-request admitted + draft counts. + + Trailing (padding) query_start_loc entries are filled with the batch + total: the exact CPU value when known (`num_tokens`), otherwise the GPU + cumsum tail, so downstream kernels treat everything past the real tokens + as padding. + """ + num_reqs = capacities.shape[0] + cu_num_logits[:1].zero_() + torch.cumsum( + capacities + num_bonus_tokens, + dim=0, + out=cu_num_logits[1 : num_reqs + 1], + ) + query_start_loc[:1].zero_() + torch.cumsum( + capacities + num_non_draft_tokens, + dim=0, + out=query_start_loc[1 : num_reqs + 1], + ) + if num_tokens is not None: + query_start_loc[num_reqs + 1 :].fill_(num_tokens) + else: + query_start_loc[num_reqs + 1 :] = query_start_loc[num_reqs] + return cu_num_logits[: num_reqs + 1], query_start_loc + + +class VariableDraftTrimmer: + """GPU-side verification trimming for variable-length drafters (ngram). + + The drafter records per-request valid draft counts on GPU in + `num_valid_drafts`. The scheduler still schedules the full + num_speculative_tokens per request; at the next step this trimmer clamps + each request's scheduled draft slots to the recorded count and rebuilds + cu_num_logits / query_start_loc on device, so the CPU keeps only upper + bounds. Trimmed slots surface as ordinary rejections through the + existing num_rejected accounting — no scheduler round-trip and no + CPU<->GPU synchronization. + """ + + def __init__( + self, + num_valid_drafts: torch.Tensor, + query_start_loc: torch.Tensor, + num_bonus_tokens: int, + max_num_reqs: int, + max_total_logits: int, + device: torch.device, + ): + self.num_valid_drafts = num_valid_drafts + self.query_start_loc = query_start_loc + self.num_bonus_tokens = num_bonus_tokens + # Rejection sampling chunks logits by the CPU (untrimmed) offsets, + # which cannot address the compacted layout; skip trimming for + # batches that would not fit in one chunk. + self.max_total_logits = max_total_logits + self._capacities = torch.empty(max_num_reqs, dtype=torch.int32, device=device) + self._num_non_draft_tokens = torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ) + self._cu_num_logits = torch.empty( + max_num_reqs + 1, dtype=torch.int32, device=device + ) + + def trim( + self, + idx_mapping: torch.Tensor, + num_draft_tokens_per_req: np.ndarray, + num_scheduled_tokens_np: np.ndarray, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_reqs = idx_mapping.shape[0] + capacities = self._capacities[:num_reqs] + async_copy_to_gpu(num_draft_tokens_per_req, out=capacities) + torch.minimum(capacities, self.num_valid_drafts[idx_mapping], out=capacities) + num_non_draft_tokens = self._num_non_draft_tokens[:num_reqs] + async_copy_to_gpu( + num_scheduled_tokens_np - num_draft_tokens_per_req, + out=num_non_draft_tokens, + ) + return build_verification_layout( + capacities, + num_non_draft_tokens, + self.num_bonus_tokens, + self._cu_num_logits, + self.query_start_loc, + num_tokens=None, + ) + + +def maybe_create_draft_trimmer( + *, + speculator, + attn_groups: list[list["AttentionGroup"]], + attn_cg_support: "AttentionCGSupportInfo", + uses_full_cudagraphs: bool, + has_lora: bool, + uses_pipeline_parallel: bool, + uses_context_parallel: bool, + query_start_loc: torch.Tensor, + num_bonus_tokens: int, + max_total_logits: int, + max_num_reqs: int, + device: torch.device, +) -> VariableDraftTrimmer | None: + """Create a VariableDraftTrimmer when the drafter and environment support + GPU-side trimming; otherwise fall back (with a log) to verifying the full + padded drafts, which is correct but wastes verification compute.""" + num_valid_drafts = getattr(speculator, "num_valid_drafts", None) + if not getattr(speculator, "trims_drafts_on_gpu", False): + return None + assert num_valid_drafts is not None + + reason = None + backend = get_query_lens_mismatch_unsupported_backend(attn_groups) + if backend is not None: + reason = f"the {backend} attention backend" + elif ( + uses_full_cudagraphs + and attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS + ): + reason = f"varlen decode cudagraphs with {attn_cg_support.min_cg_attn_backend}" + elif has_lora: + reason = "LoRA" + elif uses_pipeline_parallel: + reason = "pipeline parallelism" + elif uses_context_parallel: + reason = "context parallelism" + + if reason is not None: + logger.info( + "GPU draft trimming is not supported with %s; invalid draft " + "slots will be verified (and rejected) instead of trimmed.", + reason, + ) + return None + + return VariableDraftTrimmer( + num_valid_drafts, + query_start_loc, + num_bonus_tokens, + max_num_reqs, + max_total_logits, + device, + ) + + def build_cost_tables_from_curves( draft_curve: list[tuple[int, float]], verify_curve: list[tuple[int, float]], @@ -415,24 +570,15 @@ def reallocate_drafts( num_non_draft_tokens, out=num_non_draft_tokens_gpu, ) - self._cu_num_logits[:1].zero_() - torch.cumsum( - capacities + self.num_bonus_tokens, - dim=0, - out=self._cu_num_logits[1 : num_reqs + 1], - ) - self.query_start_loc[:1].zero_() - torch.cumsum( - capacities + num_non_draft_tokens_gpu, - dim=0, - out=self.query_start_loc[1 : num_reqs + 1], - ) - self.query_start_loc[num_reqs + 1 :].fill_(num_tokens) - return ( - self._cu_num_logits[: num_reqs + 1], + cu_num_logits, query_start_loc = build_verification_layout( + capacities, + num_non_draft_tokens_gpu, + self.num_bonus_tokens, + self._cu_num_logits, self.query_start_loc, - draft_budget, + num_tokens, ) + return cu_num_logits, query_start_loc, draft_budget def maybe_create_adaptive_verification_manager( diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index d6ab963bc62a..57430529f61a 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -224,7 +224,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + ) -> torch.Tensor: num_tokens = input_batch.num_tokens num_tokens_padded = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs @@ -313,7 +313,7 @@ def propose( if self.num_speculative_steps == 1: # Early exit. - return self.draft_tokens[:num_reqs, :1], None + return self.draft_tokens[:num_reqs, :1] # Prepare the inputs for the decode steps. prepare_decode_inputs( @@ -354,7 +354,7 @@ def propose( ) self.on_multi_step_decode_end(num_reqs) - return self.draft_tokens[:num_reqs], None + return self.draft_tokens[:num_reqs] @torch.inference_mode() def _run_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 85ca97fee6b8..97c284c03d4a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -365,7 +365,7 @@ def propose( num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=CUDAGraphMode.NONE, ) - return self.draft_tokens[:num_reqs], None + return self.draft_tokens[:num_reqs] # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. @@ -465,7 +465,7 @@ def propose( cudagraph_runtime_mode=batch_desc.cg_mode, ) - return self.draft_tokens[:num_reqs], None + return self.draft_tokens[:num_reqs] @triton.jit diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index ada9a3c81fe0..25fae8dfc1ed 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -7,10 +7,8 @@ import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import set_forward_context from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.worker.gpu.input_batch import InputBatch @@ -18,343 +16,161 @@ from vllm.v1.worker.gpu.states import RequestState -@support_torch_compile( - dynamic_arg_dims={ - "token_ids": 0, - "seq_lens": 0, - "valid_mask": 0, - "last_sampled": 0, - } -) -class _NgramKernel(nn.Module): - """GPU-accelerated N-gram proposer using fully async tensor operations.""" - - def __init__(self, min_n: int, max_n: int, k: int): - super().__init__() - assert 1 <= min_n <= max_n, ( - f"min_n must be in [1, max_n]; got min_n={min_n}, max_n={max_n}" - ) - assert k >= 1 - self.min_n = min_n - self.max_n = max_n - self.k = k - self.num_sizes = max_n - min_n + 1 - - def forward( - self, - token_ids: torch.Tensor, # [B, L] int32 - seq_lens: torch.Tensor, # [B] int32 (current total_len per req) - valid_mask: torch.Tensor, # [B] bool (row eligible for n-gram lookup) - last_sampled: torch.Tensor, # [B] int64 (fallback for -1 positions) - ) -> tuple[torch.Tensor, torch.Tensor]: - """For each row, find the longest n-gram suffix match in its context - and propose the next k tokens. Fully vectorized; no data-dependent - control flow so the kernel stays torch.compile / CUDA-graph friendly. - """ - B, L = token_ids.shape - device = token_ids.device - - # Phase 1: For each n in [min_n, max_n], find the right-most position - # of the length-n suffix inside each row (excluding the suffix itself). - # first_match_pos[b, i] = match position for n=min_n+i, or -1. - first_match_pos = torch.full( - (B, self.num_sizes), -1, dtype=torch.long, device=device - ) - batch_idx = torch.arange(B, device=device) - - for i, n in enumerate(range(self.min_n, self.max_n + 1)): - # Sliding length-n windows; needle = length-n suffix of each row. - windows = token_ids.unfold(1, n, 1) - num_windows = windows.shape[1] - suffix_start = (seq_lens.long() - n).clamp(min=0) - offsets = torch.arange(n, device=device) - suffix_idx = suffix_start.unsqueeze(1) + offsets - suffix = torch.gather(token_ids, 1, suffix_idx) - matches = (windows == suffix.unsqueeze(1)).all(dim=-1) - - # Mask out windows that overlap (or live past) the suffix. - max_valid_pos = seq_lens.long() - n - 1 - window_pos = torch.arange(num_windows, device=device) - matches = matches & (window_pos.unsqueeze(0) <= max_valid_pos.unsqueeze(1)) - - # Right-most match via argmax on (pos if match else -1); re-check - # the chosen index to distinguish a real match from the fallback. - matched_indices = torch.where(matches, window_pos.unsqueeze(0), -1) - idx = matched_indices.argmax(dim=1) - has_match = matches[batch_idx, idx] - first_match_pos[:, i] = torch.where(has_match, idx, -1) - - # Phase 2: Pick the largest n that produced a match (right-most True - # along the n-axis). - best_i = (first_match_pos >= 0).int().flip(dims=[1]).argmax(dim=1) - best_i = self.num_sizes - 1 - best_i - best_pos = first_match_pos[batch_idx, best_i] - ngram_lens_table = torch.arange( - self.min_n, self.max_n + 1, device=device, dtype=torch.long - ) - best_n = ngram_lens_table[best_i] - has_any = best_pos >= 0 - - # Phase 3: Gather the next k tokens after the match. No-match rows - # use draft_start=0 to keep indices in-bounds; they're masked below. - draft_start = torch.where( - has_any, - best_pos + best_n, - torch.zeros_like(best_pos), - ) - k_range = torch.arange(self.k, device=device) - draft_idx = (draft_start.unsqueeze(1) + k_range).clamp_(0, L - 1) - drafts = torch.gather(token_ids, 1, draft_idx).to(torch.int64) - - # Phase 4: A slot j is valid iff the row has a match, valid_mask is - # True, and j < seq_len - draft_start (so we don't read past context). - # num_valid = length of the leading run of True values per row. - tokens_available = (seq_lens.long() - draft_start).clamp_(min=0) - valid_positions = k_range.unsqueeze(0) < tokens_available.unsqueeze(1) - row_valid = has_any & valid_mask - leading_valid_mask = valid_positions & row_valid.unsqueeze(1) - cum_valid = leading_valid_mask.int().cumsum(dim=1) - positions = torch.arange(1, self.k + 1, device=device) - num_valid = (cum_valid == positions.unsqueeze(0)).int().sum(dim=1) - - # Phase 5: Replace invalid slots with last_sampled. - safe_drafts = torch.where( - leading_valid_mask, - drafts, - last_sampled.view(-1, 1).expand(B, self.k), - ) - return safe_drafts, num_valid.to(torch.int32) - - -if HAS_TRITON: - - @triton.jit - def _ngram_scan_kernel( - token_ids_ptr, # *int32 [B, L] - seq_lens_ptr, # *int32 [B] - valid_mask_ptr, # *int8 [B] - scratch_ptr, # *int64 [B, N_BLOCKS] (output) - L, # int64 scalar - L_PLUS_1, # int64 scalar (= L + 1, used for packing) - N_BLOCKS, # int64 scalar (stride of scratch's second dim) - MIN_N: tl.constexpr, - MAX_N: tl.constexpr, - MAX_N_PO2: tl.constexpr, - BLOCK_L: tl.constexpr, - ): - b = tl.program_id(0).to(tl.int64) - blk = tl.program_id(1).to(tl.int64) - L_ = tl.cast(L, tl.int64) - Lp1 = tl.cast(L_PLUS_1, tl.int64) - NB = tl.cast(N_BLOCKS, tl.int64) - - seq_len = tl.load(seq_lens_ptr + b).to(tl.int64) - row_valid = tl.load(valid_mask_ptr + b).to(tl.int1) - eligible_row = row_valid & (seq_len >= MIN_N) - - scratch_off = b * NB + blk - - # Ineligible rows (or blocks past the last valid pos) write 0. - if not eligible_row: - tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64)) - return - - row_off = b * L_ - - # Load the length-MAX_N suffix once into registers. - suf_iota = tl.arange(0, MAX_N_PO2).to(tl.int64) - suf_pos = seq_len - MAX_N + suf_iota - suf_in_range = (suf_iota < MAX_N) & (suf_pos >= 0) & (suf_pos < seq_len) - suffix = tl.load( - token_ids_ptr + row_off + suf_pos, - mask=suf_in_range, - other=-1, - ).to(tl.int32) - - pos_iota = tl.arange(0, BLOCK_L).to(tl.int64) - pos = blk * BLOCK_L + pos_iota # ascending - - best_score = tl.zeros([BLOCK_L], dtype=tl.int64) - - for n_iter in tl.static_range(MIN_N, MAX_N + 1): - max_pos_n = seq_len - n_iter - 1 - match = (pos >= 0) & (pos <= max_pos_n) - for j in tl.static_range(0, n_iter): - tok = tl.load( - token_ids_ptr + row_off + (pos + j), - mask=match, - other=0, - ).to(tl.int32) - suf_idx = (MAX_N - n_iter) + j - suf_val = tl.sum(tl.where(suf_iota == suf_idx, suffix, 0)) - match = match & (tok == suf_val) - - cand = n_iter * Lp1 + pos + 1 - best_score = tl.where(match, cand, best_score) - - block_best = tl.max(best_score, axis=0) - tl.store(scratch_ptr + scratch_off, block_best) - - @triton.jit - def _ngram_finalize_kernel( - token_ids_ptr, # *int32 [B, L] - seq_lens_ptr, # *int32 [B] - valid_mask_ptr, # *int8 [B] - last_sampled_ptr, # *int64 [B] - scratch_ptr, # *int64 [B, N_BLOCKS] - drafts_ptr, # *int64 [B, K] (output) - num_valid_ptr, # *int32 [B] (output) - L, - L_PLUS_1, - N_BLOCKS, - K: tl.constexpr, - K_PO2: tl.constexpr, - N_BLOCKS_PO2: tl.constexpr, - ): - b = tl.program_id(0).to(tl.int64) - L_ = tl.cast(L, tl.int64) - Lp1 = tl.cast(L_PLUS_1, tl.int64) - NB = tl.cast(N_BLOCKS, tl.int64) - - nb_iota = tl.arange(0, N_BLOCKS_PO2).to(tl.int64) - nb_in_range = nb_iota < NB - block_scores = tl.load( - scratch_ptr + b * NB + nb_iota, - mask=nb_in_range, - other=0, - ) - score = tl.max(block_scores, axis=0) - - seq_len = tl.load(seq_lens_ptr + b).to(tl.int64) - row_valid = tl.load(valid_mask_ptr + b).to(tl.int1) - last_tok = tl.load(last_sampled_ptr + b) - - has_match = score > 0 - s1 = score - 1 - best_n = tl.where(has_match, s1 // Lp1, tl.zeros_like(s1)) - best_pos = tl.where(has_match, s1 - best_n * Lp1, tl.zeros_like(s1)) - draft_start = tl.where(has_match, best_pos + best_n, tl.zeros_like(s1)) - - tokens_avail = tl.maximum(seq_len - draft_start, 0) - write_ok = row_valid & has_match - nv = tl.where(write_ok, tl.minimum(tl.cast(K, tl.int64), tokens_avail), 0) - tl.store(num_valid_ptr + b, nv.to(tl.int32)) - - row_off = b * L_ - k_iota = tl.arange(0, K_PO2).to(tl.int64) - k_in_range = k_iota < K - gather_idx = tl.minimum(draft_start + k_iota, L_ - 1) - slot_valid = (k_iota < tokens_avail) & write_ok & k_in_range - gathered = tl.load( - token_ids_ptr + row_off + gather_idx, - mask=slot_valid, - other=0, - ).to(tl.int64) - out = tl.where(slot_valid, gathered, last_tok) - tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) - - -_NGRAM_SCRATCH: dict[tuple, torch.Tensor] = {} - - -def _get_ngram_scratch(B: int, n_blocks: int, device: torch.device) -> torch.Tensor: - key = (device, B, n_blocks) - buf = _NGRAM_SCRATCH.get(key) - if buf is None: - buf = torch.empty((B, n_blocks), dtype=torch.int64, device=device) - _NGRAM_SCRATCH[key] = buf - return buf - - -def _ngram_propose_triton( - token_ids: torch.Tensor, - seq_lens: torch.Tensor, - valid_mask: torch.Tensor, - last_sampled: torch.Tensor, - min_n: int, - max_n: int, - k: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """ - For each row, find the longest n-gram suffix match in its context - and propose the next k tokens. - """ - B, L = token_ids.shape - device = token_ids.device - - drafts = torch.empty((B, k), dtype=torch.int64, device=device) - num_valid = torch.empty((B,), dtype=torch.int32, device=device) - - if B == 0: - return drafts, num_valid - - tok = token_ids.contiguous().to(torch.int32) - seq = seq_lens.contiguous().to(torch.int32) - vmask = valid_mask.contiguous().to(torch.bool).to(torch.int8) - last = last_sampled.contiguous().to(torch.int64).view(-1) - - if L >= 1024: - BLOCK_L = 256 - elif L >= 256: - BLOCK_L = 128 - elif L >= 64: - BLOCK_L = 64 - else: - BLOCK_L = max(16, triton.next_power_of_2(max(L, 1))) - - K_PO2 = max(1, triton.next_power_of_2(k)) - MAX_N_PO2 = max(1, triton.next_power_of_2(max_n)) - n_blocks = (L + BLOCK_L - 1) // BLOCK_L - n_blocks_po2 = max(1, triton.next_power_of_2(n_blocks)) - - scratch = _get_ngram_scratch(B, n_blocks, device) - - L_plus_1 = L + 1 - _ngram_scan_kernel[(B, n_blocks)]( - tok, - seq, - vmask, - scratch, - L, - L_plus_1, - n_blocks, - min_n, - max_n, - MAX_N_PO2, - BLOCK_L, - num_warps=4, - num_stages=2, +@triton.jit +def _ngram_scan_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] batch_idx -> req_state_idx + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + scratch_ptr, # *int64 [B, scratch_stride] (output) + scratch_stride, + L, # int64 scalar (= max_model_len) + MIN_N: tl.constexpr, + MAX_N: tl.constexpr, + MAX_N_PO2: tl.constexpr, + BLOCK_L: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + blk = tl.program_id(1).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + eligible_row = (num_sampled > 0) & (seq_len >= MIN_N) + + scratch_off = b * scratch_stride + blk + + # Ineligible rows, and blocks fully past the last candidate match + # position, write 0 and exit. + if not (eligible_row & (blk * BLOCK_L <= seq_len - MIN_N - 1)): + tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64)) + return + + row_off = req_state_idx * token_ids_stride + + # Load the length-MAX_N suffix once into registers. + suf_iota = tl.arange(0, MAX_N_PO2).to(tl.int64) + suf_pos = seq_len - MAX_N + suf_iota + suf_in_range = (suf_iota < MAX_N) & (suf_pos >= 0) & (suf_pos < seq_len) + suffix = tl.load( + token_ids_ptr + row_off + suf_pos, + mask=suf_in_range, + other=-1, + ).to(tl.int32) + + pos_iota = tl.arange(0, BLOCK_L).to(tl.int64) + pos = blk * BLOCK_L + pos_iota # ascending + + best_score = tl.zeros([BLOCK_L], dtype=tl.int64) + + for n_iter in tl.static_range(MIN_N, MAX_N + 1): + max_pos_n = seq_len - n_iter - 1 + match = (pos >= 0) & (pos <= max_pos_n) + for j in tl.static_range(0, n_iter): + tok = tl.load( + token_ids_ptr + row_off + (pos + j), + mask=match, + other=0, + ).to(tl.int32) + suf_idx = (MAX_N - n_iter) + j + suf_val = tl.sum(tl.where(suf_iota == suf_idx, suffix, 0)) + match = match & (tok == suf_val) + + # Pack (n, pos) so a single max yields longest-n, rightmost-pos. + cand = n_iter * Lp1 + pos + 1 + best_score = tl.where(match, cand, best_score) + + block_best = tl.max(best_score, axis=0) + tl.store(scratch_ptr + scratch_off, block_best) + + +@triton.jit +def _ngram_finalize_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + last_sampled_ptr, # *int64 [max_num_reqs] + scratch_ptr, # *int64 [B, scratch_stride] + scratch_stride, + drafts_ptr, # *int64 [B, K] (output, batch indexed) + num_valid_ptr, # *int32 [max_num_reqs] (output, req-slot indexed) + L, + N_BLOCKS, + K: tl.constexpr, + K_PO2: tl.constexpr, + N_BLOCKS_PO2: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + NB = tl.cast(N_BLOCKS, tl.int64) + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + + nb_iota = tl.arange(0, N_BLOCKS_PO2).to(tl.int64) + nb_in_range = nb_iota < NB + block_scores = tl.load( + scratch_ptr + b * scratch_stride + nb_iota, + mask=nb_in_range, + other=0, ) - - _ngram_finalize_kernel[(B,)]( - tok, - seq, - vmask, - last, - scratch, - drafts, - num_valid, - L, - L_plus_1, - n_blocks, - k, - K_PO2, - n_blocks_po2, - num_warps=2, - num_stages=1, - ) - return drafts, num_valid + score = tl.max(block_scores, axis=0) + + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + last_tok = tl.load(last_sampled_ptr + req_state_idx) + + has_match = score > 0 + s1 = score - 1 + best_n = tl.where(has_match, s1 // Lp1, tl.zeros_like(s1)) + best_pos = tl.where(has_match, s1 - best_n * Lp1, tl.zeros_like(s1)) + draft_start = tl.where(has_match, best_pos + best_n, tl.zeros_like(s1)) + + tokens_avail = tl.maximum(seq_len - draft_start, 0) + write_ok = (num_sampled > 0) & has_match + nv = tl.where(write_ok, tl.minimum(tl.cast(K, tl.int64), tokens_avail), 0) + tl.store(num_valid_ptr + req_state_idx, nv.to(tl.int32)) + + row_off = req_state_idx * token_ids_stride + k_iota = tl.arange(0, K_PO2).to(tl.int64) + k_in_range = k_iota < K + gather_idx = tl.minimum(draft_start + k_iota, tl.cast(L, tl.int64) - 1) + slot_valid = (k_iota < tokens_avail) & write_ok & k_in_range + gathered = tl.load( + token_ids_ptr + row_off + gather_idx, + mask=slot_valid, + other=0, + ).to(tl.int64) + # Invalid slots fall back to the last sampled token; they are either + # trimmed from the verification batch on GPU or verified as ordinary + # (rejectable) drafts, so the fill value only affects efficiency. + out = tl.where(slot_valid, gathered, last_tok) + tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) class NgramGPUSpeculator: - """ - V2-compatible GPU n-gram speculator. + """V2-compatible GPU n-gram speculator. + + Drafts are proposed entirely on GPU by scanning each request's token + history (`RequestState.all_token_ids`) in place for the longest n-gram + suffix match. The per-request count of usable drafts is written to the + persistent `num_valid_drafts` tensor, which the model runner's draft + trimmer consumes at the next step to shrink verification on GPU without + any CPU<->GPU synchronization. """ supports_mm_inputs = False draft_logits = None + # Signals that num_valid_drafts holds per-request valid draft counts + # for GPU-side verification trimming. + trims_drafts_on_gpu = True def __init__(self, vllm_config: VllmConfig, device: torch.device): + if not HAS_TRITON: + raise RuntimeError("ngram_gpu speculative decoding requires Triton.") spec = vllm_config.speculative_config assert spec is not None assert spec.prompt_lookup_min is not None, ( @@ -363,6 +179,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): assert spec.prompt_lookup_max is not None, ( "prompt_lookup_max must be configured for ngram_gpu" ) + assert 1 <= spec.prompt_lookup_min <= spec.prompt_lookup_max self.vllm_config = vllm_config self.device = device @@ -375,21 +192,32 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs self.max_model_len: int = vllm_config.model_config.max_model_len - # Triton is the default fast path; the torch.compile kernel is - # only constructed (and used) when Triton is unavailable. - self.use_triton: bool = HAS_TRITON - if self.use_triton: - self.kernel: _NgramKernel + L = self.max_model_len + if L >= 1024: + self.block_l = 256 + elif L >= 256: + self.block_l = 128 + elif L >= 64: + self.block_l = 64 else: - self.kernel = ( - _NgramKernel( - min_n=self.min_n, - max_n=self.max_n, - k=self.num_speculative_steps, - ) - .to(device) - .eval() - ) + self.block_l = max(16, triton.next_power_of_2(max(L, 1))) + self.n_blocks = triton.cdiv(L, self.block_l) + + self.scratch = torch.zeros( + (self.max_num_reqs, self.n_blocks), dtype=torch.int64, device=device + ) + # Per request-slot count of usable drafts from the latest proposal, + # consumed by the model runner's GPU draft trimmer. + self.num_valid_drafts = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + # Batch-ordered draft output, scattered into RequestState.draft_tokens + # by the model runner (same contract as the model-based speculators). + self.drafts = torch.zeros( + (self.max_num_reqs, self.num_speculative_steps), + dtype=torch.int64, + device=device, + ) self.req_states: RequestState | None = None @@ -428,40 +256,54 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - assert self.req_states is not None, ( - "NgramGPUSpeculator.req_states was not injected by the model " - "runner. Ensure model_runner sets `speculator.req_states = " - "self.req_states` after RequestState is constructed." + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + if dummy_run: + # No persistent request state may be touched during dummy runs. + return self.drafts[:num_reqs] + + req_states = self.req_states + assert req_states is not None, ( + "NgramGPUSpeculator.req_states was not injected by the model runner." ) - idx_mapping_long = input_batch.idx_mapping.long() - - active_tokens: torch.Tensor = self.req_states.all_token_ids.gpu[ - idx_mapping_long - ] - active_seq_lens: torch.Tensor = self.req_states.total_len.gpu[idx_mapping_long] - active_last_sampled: torch.Tensor = last_sampled.view(-1)[idx_mapping_long] - - valid_mask = (num_sampled > 0) & (active_seq_lens >= self.min_n) - - if self.use_triton: - drafts, num_valid = _ngram_propose_triton( - active_tokens, - active_seq_lens, - valid_mask, - active_last_sampled, - self.min_n, - self.max_n, - self.num_speculative_steps, - ) - else: - with set_forward_context(None, self.vllm_config): - drafts, num_valid = self.kernel( - active_tokens, - active_seq_lens, - valid_mask, - active_last_sampled, - ) - - return drafts, num_valid + token_ids = req_states.all_token_ids.gpu + idx_mapping = input_batch.idx_mapping + + _ngram_scan_kernel[(num_reqs, self.n_blocks)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + self.scratch, + self.scratch.stride(0), + self.max_model_len, + self.min_n, + self.max_n, + max(1, triton.next_power_of_2(self.max_n)), + self.block_l, + num_warps=4, + num_stages=2, + ) + + _ngram_finalize_kernel[(num_reqs,)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + last_sampled.view(-1), + self.scratch, + self.scratch.stride(0), + self.drafts, + self.num_valid_drafts, + self.max_model_len, + self.n_blocks, + self.num_speculative_steps, + max(1, triton.next_power_of_2(self.num_speculative_steps)), + max(1, triton.next_power_of_2(self.n_blocks)), + num_warps=2, + num_stages=1, + ) + return self.drafts[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 709783b848e0..e25672e953bd 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -17,23 +17,17 @@ def __init__(self, device: torch.device | None = None): self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None - self.num_valid_draft_tokens_np: np.ndarray | None = None self.num_draft_tokens: int = 0 def set_draft_tokens( - self, - input_batch: InputBatch, - draft_tokens: torch.Tensor, - num_valid_draft_tokens: torch.Tensor | None = None, + self, input_batch: InputBatch, draft_tokens: torch.Tensor ) -> None: self.req_ids = input_batch.req_ids self.num_draft_tokens = draft_tokens.shape[1] - - needs_draft_copy = input_batch.has_structured_output_reqs - - if not needs_draft_copy and num_valid_draft_tokens is None: + if not input_batch.has_structured_output_reqs: + # No draft token validation needs to be performed by + # the scheduler for this batch. self.draft_tokens_np = None - self.num_valid_draft_tokens_np = None return # For spec decoding + structured outputs, we must transfer the @@ -41,46 +35,21 @@ def set_draft_tokens( current_stream = torch.cuda.current_stream(self.device) self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): - # draft_tokens / num_valid_draft_tokens are temporary allocations on - # the main stream and read here on copy_stream; without record_stream, - # the caching allocator may reuse their memory before the async copy - # executes. - if needs_draft_copy: - self.draft_tokens_np = async_copy_to_np(draft_tokens) - draft_tokens.record_stream(self.copy_stream) - else: - self.draft_tokens_np = None - if num_valid_draft_tokens is not None: - self.num_valid_draft_tokens_np = async_copy_to_np( - num_valid_draft_tokens - ) - num_valid_draft_tokens.record_stream(self.copy_stream) - else: - self.num_valid_draft_tokens_np = None + self.draft_tokens_np = async_copy_to_np(draft_tokens) + # draft_tokens is a temporary allocation on the main stream and read here on + # copy_stream; without record_stream, the caching allocator may reuse its + # memory before the async copy executes. + draft_tokens.record_stream(self.copy_stream) self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: - if ( - self.draft_tokens_np is not None - or self.num_valid_draft_tokens_np is not None - ): - self.copy_event.synchronize() - if self.draft_tokens_np is not None: + self.copy_event.synchronize() draft_token_ids = self.draft_tokens_np.tolist() else: # This case only happens when async scheduling is disabled. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] - - num_valid_list: list[int] | None = None - if self.num_valid_draft_tokens_np is not None: - num_valid_list = self.num_valid_draft_tokens_np.tolist() - - return DraftTokenIds( - req_ids=self.req_ids, - draft_token_ids=draft_token_ids, - num_valid_draft_tokens=num_valid_list, - ) + return DraftTokenIds(self.req_ids, draft_token_ids) def get_parallel_drafting_token_id(hf_config) -> int: From d84087e863339b291c996aaeab9635ec5d66a982 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 14 Aug 2026 18:18:24 -0700 Subject: [PATCH 26/34] [Spec Decode] ngram_gpu: conservative trim gating for auto cudagraph mode Unset (auto) cudagraph_mode may resolve to full graphs after the trimmer is created, so treat it as full-graph usage when checking varlen decode capture support. Also log when GPU draft trimming is enabled. Co-authored-by: Claude Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/model_runner.py | 17 ++++++++++------- .../gpu/spec_decode/adaptive_verification.py | 1 + 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 25ce7b5aba62..f6be5e20202a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -562,8 +562,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: speculator=self.speculator, attn_groups=self.attn_groups, attn_cg_support=attn_cg_support, - uses_full_cudagraphs=self.compilation_config.cudagraph_mode is not None - and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE, + # Unset (auto) mode may resolve to full graphs, so treat it + # as full-graph usage. + uses_full_cudagraphs=self.compilation_config.cudagraph_mode is None + or self.compilation_config.cudagraph_mode.has_full_cudagraphs(), has_lora=self.lora_config is not None, uses_pipeline_parallel=self.use_pp, uses_context_parallel=self.dcp_size > 1 @@ -599,13 +601,14 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE - elif ( - self.draft_trimmer is not None - and self.compilation_config.cudagraph_mode is not None - and self.compilation_config.cudagraph_mode.has_full_cudagraphs() + elif self.draft_trimmer is not None and ( + self.compilation_config.cudagraph_mode is None + or self.compilation_config.cudagraph_mode.has_full_cudagraphs() ): # Trimmed decode batches have per-request varlen queries, which - # uniform-decode full graphs cannot replay. + # uniform-decode full graphs cannot replay. The trimmer is only + # created with full cudagraphs when the attention backend + # supports varlen decode capture (AttentionCGSupport.ALWAYS). self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 7bbcf049bb00..6915a91cd8ee 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -210,6 +210,7 @@ def maybe_create_draft_trimmer( ) return None + logger.info("GPU draft trimming enabled for variable-length drafts.") return VariableDraftTrimmer( num_valid_drafts, query_start_loc, From b8a3834cdbc5787dc18161160cb6ba144025f4e2 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 14 Aug 2026 21:34:51 -0700 Subject: [PATCH 27/34] simplify Signed-off-by: Nick Hill --- .../gpu/spec_decode/ngram/speculator.py | 31 ++----------------- vllm/v1/worker/gpu/spec_decode/speculator.py | 6 ++-- 2 files changed, 5 insertions(+), 32 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 25fae8dfc1ed..c59a7de5aca6 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -5,12 +5,11 @@ from typing import TYPE_CHECKING, Any import torch -import torch.nn as nn from vllm.config import VllmConfig -from vllm.config.compilation import CUDAGraphMode from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator if TYPE_CHECKING: from vllm.v1.worker.gpu.states import RequestState @@ -151,16 +150,8 @@ def _ngram_finalize_kernel( tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) -class NgramGPUSpeculator: - """V2-compatible GPU n-gram speculator. - - Drafts are proposed entirely on GPU by scanning each request's token - history (`RequestState.all_token_ids`) in place for the longest n-gram - suffix match. The per-request count of usable drafts is written to the - persistent `num_valid_drafts` tensor, which the model runner's draft - trimmer consumes at the next step to shrink verification on GPU without - any CPU<->GPU synchronization. - """ +class NgramGPUSpeculator(BaseSpeculator): + """V2-compatible GPU n-gram speculator.""" supports_mm_inputs = False draft_logits = None @@ -221,22 +212,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.req_states: RequestState | None = None - def load_model(self, target_model: nn.Module) -> None: - """No weights to load — ngram is a data-only proposer.""" - pass - - def set_attn(self, *args: Any, **kwargs: Any) -> None: - """No attention layers owned by this speculator.""" - pass - - def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - """N-gram kernels are launched directly; no explicit CG capture.""" - pass - - def capture(self, *args: Any, **kwargs: Any) -> None: - """No graph capture phase required.""" - pass - @torch.inference_mode() def propose( self, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 69e0cc109160..64f6f6b865c6 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -30,13 +30,11 @@ class BaseSpeculator(ABC): - @abstractmethod def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - pass + return None - @abstractmethod def capture(self) -> None: - pass + return None @abstractmethod def propose( From e052b8c9b3a64e1b43b5b19a34b54cf000ae2127 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 15 Aug 2026 14:36:36 -0700 Subject: [PATCH 28/34] [Spec Decode] Resolve both ngram methods to the GPU implementation on the V2 runner On the V2 model runner, method="ngram" and method="ngram_gpu" now both use NgramGPUSpeculator (via a new SpeculativeConfig.use_ngram() helper); the V1 runner keeps its separate CPU and GPU proposers. Also scope the torch.compile cache disable to the V1 ngram-gpu proposer: it exists for V1's @support_torch_compile kernel, while the V2 implementation is pure Triton and does not need it. Co-authored-by: Claude Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_max_len.py | 8 +++++--- tests/v1/spec_decode/test_ngram_gpu.py | 18 ++++++++++++++---- vllm/compilation/backends.py | 5 ++++- vllm/config/speculative.py | 3 +++ vllm/config/vllm.py | 8 ++++---- vllm/v1/worker/gpu/model_runner.py | 3 +-- vllm/v1/worker/gpu/spec_decode/__init__.py | 2 +- 7 files changed, 32 insertions(+), 15 deletions(-) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 91d53b5b57e3..a20d5caecfc0 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -37,10 +37,12 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) -def test_ngram_gpu_max_len(num_speculative_tokens: int, vllm_runner): +@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) +def test_ngram_gpu_max_len(method: str, num_speculative_tokens: int, vllm_runner): """V2 GPU n-gram counterpart of ``test_ngram_max_len``. - Verifies that the V2 model runner with ``method="ngram_gpu"`` correctly + Verifies that the V2 model runner (where "ngram" and "ngram_gpu" both + resolve to the GPU implementation) correctly handles the ``max_model_len`` boundary across various speculative-token counts. """ @@ -51,7 +53,7 @@ def test_ngram_gpu_max_len(num_speculative_tokens: int, vllm_runner): enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ - "method": "ngram_gpu", + "method": method, "prompt_lookup_max": 5, "prompt_lookup_min": 3, "num_speculative_tokens": num_speculative_tokens, diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 470caf9618f9..ce55fcfa3d2b 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -49,6 +49,7 @@ def _make_vllm_config( k: int, max_num_seqs: int = 8, max_model_len: int = 64, + method: str = "ngram_gpu", ) -> VllmConfig: model_config = ModelConfig( model="facebook/opt-125m", @@ -60,7 +61,7 @@ def _make_vllm_config( max_model_len=max_model_len, ) speculative_config = SpeculativeConfig( - method="ngram_gpu", + method=method, prompt_lookup_min=min_n, prompt_lookup_max=max_n, num_speculative_tokens=k, @@ -360,14 +361,23 @@ def test_propose_requires_req_states(): ) +@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) +def test_both_ngram_methods_resolve_to_gpu_speculator(method: str): + """On the V2 runner, "ngram" and "ngram_gpu" use the same implementation.""" + from vllm.v1.worker.gpu.spec_decode import init_speculator + + cfg = _make_vllm_config(min_n=2, max_n=3, k=2, method=method) + spec = init_speculator(cfg, DEVICE) + assert isinstance(spec, NgramGPUSpeculator) + + def test_construction_validates_speculative_config(): spec = _make_speculator(min_n=2, max_n=3, k=2) assert spec.min_n == 2 assert spec.max_n == 3 assert spec.num_speculative_steps == 2 - # No-op hooks must not raise. - spec.load_model(target_model=None) - spec.set_attn() + # Inherited no-op hooks must not raise. + spec.init_cudagraph_manager(None) spec.capture() diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 8fbf5b41f747..56b78cfab2ac 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1081,10 +1081,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE. disable_cache = not is_compile_cache_enabled(self.inductor_config) - # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors. + # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm + # torch compile cache errors. The V2 implementation is pure Triton and + # does not need the cache disabled. is_ngram_gpu_enabled = ( vllm_config.speculative_config is not None and vllm_config.speculative_config.use_ngram_gpu() + and not vllm_config.use_v2_model_runner ) disable_cache = disable_cache or is_ngram_gpu_enabled diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 2b93113b7ed3..0e09566aaccc 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1498,6 +1498,9 @@ def uses_extract_hidden_states(self) -> bool: def use_ngram_gpu(self) -> bool: return self.method == "ngram_gpu" + def use_ngram(self) -> bool: + return self.method in ("ngram", "ngram_gpu") + def use_multi_module_mtp(self) -> bool: if self.method != "mtp" or self.draft_model_config is None: return False diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index aa401fbbbe92..ffa852a31c32 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2395,15 +2395,15 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported.append("pipeline parallelism with external_launcher") if speculative_config is not None: - # CPU ngram is not supported by the V2 model runner yet. - if speculative_config.method == "ngram": - unsupported.append("ngram speculative decoding") - elif speculative_config.method not in ( + # Both ngram methods resolve to the same GPU implementation on + # the V2 model runner. + if speculative_config.method not in ( "eagle", "eagle3", "mtp", "dflash", "dspark", + "ngram", "ngram_gpu", ): unsupported.append(f"speculative method '{speculative_config.method}'") diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f6be5e20202a..1b4031e2c4b8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -271,8 +271,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # General request states. use_dense_all_token_ids = ( - self.speculative_config is not None - and self.speculative_config.use_ngram_gpu() + self.speculative_config is not None and self.speculative_config.use_ngram() ) self.req_states = RequestState( max_num_reqs=self.max_num_reqs, diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index b84f30bfcc48..4ae9bb70b569 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -42,7 +42,7 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): ) return EagleSpeculator(vllm_config, device) - elif speculative_config.use_ngram_gpu(): + elif speculative_config.use_ngram(): from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( NgramGPUSpeculator, ) From 6c5dfc956872079608bcd2987dbead62be9925aa Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 17 Aug 2026 18:07:36 -0700 Subject: [PATCH 29/34] [Spec Decode] Resolve the varlen-decode cudagraph constraint inside resolve_cudagraph_mode_and_sizes Both adaptive verification and variable-length drafters decide per-request query lengths on device, so decode batches are varlen and cudagraph capture needs a separate decode routine. Express that as a varlen_decode flag on resolve_cudagraph_mode_and_sizes, alongside the other backend-support downgrades, instead of mutating compilation_config.cudagraph_mode from the model runner beforehand. Only CUDAGraphMode.FULL actually needs to change (it has full cudagraphs but no separate decode routine); PIECEWISE/NONE capture no full decode graphs and FULL_DECODE_ONLY/FULL_AND_PIECEWISE already have one. This drops the adaptive-verification override of an explicitly requested PIECEWISE or FULL_DECODE_ONLY mode, which was a no-op for the default FULL_AND_PIECEWISE. Co-authored-by: Claude Signed-off-by: Nick Hill --- tests/test_config.py | 32 ++++++++++++++++++++++++++++++ vllm/config/compilation.py | 18 +++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 16 +++------------ 3 files changed, 53 insertions(+), 13 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index a2797e52126a..6b5e445216a1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -157,6 +157,38 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes +@pytest.mark.parametrize( + ("requested", "expected"), + [ + # Only a mode without a separate decode routine has to change: varlen + # decode batches would otherwise replay on a mixed full graph. + (CUDAGraphMode.FULL, CUDAGraphMode.FULL_AND_PIECEWISE), + (CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE), + (CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_DECODE_ONLY), + (CUDAGraphMode.PIECEWISE, CUDAGraphMode.PIECEWISE), + (CUDAGraphMode.NONE, CUDAGraphMode.NONE), + ], +) +def test_resolve_cudagraph_mode_varlen_decode(requested, expected): + """varlen_decode requires a separate decode routine for full cudagraphs.""" + compilation_config = CompilationConfig( + cudagraph_mode=requested, + cudagraph_capture_sizes=[1, 2, 4, 8], + ) + compilation_config.max_cudagraph_capture_size = 8 + compilation_config.post_init_cudagraph_sizes() + + cudagraph_mode = compilation_config.resolve_cudagraph_mode_and_sizes( + AttentionCGSupport.ALWAYS, + "FakeAttentionBackend", + use_v2_model_runner=True, + varlen_decode=True, + ) + + assert cudagraph_mode == expected + assert compilation_config.cudagraph_mode == expected + + @pytest.mark.parametrize( ("model_config", "expected"), [ diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 3cd227d72ce4..351f8cce8e6c 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1376,6 +1376,7 @@ def resolve_cudagraph_mode_and_sizes( kv_cache_config: "KVCacheConfig | None" = None, max_num_reqs: int | None = None, is_profiling: bool = False, + varlen_decode: bool = False, ) -> CUDAGraphMode: from vllm.v1.attention.backend import AttentionCGSupport @@ -1384,6 +1385,23 @@ def resolve_cudagraph_mode_and_sizes( self.cudagraph_mode = CUDAGraphMode.NONE return CUDAGraphMode.NONE + # Decode batches whose per-request query lengths are decided on device + # (adaptive verification, variable-length drafters) are captured as + # varlen decode graphs, which requires a separate decode routine. + # Modes without one would replay such a batch on a mixed graph. + if ( + varlen_decode + and cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + logger.warning( + "CUDAGraphMode.%s cannot capture decode batches with varying " + "per-request query lengths; setting " + "cudagraph_mode=FULL_AND_PIECEWISE", + cudagraph_mode.name, + ) + cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + # Check cudagraph for mixed batch is supported if ( cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 00e0d2d953d6..4f682ba6d1ac 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -607,17 +607,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config ) - if self.adaptive_verification is not None: - self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE - elif self.draft_trimmer is not None and ( - self.compilation_config.cudagraph_mode is None - or self.compilation_config.cudagraph_mode.has_full_cudagraphs() - ): - # Trimmed decode batches have per-request varlen queries, which - # uniform-decode full graphs cannot replay. The trimmer is only - # created with full cudagraphs when the attention backend - # supports varlen decode capture (AttentionCGSupport.ALWAYS). - self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + varlen_decode = bool(self.adaptive_verification or self.draft_trimmer) cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, @@ -626,6 +616,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, + varlen_decode=varlen_decode, ) self.cudagraph_manager = ModelCudaGraphManager( self.vllm_config, @@ -633,8 +624,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cudagraph_mode, decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, - varlen_decode=self.adaptive_verification is not None - or self.draft_trimmer is not None, + varlen_decode=varlen_decode, ) check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): From bc5e16d7e922bd078e1698b9e5adaedad20ef933 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 17 Aug 2026 18:39:52 -0700 Subject: [PATCH 30/34] [Spec Decode] Pass RequestState to speculators at construction Replace the getattr/hasattr injection of RequestState into the speculator with an explicit init_speculator parameter, handed to the speculators that draft from the persistent token store (currently only NgramGPUSpeculator). RequestState is now built before the speculator, which only needs config values that were already available at that point. NgramGPUSpeculator.req_states is consequently non-optional, so propose() drops its injection assert, and its tests exercise a real RequestState instead of a duck-typed stand-in. Co-authored-by: Claude Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_ngram_gpu.py | 68 +++++-------------- vllm/v1/worker/gpu/model_runner.py | 60 ++++++++-------- vllm/v1/worker/gpu/spec_decode/__init__.py | 14 +++- .../gpu/spec_decode/ngram/speculator.py | 14 ++-- 4 files changed, 64 insertions(+), 92 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index ce55fcfa3d2b..d92bd7497f42 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -33,6 +33,7 @@ build_verification_layout, ) from vllm.v1.worker.gpu.spec_decode.ngram.speculator import NgramGPUSpeculator +from vllm.v1.worker.gpu.states import RequestState if not torch.cuda.is_available(): pytest.skip( @@ -73,19 +74,16 @@ def _make_vllm_config( ) -class _FakeStaged: - """Stand-in for ``StagedWriteTensor``; only ``.gpu`` is read by propose().""" - - def __init__(self, tensor: torch.Tensor): - self.gpu = tensor - - -class _FakeRequestState: - """Minimal duck-typed ``RequestState`` exposing what propose() reads.""" - - def __init__(self, all_token_ids: torch.Tensor, total_len: torch.Tensor): - self.all_token_ids = _FakeStaged(all_token_ids) - self.total_len = _FakeStaged(total_len) +def _make_request_state(cfg: VllmConfig) -> RequestState: + return RequestState( + max_num_reqs=cfg.scheduler_config.max_num_seqs, + max_model_len=cfg.model_config.max_model_len, + max_num_batched_tokens=cfg.scheduler_config.max_num_batched_tokens, + num_speculative_steps=cfg.speculative_config.num_speculative_tokens, + vocab_size=cfg.model_config.get_vocab_size(), + device=DEVICE, + use_dense_all_token_ids=True, + ) def _make_speculator( @@ -102,7 +100,7 @@ def _make_speculator( max_num_seqs=max_num_seqs, max_model_len=max_model_len, ) - return NgramGPUSpeculator(vllm_config=cfg, device=DEVICE) + return NgramGPUSpeculator(cfg, DEVICE, _make_request_state(cfg)) def _propose( @@ -128,9 +126,10 @@ def _propose( slots = list(range(B)) max_num_reqs = spec.max_num_reqs - L = spec.max_model_len - all_token_ids = torch.zeros((max_num_reqs, L), dtype=torch.int32, device=DEVICE) - total_len = torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + all_token_ids = spec.req_states.all_token_ids.gpu + total_len = spec.req_states.total_len.gpu + all_token_ids.zero_() + total_len.zero_() last_sampled_t = torch.zeros((max_num_reqs, 1), dtype=torch.int64, device=DEVICE) for row, slot, seq_len, last in zip(rows, slots, seq_lens, last_sampled): if row: @@ -140,7 +139,6 @@ def _propose( total_len[slot] = seq_len last_sampled_t[slot, 0] = last - spec.req_states = _FakeRequestState(all_token_ids, total_len) idx_mapping = torch.tensor(slots, dtype=torch.int64, device=DEVICE) input_batch = SimpleNamespace(num_reqs=B, idx_mapping=idx_mapping) @@ -337,40 +335,6 @@ def test_dummy_run_does_not_touch_state(): assert torch.equal(spec.num_valid_drafts.cpu(), before.cpu()) -def test_propose_requires_req_states(): - """propose() must assert that req_states was injected by the model runner.""" - spec = _make_speculator(min_n=2, max_n=2, k=2) - assert spec.req_states is None - input_batch = SimpleNamespace( - num_reqs=1, - idx_mapping=torch.tensor([0], dtype=torch.int64, device=DEVICE), - ) - with pytest.raises(AssertionError, match="req_states"): - spec.propose( - input_batch=input_batch, - attn_metadata=None, - slot_mappings=None, - last_hidden_states=torch.empty(0, device=DEVICE), - aux_hidden_states=None, - num_sampled=torch.ones(1, dtype=torch.int32, device=DEVICE), - num_rejected=torch.zeros(1, dtype=torch.int32, device=DEVICE), - last_sampled=torch.zeros((8, 1), dtype=torch.int64, device=DEVICE), - next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device=DEVICE), - temperature=torch.zeros(1, dtype=torch.float32, device=DEVICE), - seeds=torch.zeros(1, dtype=torch.int64, device=DEVICE), - ) - - -@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) -def test_both_ngram_methods_resolve_to_gpu_speculator(method: str): - """On the V2 runner, "ngram" and "ngram_gpu" use the same implementation.""" - from vllm.v1.worker.gpu.spec_decode import init_speculator - - cfg = _make_vllm_config(min_n=2, max_n=3, k=2, method=method) - spec = init_speculator(cfg, DEVICE) - assert isinstance(spec, NgramGPUSpeculator) - - def test_construction_validates_speculative_config(): spec = _make_speculator(min_n=2, max_n=3, k=2) assert spec.min_n == 2 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4f682ba6d1ac..33f5688bed8a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -238,31 +238,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.encoder_cache = EncoderCache() self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache) - # Speculative decoding. - self.speculator = None - self.use_aux_hidden_state_outputs = False self.num_speculative_steps = vllm_config.num_speculative_tokens - if self.speculative_config is not None: - if self.is_last_pp_rank: - self.speculator = init_speculator(self.vllm_config, self.device) - - if self.speculative_config.method in ("eagle3", "dflash", "dspark"): - # Drafting may require auxiliary hidden states from target model outputs - self.use_aux_hidden_state_outputs = True - if self.use_pp: - raise ValueError( - f"{self.speculative_config.method} with pipeline parallel " - "is not supported." - ) - - # Draft tokens propagation - for spec-dec + struct outputs. - self.draft_tokens_handler = DraftTokensHandler(self.device) - - self.pcp_manager: pcp.PCPManager | None = None - - # Pooling models. - self.is_pooling_model = self.model_config.runner_type == "pooling" - self.pooling_runner: PoolingRunner | None = None # Multi-module MTP feeds its modules the next num_speculative_steps prefill # tokens during chunked prefill. Other speculators only read the immediate @@ -274,8 +250,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): else 1 ) - self.step_timing = StepTimingCollector() - # General request states. use_dense_all_token_ids = ( self.speculative_config is not None and self.speculative_config.use_ngram() @@ -290,6 +264,35 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): num_prefill_lookahead=num_prefill_lookahead, use_dense_all_token_ids=use_dense_all_token_ids, ) + + # Speculative decoding. + self.speculator = None + self.use_aux_hidden_state_outputs = False + if self.speculative_config is not None: + if self.is_last_pp_rank: + self.speculator = init_speculator( + self.vllm_config, self.device, self.req_states + ) + + if self.speculative_config.method in ("eagle3", "dflash", "dspark"): + # Drafting may require auxiliary hidden states from target model outputs + self.use_aux_hidden_state_outputs = True + if self.use_pp: + raise ValueError( + f"{self.speculative_config.method} with pipeline parallel " + "is not supported." + ) + + # Draft tokens propagation - for spec-dec + struct outputs. + self.draft_tokens_handler = DraftTokensHandler(self.device) + + self.pcp_manager: pcp.PCPManager | None = None + + # Pooling models. + self.is_pooling_model = self.model_config.runner_type == "pooling" + self.pooling_runner: PoolingRunner | None = None + + self.step_timing = StepTimingCollector() self.adaptive_verification: AdaptiveVerificationManager | None = None self.draft_trimmer: VariableDraftTrimmer | None = None self.input_buffers = InputBuffers( @@ -304,11 +307,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=self.device, ) - # Inject RequestState into speculators that consume the persistent - # token store directly (e.g. NgramGPUSpeculator). - if self.speculator is not None and hasattr(self.speculator, "req_states"): - self.speculator.req_states = self.req_states - # Samplers and decode_query_len created in load_model() after # model_state exists (num_new_sampled_tokens_per_step from ModelState). self.sampler: Sampler | None = None diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 4ae9bb70b569..1b1d7b341b50 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,11 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import torch from vllm.config import VllmConfig +if TYPE_CHECKING: + from vllm.v1.worker.gpu.states import RequestState + -def init_speculator(vllm_config: VllmConfig, device: torch.device): +def init_speculator( + vllm_config: VllmConfig, + device: torch.device, + req_states: "RequestState", +): + """Build the speculator for this config.""" speculative_config = vllm_config.speculative_config assert speculative_config is not None if speculative_config.method == "dflash": @@ -47,6 +57,6 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): NgramGPUSpeculator, ) - return NgramGPUSpeculator(vllm_config, device) + return NgramGPUSpeculator(vllm_config, device, req_states) else: raise NotImplementedError(f"{speculative_config.method} is not supported yet.") diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index c59a7de5aca6..c3dcdffa92ae 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -159,7 +159,12 @@ class NgramGPUSpeculator(BaseSpeculator): # for GPU-side verification trimming. trims_drafts_on_gpu = True - def __init__(self, vllm_config: VllmConfig, device: torch.device): + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + req_states: RequestState, + ): if not HAS_TRITON: raise RuntimeError("ngram_gpu speculative decoding requires Triton.") spec = vllm_config.speculative_config @@ -174,6 +179,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.device = device + self.req_states = req_states self.speculative_config = spec self.num_speculative_steps: int = spec.num_speculative_tokens @@ -210,8 +216,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=device, ) - self.req_states: RequestState | None = None - @torch.inference_mode() def propose( self, @@ -238,10 +242,6 @@ def propose( return self.drafts[:num_reqs] req_states = self.req_states - assert req_states is not None, ( - "NgramGPUSpeculator.req_states was not injected by the model runner." - ) - token_ids = req_states.all_token_ids.gpu idx_mapping = input_batch.idx_mapping From 2e39f8c81c8218e49fd84653e530ad8f4b13938e Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 17 Aug 2026 19:00:21 -0700 Subject: [PATCH 31/34] [Spec Decode] Derive draft-trimmer prerequisites inside maybe_create_draft_trimmer Pass vllm_config and RequestState instead of eight individually-derived values: LoRA/PP/CP/cudagraph-mode support all come from the config, and max_num_reqs, device and the logit chunk limit from RequestState. Declare trims_drafts_on_gpu and num_valid_drafts on BaseSpeculator so the factory reads them directly rather than through getattr, which also lets it accept a None speculator and drop that check from the call site. Co-authored-by: Claude Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/model_runner.py | 21 +++------- .../gpu/spec_decode/adaptive_verification.py | 39 +++++++++++-------- .../gpu/spec_decode/ngram/speculator.py | 2 - vllm/v1/worker/gpu/spec_decode/speculator.py | 6 +++ 4 files changed, 33 insertions(+), 35 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 33f5688bed8a..6c14f7bc09e7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -563,24 +563,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: # Variable-length drafters (ngram_gpu) trim scheduled draft slots to # the drafter's valid counts on GPU, when supported. self.draft_trimmer = None - if self.speculator is not None and self.adaptive_verification is None: + if self.adaptive_verification is None: self.draft_trimmer = maybe_create_draft_trimmer( + vllm_config=self.vllm_config, speculator=self.speculator, attn_groups=self.attn_groups, attn_cg_support=attn_cg_support, - # Unset (auto) mode may resolve to full graphs, so treat it - # as full-graph usage. - uses_full_cudagraphs=self.compilation_config.cudagraph_mode is None - or self.compilation_config.cudagraph_mode.has_full_cudagraphs(), - has_lora=self.lora_config is not None, - uses_pipeline_parallel=self.use_pp, - uses_context_parallel=self.dcp_size > 1 - or self.parallel_config.prefill_context_parallel_size > 1, + req_states=self.req_states, query_start_loc=self.input_buffers.query_start_loc, num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, - max_total_logits=get_max_chunk_logits(self.vocab_size), - max_num_reqs=self.max_num_reqs, - device=self.device, ) self.block_tables = BlockTables( @@ -1904,11 +1895,9 @@ def sample_tokens( ) if self.num_speculative_steps > 0: - # Spec-decode and diffusion LLMs both use draft tokens but the latter does - # not have a speculator (i.e. self.speculator is None) + # Spec-decode and diffusion LLMs both use draft tokens. self.draft_tokens_handler.set_draft_tokens( - input_batch, - self.req_states.draft_tokens[input_batch.idx_mapping], + input_batch, self.req_states.draft_tokens[input_batch.idx_mapping] ) # Post-step KV connector related operations. diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 6915a91cd8ee..a779d35e370a 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -25,8 +25,10 @@ _PROFILE_REPLAYS = 5 if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup @@ -165,41 +167,44 @@ def trim( def maybe_create_draft_trimmer( *, - speculator, + vllm_config: "VllmConfig", + speculator: "BaseSpeculator | None", attn_groups: list[list["AttentionGroup"]], attn_cg_support: "AttentionCGSupportInfo", - uses_full_cudagraphs: bool, - has_lora: bool, - uses_pipeline_parallel: bool, - uses_context_parallel: bool, + req_states: "RequestState", query_start_loc: torch.Tensor, num_bonus_tokens: int, - max_total_logits: int, - max_num_reqs: int, - device: torch.device, ) -> VariableDraftTrimmer | None: """Create a VariableDraftTrimmer when the drafter and environment support GPU-side trimming; otherwise fall back (with a log) to verifying the full padded drafts, which is correct but wastes verification compute.""" - num_valid_drafts = getattr(speculator, "num_valid_drafts", None) - if not getattr(speculator, "trims_drafts_on_gpu", False): + from vllm.v1.worker.gpu.spec_decode.rejection_sampler import get_max_chunk_logits + + if speculator is None or not speculator.trims_drafts_on_gpu: return None + num_valid_drafts = speculator.num_valid_drafts assert num_valid_drafts is not None + parallel_config = vllm_config.parallel_config + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + reason = None backend = get_query_lens_mismatch_unsupported_backend(attn_groups) if backend is not None: reason = f"the {backend} attention backend" elif ( - uses_full_cudagraphs + cudagraph_mode.has_full_cudagraphs() and attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS ): reason = f"varlen decode cudagraphs with {attn_cg_support.min_cg_attn_backend}" - elif has_lora: + elif vllm_config.lora_config is not None: reason = "LoRA" - elif uses_pipeline_parallel: + elif parallel_config.pipeline_parallel_size > 1: reason = "pipeline parallelism" - elif uses_context_parallel: + elif ( + parallel_config.decode_context_parallel_size > 1 + or parallel_config.prefill_context_parallel_size > 1 + ): reason = "context parallelism" if reason is not None: @@ -215,9 +220,9 @@ def maybe_create_draft_trimmer( num_valid_drafts, query_start_loc, num_bonus_tokens, - max_num_reqs, - max_total_logits, - device, + req_states.max_num_reqs, + get_max_chunk_logits(req_states.vocab_size), + req_states.device, ) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index c3dcdffa92ae..f2c361fc193a 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -155,8 +155,6 @@ class NgramGPUSpeculator(BaseSpeculator): supports_mm_inputs = False draft_logits = None - # Signals that num_valid_drafts holds per-request valid draft counts - # for GPU-side verification trimming. trims_drafts_on_gpu = True def __init__( diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 8cd7454ad7c0..d9978e27f895 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -30,6 +30,12 @@ class BaseSpeculator(ABC): + # Variable-length drafters set this and publish per-request counts of + # usable drafts in num_valid_drafts, which the model runner's draft + # trimmer reads to shrink verification on device. + trims_drafts_on_gpu: bool = False + num_valid_drafts: torch.Tensor | None = None + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: return None From 8797fe6e73ce1d9fa9f21295d9c4f35c5f592aa5 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 17 Aug 2026 19:17:33 -0700 Subject: [PATCH 32/34] [Spec Decode] Collapse the draft-trim contract into one speculator field trims_drafts_on_gpu carried no information beyond "num_valid_drafts is set", so replace both with a single optional num_valid_drafts_for_trim tensor on BaseSpeculator: None means verify every scheduled draft, a tensor opts the drafter into device-side trimming. Co-authored-by: Claude Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_ngram_gpu.py | 10 +++++----- .../v1/worker/gpu/spec_decode/adaptive_verification.py | 8 +++----- vllm/v1/worker/gpu/spec_decode/ngram/speculator.py | 5 ++--- vllm/v1/worker/gpu/spec_decode/speculator.py | 10 +++++----- 4 files changed, 15 insertions(+), 18 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index d92bd7497f42..95303e6c9b9d 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -155,7 +155,7 @@ def _propose( temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), ) - num_valid = spec.num_valid_drafts[idx_mapping] + num_valid = spec.num_valid_drafts_for_trim[idx_mapping] return drafts.cpu().tolist(), num_valid.cpu().tolist() @@ -295,14 +295,14 @@ def test_noncontiguous_idx_mapping(): def test_num_valid_written_to_request_slots(): - """num_valid_drafts is req-slot indexed for the GPU draft trimmer.""" + """num_valid_drafts_for_trim is req-slot indexed for the draft trimmer.""" spec = _make_speculator(min_n=2, max_n=2, k=2) _propose( spec, [[7, 8, 9, 7, 8], [1, 2, 3, 4, 5]], slots=[5, 2], ) - nv = spec.num_valid_drafts.cpu() + nv = spec.num_valid_drafts_for_trim.cpu() assert nv[5].item() == 2 # match assert nv[2].item() == 0 # no match @@ -311,7 +311,7 @@ def test_dummy_run_does_not_touch_state(): """Dummy runs must not mutate persistent request or drafter state.""" spec = _make_speculator(min_n=2, max_n=2, k=2) _propose(spec, [[1, 2, 3, 1, 2]], slots=[1]) - before = spec.num_valid_drafts.clone() + before = spec.num_valid_drafts_for_trim.clone() input_batch = SimpleNamespace( num_reqs=1, @@ -332,7 +332,7 @@ def test_dummy_run_does_not_touch_state(): dummy_run=True, ) assert drafts.shape == (1, 2) - assert torch.equal(spec.num_valid_drafts.cpu(), before.cpu()) + assert torch.equal(spec.num_valid_drafts_for_trim.cpu(), before.cpu()) def test_construction_validates_speculative_config(): diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index a779d35e370a..c8fab9b555f2 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -107,7 +107,7 @@ class VariableDraftTrimmer: """GPU-side verification trimming for variable-length drafters (ngram). The drafter records per-request valid draft counts on GPU in - `num_valid_drafts`. The scheduler still schedules the full + `num_valid_drafts_for_trim`. The scheduler still schedules the full num_speculative_tokens per request; at the next step this trimmer clamps each request's scheduled draft slots to the recorded count and rebuilds cu_num_logits / query_start_loc on device, so the CPU keeps only upper @@ -180,10 +180,8 @@ def maybe_create_draft_trimmer( padded drafts, which is correct but wastes verification compute.""" from vllm.v1.worker.gpu.spec_decode.rejection_sampler import get_max_chunk_logits - if speculator is None or not speculator.trims_drafts_on_gpu: + if speculator is None or speculator.num_valid_drafts_for_trim is None: return None - num_valid_drafts = speculator.num_valid_drafts - assert num_valid_drafts is not None parallel_config = vllm_config.parallel_config cudagraph_mode = vllm_config.compilation_config.cudagraph_mode @@ -217,7 +215,7 @@ def maybe_create_draft_trimmer( logger.info("GPU draft trimming enabled for variable-length drafts.") return VariableDraftTrimmer( - num_valid_drafts, + speculator.num_valid_drafts_for_trim, query_start_loc, num_bonus_tokens, req_states.max_num_reqs, diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index f2c361fc193a..1ab5ba15b95c 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -155,7 +155,6 @@ class NgramGPUSpeculator(BaseSpeculator): supports_mm_inputs = False draft_logits = None - trims_drafts_on_gpu = True def __init__( self, @@ -203,7 +202,7 @@ def __init__( ) # Per request-slot count of usable drafts from the latest proposal, # consumed by the model runner's GPU draft trimmer. - self.num_valid_drafts = torch.zeros( + self.num_valid_drafts_for_trim = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=device ) # Batch-ordered draft output, scattered into RequestState.draft_tokens @@ -270,7 +269,7 @@ def propose( self.scratch, self.scratch.stride(0), self.drafts, - self.num_valid_drafts, + self.num_valid_drafts_for_trim, self.max_model_len, self.n_blocks, self.num_speculative_steps, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d9978e27f895..4f6725e005b8 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -30,11 +30,11 @@ class BaseSpeculator(ABC): - # Variable-length drafters set this and publish per-request counts of - # usable drafts in num_valid_drafts, which the model runner's draft - # trimmer reads to shrink verification on device. - trims_drafts_on_gpu: bool = False - num_valid_drafts: torch.Tensor | None = None + # Variable-length drafters publish per-request counts of usable drafts + # here, [max_num_reqs] int32 indexed by request slot. Leaving it None + # means every scheduled draft is verified; setting it opts the drafter + # into device-side trimming by the model runner's draft trimmer. + num_valid_drafts_for_trim: torch.Tensor | None = None def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: return None From ac608bb61f31fa6eee3923e1c8502f2904dd0426 Mon Sep 17 00:00:00 2001 From: Xuanan Chen Date: Tue, 15 Sep 2026 01:38:54 -0700 Subject: [PATCH 33/34] bug fix Signed-off-by: Xuanan Chen --- tests/test_config.py | 5 +- tests/v1/spec_decode/test_max_len.py | 24 ++++--- tests/v1/spec_decode/test_ngram_gpu.py | 62 +++++++++++++++---- tests/v1/worker/test_gpu_model_runner_v2.py | 20 ++++++ .../test_gpu_rejection_sampler_chunking.py | 19 +++--- vllm/config/vllm.py | 3 - vllm/v1/worker/gpu/model_runner.py | 36 ++++++----- .../gpu/spec_decode/adaptive_verification.py | 4 ++ .../gpu/spec_decode/ngram/speculator.py | 3 +- .../gpu/spec_decode/rejection_sampler.py | 19 +++--- 10 files changed, 134 insertions(+), 61 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 4d39b56d265f..e3dfa303c3a8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -538,12 +538,13 @@ def test_dsa_models_select_matching_mtp(model_type, expected_architecture): assert hf_config.architectures == [expected_architecture] -def test_v2_model_runner_supports_extract_hidden_states(): +@pytest.mark.parametrize("method", ["extract_hidden_states", "ngram", "ngram_gpu"]) +def test_v2_model_runner_supports_speculative_method(method): config = VllmConfig() config.speculative_config = cast( SpeculativeConfig, SimpleNamespace( - method="extract_hidden_states", + method=method, parallel_drafting=False, enable_adaptive_verification=False, ), diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index a20d5caecfc0..17dd318563b7 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -38,14 +38,14 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) -def test_ngram_gpu_max_len(method: str, num_speculative_tokens: int, vllm_runner): - """V2 GPU n-gram counterpart of ``test_ngram_max_len``. - - Verifies that the V2 model runner (where "ngram" and "ngram_gpu" both - resolve to the GPU implementation) correctly - handles the ``max_model_len`` boundary across various speculative-token - counts. - """ +def test_ngram_gpu_max_len( + method: str, + num_speculative_tokens: int, + vllm_runner, + monkeypatch: pytest.MonkeyPatch, +): + """V2 n-gram decoding stops at max_model_len.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") with vllm_runner( "facebook/opt-125m", trust_remote_code=False, @@ -59,8 +59,14 @@ def test_ngram_gpu_max_len(method: str, num_speculative_tokens: int, vllm_runner "num_speculative_tokens": num_speculative_tokens, }, ) as runner: + assert runner.llm.llm_engine.vllm_config.use_v2_model_runner sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) - runner.llm.generate(_PROMPTS, sampling_params) + outputs = runner.llm.generate(_PROMPTS, sampling_params) + for output in outputs: + assert output.prompt_token_ids is not None + assert ( + len(output.prompt_token_ids) + len(output.outputs[0].token_ids) == 100 + ) @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 95303e6c9b9d..69393c403f98 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -28,9 +28,11 @@ SpeculativeConfig, VllmConfig, ) +from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( VariableDraftTrimmer, build_verification_layout, + maybe_create_draft_trimmer, ) from vllm.v1.worker.gpu.spec_decode.ngram.speculator import NgramGPUSpeculator from vllm.v1.worker.gpu.states import RequestState @@ -154,6 +156,7 @@ def _propose( next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), + dp_sync=None, ) num_valid = spec.num_valid_drafts_for_trim[idx_mapping] return drafts.cpu().tolist(), num_valid.cpu().tolist() @@ -164,9 +167,11 @@ def _propose( # --------------------------------------------------------------------------- -def test_no_match_returns_zero_valid(): - """No 2-gram match in [1,2,3,4,5] → 0 valid drafts, last_sampled fill.""" - spec = _make_speculator(min_n=2, max_n=2, k=2) +@pytest.mark.parametrize("max_model_len", [32, 300]) +def test_no_match_clears_previous_proposal(max_model_len): + spec = _make_speculator(min_n=2, max_n=2, k=2, max_model_len=max_model_len) + row = [0] * (max_model_len - 5) + [1, 2, 3, 1, 2] + assert _propose(spec, [row]) == ([[3, 1]], [2]) drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 5]], last_sampled=[42]) assert num_valid == [0] assert drafts == [[42, 42]] @@ -189,26 +194,31 @@ def test_falls_back_to_3gram_when_4gram_missing(): def test_prefers_longer_ngram(): - """Both a 4-gram and a 3-gram match exist → prefer the 4-gram match.""" + """Prefer a 4-gram match over a more recent 3-gram match.""" spec = _make_speculator(min_n=3, max_n=4, k=2) - drafts, num_valid = _propose(spec, [[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]]) + drafts, num_valid = _propose( + spec, [[1, 2, 3, 4, 50, 51, 2, 3, 4, 60, 61, 1, 2, 3, 4]] + ) assert num_valid == [2] - assert drafts == [[1, 2]] + assert drafts == [[50, 51]] def test_picks_longest_match_among_2_3_4_grams(): - """2-gram and 3-gram match, 4-gram does not → propose 3-gram match [1, 2].""" + """Prefer a 3-gram match over a more recent 2-gram match.""" spec = _make_speculator(min_n=2, max_n=4, k=2) - drafts, num_valid = _propose(spec, [[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]]) + drafts, num_valid = _propose(spec, [[2, 3, 4, 50, 51, 3, 4, 60, 61, 1, 2, 3, 4]]) assert num_valid == [2] - assert drafts == [[1, 2]] + assert drafts == [[50, 51]] -def test_picks_rightmost_when_multiple_matches(): - """Multiple 3-gram matches for suffix (1,2,3) → pick the right-most.""" - spec = _make_speculator(min_n=3, max_n=3, k=2) +@pytest.mark.parametrize("max_model_len", [32, 128, 257, 1025]) +def test_picks_rightmost_when_multiple_matches(max_model_len): + """Pick the last valid match across blocks, ignoring trailing tokens.""" + spec = _make_speculator(min_n=3, max_n=3, k=2, max_model_len=max_model_len) + padding = [0] * (spec.block_l - 5) if spec.n_blocks > 1 else [] + row = [1, 2, 3, 100] + padding + [1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3] drafts, num_valid = _propose( - spec, [[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]] + spec, [row + [1, 2, 3, 999, 1, 2, 3]], seq_lens=[len(row)] ) assert num_valid == [2] assert drafts == [[300, 1]] @@ -329,6 +339,7 @@ def test_dummy_run_does_not_touch_state(): next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device=DEVICE), temperature=torch.zeros(1, dtype=torch.float32, device=DEVICE), seeds=torch.zeros(1, dtype=torch.int64, device=DEVICE), + dp_sync=None, dummy_run=True, ) assert drafts.shape == (1, 2) @@ -399,3 +410,28 @@ def test_variable_draft_trimmer_clamps_to_num_valid(): assert qsl.cpu().tolist()[:4] == [0, 2, 5, 12] # Padding tail equals the (GPU) batch total. assert qsl.cpu().tolist()[4:] == [12] * (max_num_reqs - 3) + + +@pytest.mark.parametrize("batch_sharded_sampling", [False, True]) +def test_draft_trimmer_disabled_with_batch_sharded_sampling(batch_sharded_sampling): + """The sharder plans from CPU logits boundaries, so GPU trimming must stay off.""" + cfg = _make_vllm_config(min_n=2, max_n=2, k=2) + cfg.parallel_config.enable_batch_sharded_sampling = batch_sharded_sampling + backend = SimpleNamespace( + __name__="FakeBackend", + supports_device_cpu_query_lens_mismatch=lambda: True, + ) + trimmer = maybe_create_draft_trimmer( + vllm_config=cfg, + speculator=SimpleNamespace( + num_valid_drafts_for_trim=torch.zeros(8, dtype=torch.int32, device=DEVICE) + ), + attn_groups=[[SimpleNamespace(backend=backend, layer_names=set())]], + attn_cg_support=AttentionCGSupportInfo(), + req_states=SimpleNamespace(max_num_reqs=8, vocab_size=32, device=DEVICE), + query_start_loc=torch.empty(9, dtype=torch.int32, device=DEVICE), + num_bonus_tokens=1, + ) + assert (trimmer is None) == batch_sharded_sampling + if trimmer is not None: + assert isinstance(trimmer, VariableDraftTrimmer) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 86ff1a074538..c029db73983c 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -298,3 +298,23 @@ def test_capture_model_profile_only_skips_lock(monkeypatch): runner.capture_model(profile_only=True) assert lock_calls == [] + + +@pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"]) +def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer): + """Targets allocate the MTP hidden buffer only for hidden-state drafters.""" + runner = GPUModelRunner.__new__(GPUModelRunner) + hidden_states = torch.zeros(4, 8) + buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8) + if target_buffer == "absent": + runner.model = SimpleNamespace() + else: + returned = buffer if target_buffer == "tensor" else None + runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) + + out = runner._get_drafter_hidden_states(hidden_states) + + if target_buffer == "tensor": + assert torch.equal(out, buffer[:4]) + else: + assert out is hidden_states diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py index 2e5adfabd13c..3ab54be3db6f 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -28,15 +28,19 @@ def test_iter_request_chunks_preserves_request_boundaries(): @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) -def test_chunked_scores_match_full_batch(logprobs_mode: str): +@pytest.mark.parametrize("trim_drafts", [False, True]) +def test_chunked_scores_match_full_batch(logprobs_mode: str, trim_drafts: bool): device = torch.device("cuda") cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32) - num_logits_per_req = np.diff(cu_num_logits_np) + expected_offsets = cu_num_logits_np.copy() + if trim_drafts: + expected_offsets[1:] = [1, 2, 5, 6] + num_logits_per_req = np.diff(expected_offsets) idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32) input_batch = SimpleNamespace( num_reqs=4, cu_num_logits_np=cu_num_logits_np, - cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device), + cu_num_logits=torch.from_numpy(expected_offsets).to(device), idx_mapping_np=idx_mapping_np, idx_mapping=torch.from_numpy(idx_mapping_np).to(device), expanded_idx_mapping=torch.from_numpy( @@ -78,7 +82,7 @@ def fake_verify( draft_logits=None, draft_sampled=torch.arange(10, device=device), pos=torch.arange(10, device=device), - max_chunk_logits=5, + max_chunk_logits=10 if trim_drafts else 5, max_num_logprobs=2, ) score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits @@ -87,7 +91,7 @@ def fake_verify( num_sampled, score_logits, input_batch.cu_num_logits, - input_batch.cu_num_logits_np, + expected_offsets, max_num_logprobs=2, ) @@ -105,6 +109,7 @@ def fake_verify( full_logprobs.selected_token_ranks, ) assert ( - chunked_logprobs.cu_num_generated_tokens - == full_logprobs.cu_num_generated_tokens + chunked_logprobs.tolists().cu_num_generated_tokens + == full_logprobs.tolists().cu_num_generated_tokens + == expected_offsets.tolist() ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ad8f077148c3..919506ba20c3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2830,9 +2830,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: if speculative_config is not None: if speculative_config.method in ( - # https://github.com/vllm-project/vllm/pull/40704 - "ngram", - "ngram_gpu", # https://github.com/vllm-project/vllm/pull/43091 "draft_model", "suffix", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8779793f81fb..b9157b850624 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -903,14 +903,7 @@ def _dummy_run( ), ) - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = self._get_drafter_hidden_states(hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), @@ -2041,6 +2034,24 @@ def execute_model( ) return None + def _get_drafter_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Hidden states fed to the drafter. + + Targets such as DeepSeek V4 expose the pre-hc_head residual through + get_mtp_target_hidden_states(). The buffer is sized at + max_num_batched_tokens and only allocated for drafters that consume + target hidden states, so None means "use the regular hidden states". + """ + get_target_hidden_states = getattr( + self.model, "get_mtp_target_hidden_states", None + ) + if get_target_hidden_states is None: + return hidden_states + target_hidden_states = get_target_hidden_states() + if target_hidden_states is None: + return hidden_states + return target_hidden_states[: hidden_states.shape[0]] + @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2175,14 +2186,7 @@ def sample_tokens( self.speculator.observe_verification( input_batch.idx_mapping, num_sampled, num_rejected ) - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = draft_hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] + spec_hidden_states = self._get_drafter_hidden_states(draft_hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index f675fd7cc597..a5a5f5df7013 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -209,6 +209,10 @@ def maybe_create_draft_trimmer( reason = f"varlen decode cudagraphs with {attn_cg_support.min_cg_attn_backend}" elif vllm_config.lora_config is not None: reason = "LoRA" + elif parallel_config.enable_batch_sharded_sampling: + # The sharder plans its all-to-all splits and local buffers from the + # CPU (untrimmed) logits boundaries. + reason = "batch-sharded sampling" elif parallel_config.pipeline_parallel_size > 1: reason = "pipeline parallelism" elif ( diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 1ab5ba15b95c..3fe3c392bd8a 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -12,6 +12,7 @@ from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator if TYPE_CHECKING: + from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.states import RequestState @@ -227,7 +228,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 609924e87fd2..3d20ebd21573 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -110,7 +110,7 @@ def _get_logprobs_tensors( num_sampled: torch.Tensor, logits: torch.Tensor, cu_num_logits: torch.Tensor, - cu_num_logits_np: np.ndarray, + cu_num_logits_np: np.ndarray | None, max_num_logprobs: int, ) -> LogprobsTensors | None: if max_num_logprobs == NO_LOGPROBS: @@ -132,10 +132,7 @@ def _get_logprobs_tensors( expanded_logits = num_logits != num_reqs cu_num_generated_tokens: list[int] | torch.Tensor | None = None if expanded_logits: - if self.enable_adaptive_verification: - # Adaptive verification keeps the true per-request boundaries - # on device only; cu_num_logits_np holds the pre-compacted - # layout. + if cu_num_logits_np is None: cu_num_generated_tokens = cu_num_logits.clone() else: cu_num_generated_tokens = cu_num_logits_np.tolist() @@ -222,10 +219,8 @@ def _verify_in_chunks( num_reqs = input_batch.num_reqs if logits.shape[0] <= max_chunk_logits: - # One chunk covers the batch. Adaptive verification compacts the logits - # without updating cu_num_logits_np (it keeps the pre-compacted layout), - # so the stale sums must not pick chunk boundaries; its budget cap - # guarantees the compacted batch always lands here. + # GPU trimming can leave the CPU boundaries stale. + # Trimmed batches fit in one chunk. request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),) else: assert not self.enable_adaptive_verification @@ -257,7 +252,11 @@ def _verify_in_chunks( num_sampled, processed_logits if use_processed_logits else logits[lo:hi], chunk_cu_num_logits, - chunk_cu_num_logits_np, + ( + chunk_cu_num_logits_np + if logits.shape[0] > max_chunk_logits + else None + ), max_num_logprobs, ) if chunk_logprobs is not None: From ed490a0ef55cd497e033d7c436bfa0d506655f46 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 15 Sep 2026 11:50:25 -0700 Subject: [PATCH 34/34] minor simplifications Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/model_runner.py | 9 ++++----- .../gpu/spec_decode/adaptive_verification.py | 14 ++++---------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b9157b850624..2a8d9842efe4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -286,11 +286,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): and self.speculative_config.use_multi_module_mtp() else 1 ) - - # General request states. use_dense_all_token_ids = ( self.speculative_config is not None and self.speculative_config.use_ngram() ) + + # General request states. self.req_states = RequestState( max_num_reqs=self.max_num_reqs, max_model_len=self.max_model_len, @@ -1389,9 +1389,8 @@ def prepare_inputs( ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens elif draft_trimmer is not None: - # Clamp scheduled draft slots to the drafter's valid counts on - # GPU. CPU-side totals remain upper bounds; the trimmed gap is - # treated as padding downstream. + # Clamp scheduled draft slots to the drafter's valid counts on GPU. + # CPU-side totals remain upper bounds; trimmed gap treated as padding. cu_num_logits, query_start_loc = draft_trimmer.trim( idx_mapping, num_draft_tokens_per_req, num_scheduled_tokens_np ) diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index a5a5f5df7013..1840a4eda7c7 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -98,20 +98,14 @@ def build_verification_layout( num_reqs = capacities.shape[0] cu_num_logits[:1].zero_() torch.cumsum( - capacities + num_bonus_tokens, - dim=0, - out=cu_num_logits[1 : num_reqs + 1], + capacities + num_bonus_tokens, dim=0, out=cu_num_logits[1 : num_reqs + 1] ) query_start_loc[:1].zero_() torch.cumsum( - capacities + num_non_draft_tokens, - dim=0, - out=query_start_loc[1 : num_reqs + 1], + capacities + num_non_draft_tokens, dim=0, out=query_start_loc[1 : num_reqs + 1] ) - if num_tokens is not None: - query_start_loc[num_reqs + 1 :].fill_(num_tokens) - else: - query_start_loc[num_reqs + 1 :] = query_start_loc[num_reqs] + tail = num_tokens if num_tokens is not None else query_start_loc[num_reqs] + query_start_loc[num_reqs + 1 :] = tail return cu_num_logits[: num_reqs + 1], query_start_loc