[ModelRunner V2] Speculative Decoding NGram GPU Implementations

[PatchouliTIS]

Purpose

1. Added a new NGram GPU speculator.
   The main feature is a new implementation at: vllm/v1/worker/gpu/spec_decode/ngram/speculator.py, similiar to [Core] NGram GPU Implementation compatible with Async Scheduler #29184.

2. Updated request state storage for NGram GPU.
   vllm/v1/worker/gpu/states.py
vllm/v1/worker/gpu/model_runner.py
   This changes how RequestState is initialized so that all_token_ids can stay densely resident on GPU instead of defaulting to UVA when ngram_gpu is active. As discussed in [Core] NGram GPU Implementation compatible with Async Scheduler #29184, the new n-gram speculator repeatedly scans active request token history, doing that from GPU-resident dense storage is much more appropriate than pulling through UVA-backed memory, this is a performance-oriented architectural change supporting the new feature.
   The model_runner.py also injects req_states into speculators that need direct access to the persistent token store.

3. Added variable-length draft token plumbing
   Several files were updated to support draft proposals where different requests may have different numbers of valid draft tokens:
   vllm/v1/outputs.py
vllm/v1/worker/gpu/spec_decode/utils.py
vllm/v1/core/sched/scheduler.py
vllm/v1/engine/core.py
vllm/v1/worker/gpu/model_runner.py

DraftTokenIds now includes:
num_valid_draft_tokens: list[int] | None.
Scheduler logic now truncates speculative tokens based on num_valid_draft_tokens.
EngineCore adds _maybe_update_async_draft_token_ids() to consume draft metadata from async execution and update scheduler state at the right time.

Test Plan

vllm bench cmd:

vllm bench serve \
--port 8000 \
--backend vllm \
--model Qwen3-8B \
--endpoint /v1/completions \
--dataset-name sonnet \
--dataset-path sonnet.txt \
--max-concurrency 128 \
--sonnet-input-len 128 \
--sonnet-output-len 100 \
--sonnet-prefix-len 10 \
--num-prompts 256 \
--ignore-eos \
--percentile-metrics "ttft,tpot,itl,e2el" \
--seed 1234

Test Result

Async NGram GPU V1 results:

============ Serving Benchmark Result ============
Successful requests:                     256       
Failed requests:                         0         
Maximum request concurrency:             128       
Benchmark duration (s):                  10.42     
Total input tokens:                      31693     
Total generated tokens:                  25600     
Request throughput (req/s):              24.56     
Output token throughput (tok/s):         2456.33   
Peak output token throughput (tok/s):    4277.00   
Peak concurrent requests:                186.00    
Total token throughput (tok/s):          5497.29   
---------------Time to First Token----------------
Mean TTFT (ms):                          1348.49   
Median TTFT (ms):                        1083.36   
P99 TTFT (ms):                           2778.11   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          34.86     
Median TPOT (ms):                        34.45     
P99 TPOT (ms):                           51.03     
---------------Inter-token Latency----------------
Mean ITL (ms):                           40.09     
Median ITL (ms):                         30.44     
P99 ITL (ms):                            236.84    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          4799.33   
Median E2EL (ms):                        4419.59   
P99 E2EL (ms):                           7136.50   
---------------Speculative Decoding---------------
Acceptance rate (%):                     5.16      
Acceptance length:                       1.15      
Drafts:                                  22038     
Draft tokens:                            66114     
Accepted tokens:                         3409      
Per-position acceptance (%):
  Position 0:                            6.27      
  Position 1:                            5.34      
  Position 2:                            3.86      
==================================================

Async NGram GPU V2 results:

============ Serving Benchmark Result ============
Successful requests:                     256       
Failed requests:                         0         
Maximum request concurrency:             128       
Benchmark duration (s):                  8.41      
Total input tokens:                      31693     
Total generated tokens:                  25600     
Request throughput (req/s):              30.43     
Output token throughput (tok/s):         3042.78   
Peak output token throughput (tok/s):    4471.00   
Peak concurrent requests:                187.00    
Total token throughput (tok/s):          6809.77   
---------------Time to First Token----------------
Mean TTFT (ms):                          980.36    
Median TTFT (ms):                        713.87    
P99 TTFT (ms):                           1685.63   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          31.17     
Median TPOT (ms):                        31.34     
P99 TPOT (ms):                           42.63     
---------------Inter-token Latency----------------
Mean ITL (ms):                           35.77     
Median ITL (ms):                         28.55     
P99 ITL (ms):                            154.67    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          4065.85   
Median E2EL (ms):                        3954.04   
P99 E2EL (ms):                           5862.69   
---------------Speculative Decoding---------------
Acceptance rate (%):                     52.64     
Acceptance length:                       2.58      
Drafts:                                  2136      
Draft tokens:                            6408      
Accepted tokens:                         3373      
Per-position acceptance (%):
  Position 0:                            64.04     
  Position 1:                            54.63     
  Position 2:                            39.23     
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request introduces a GPU-accelerated N-gram speculator (NgramGPUSpeculator) for speculative decoding in vLLM V2. Key changes include infrastructure to handle variable-length draft tokens, memory management updates to keep token IDs on the GPU for faster scanning, and performance optimizations for the Gumbel sampler by making FP64 precision optional. A review comment suggests improving the N-gram matching logic in the kernel to select the most recent occurrence of a pattern instead of the first, which better captures local context for prompt lookup.

[gemini-code-assist]

The current implementation uses argmax(dim=1) on the boolean matches tensor, which finds the first occurrence of the n-gram pattern in the sequence. In speculative decoding (specifically prompt lookup), it is standard practice and significantly more effective to use the most recent (last) occurrence of the pattern, as it better captures the local context.

You can find the last match by applying argmax to a tensor of indices where matches occur, which will return the largest index for each row.

Suggested change

	idx = matches.int().argmax(dim=1)
	has_match = matches[batch_idx, idx]
	first_match_pos[:, i] = torch.where(has_match, idx.long(), -1)
	# Find the last match by using argmax on indices to get the most recent occurrence
	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)

[WoosukKwon]

@TheEpicDolphin

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[TheEpicDolphin]

@PatchouliTIS it seems like the FP64 -> FP32 gumbel sample changes are not necessary for enabling the ngram functionality. Would it be possible to separate those changes out into a separate PR?

[PatchouliTIS]

@PatchouliTIS it seems like the FP64 -> FP32 gumbel sample changes are not necessary for enabling the ngram functionality. Would it be possible to separate those changes out into a separate PR?

okay, I'm on my vacation now and I will handle this next week.

[PatchouliTIS]

Removed gumbel sampling modifications from this PR, ready for review. @TheEpicDolphin

[mergify]

Hi @PatchouliTIS, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[njhill]

I pushed my changes along with a bit more rework, but I'm still not sure it's in the best state w.r.t. how the draft trimmer abstraction is structured.

Would you mind sharing the benchmark scripts and commands you used?

@PatchouliTIS here is the benchmark script that was used. I hadn't actually read it, I guess the prompts aren't ideal and we could use some better test sets for this: bench_ngram_trim.py

[mergify]

Hi @PatchouliTIS, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[ekagra-ranjan]

Thanks @PatchouliTIS for the PR and maintainers for reviewing it! Really excited to see Ngram x MRv2 getting merged!

@njhill - just wondering if the PR is likely to be merged soon, or if there’s any blocker preventing the merge?

[kosztyua]

Hi @PatchouliTIS — we tried a fixed-budget port of your GPU ngram implementation on Qwen3.8-Flash-Next across two DGX Sparks, and got useful copy/edit results. Thanks for the scan/finalize kernels; this is based on your work at ffeadf6e77cf32c37b4e2ab603d0ce1fe50c6604, not a new ngram algorithm.

Against our optimized native-MTP4 baseline (already using the FP8 proposal head from #56577), K12 measured:

Workload	MTP4 tok/s	GPU ngram K12 tok/s
Configuration copy	84.20	144.41
Configuration small edit	83.59	131.61
Python-file copy	83.38	173.03
Python-file small edit	82.09	163.34

That's 57–108% higher output throughput on these copy/edit fixtures. Fresh paired runs used concurrency 1, one warmup plus three measured requests per case, eager execution and prefix caching off. All 32 responses matched the expected outputs exactly, with matching per-case hashes/lengths across arms and natural stops (552/552/1094/1095 output tokens). An earlier configuration-only screen agreed. The GPU proposer tests passed 20 cases, with 2 explicit skips for variable-length trimming that we did not port.

The local integration needed proposer construction after request-state allocation, request-history plumbing, dp_sync compatibility, eager no-op graph hooks, and handling a non-MTP proposer at both MTP hidden-state getter sites. We kept the existing target verification/rejection and PLE/GDN state paths.

Our runtime base was 7ef4d9bfed6311e3b78a40abb4a8bb6a2fc741b0 with the FP8-head overlay, checkpoint nvidia/Qwen3.8-Flash-Next-NVFP4 at fab0aecb760cec45227f6656abcaafa11abca87a, TP2 and BF16 KV. The configured context limit was 262144, but these lookup fixtures were short—not long-context qualification.

This is an alternative proposer for copy/edit work, not simultaneous MTP+lookup or a general-generation speedup. We have not validated the full PR, variable-budget scheduler integration, CUDA graphs, concurrent requests or general sampling.

Would a minimal Qwen/V2 compatibility patch and the fixed-budget reproduction cases be useful here? We can prepare those against your preferred base; we do not plan to open a competing ngram implementation PR.

Development and testing were AI-assisted, with the experiments executed by the coding agent.

[ekagra-ranjan]

@PatchouliTIS

I see in the Test plan that the AL reported for sonnet is 1.15 for MRv1 and 2.58 for MRv2.
a. Wondering why is there such a big difference in AL when the dataset is the same?
b. Perhaps it wont make any diff in this eval setup but I think we should avoid using --ignore-eos when running SD experiments since it could pollute the AL towards the end of generation if we generate beyond EOS which could be weird to draft for.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @PatchouliTIS.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ekagra-ranjan]

@PatchouliTIS @njhill - wondering if the code correctness is established with current state of the PR and only code quality improvement is being worked on OR there is something critical missing from the PR and needs to fixed/implemented? Asking because I wanted to use ngram-gpu in MRv2 so checking if its safe to build from the current commit of the PR. Lmk otherwise if its unsafe to use the PR right now.