[ModelRunner V2] Speculative Decoding NGram GPU Implementations

vllm/envs.py

@@ -192,6 +192,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"
@@ -1417,6 +1418,10 @@ def _get_or_set_default() -> str:
     "VLLM_COMPUTE_NANS_IN_LOGITS": lambda: bool(
         int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0"))
     ),
+    # Optimization for Gumbel sampler.
+    "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

vllm/v1/core/sched/scheduler.py

@@ -1671,9 +1671,12 @@ 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,
+        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():
@@ -1686,6 +1689,12 @@ 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
@@ -1696,11 +1705,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():
@@ -1714,7 +1726,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

vllm/v1/engine/core.py

@@ -72,7 +72,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
@@ -427,7 +427,26 @@ 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.
+        """
+        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:
+            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,
@@ -531,6 +550,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.
@@ -539,7 +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 = 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

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(

vllm/v1/worker/gpu/model_runner.py

@@ -186,20 +186,30 @@ 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,
             max_num_batched_tokens=self.max_num_tokens,
             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,
             max_num_tokens=self.max_num_tokens,
             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
+
         self.sampler: Sampler | None = None
         self.rejection_sampler: RejectionSampler | None = None
         self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
@@ -294,7 +304,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.
@@ -1233,7 +1243,13 @@ 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)
+            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

vllm/v1/worker/gpu/sample/gumbel.py

@@ -2,8 +2,13 @@
 # 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_TINY = tl.constexpr(float.fromhex("0x1p-126"))
+
 
 @triton.jit
 def _temperature_kernel(
@@ -18,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)
@@ -77,36 +82,39 @@ def gumbel_block_argmax(
     processed_logits_ptr,
     processed_logits_stride,
     APPLY_TEMPERATURE: tl.constexpr,
+    USE_FP64: tl.constexpr,
 ):
     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,
         )
 
-    logits = logits.to(tl.float64)
+    # 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)
 
-        # 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)
@@ -168,12 +176,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),
@@ -190,6 +202,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)

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.")