[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

[benchislett]

Suggested change

	"ngram_gpu",
	"dspark",
	"ngram_gpu",

it this an accident due to merge conflict?

[PatchouliTIS]

Yes, dspark was accidentally removed in an earlier commit. Fixed in the latest version.

[njhill]

@PatchouliTIS really sorry for the delay, we'll aim to get this merged this week!

[njhill]

@PatchouliTIS I have reworked this a bit in another branch here: main...njhill:vllm:patchy/async_ngram_v2_pr

to exploit some of what was added recently in #47808, it simplifies things quite a bit and does not require round-tripping draft tokens back to CPU. PTAL!

(note it's still not in a final state, quite a bit more simplification/cleanup can be done espec. in model_runner.py)

Output throughput (tok/s)

Workloads (max_tokens=256, ignore_eos):

• repeat — 128 verbatim-repetition prompts (high acceptance)
• prose — 128 open-ended story prompts (adversarial: matches exist but
  continuations are mostly wrong)
• longctx — 32 prompts with ~6k-token documents, repeat-section tasks

Config	repeat	prose	longctx
MRV1 baseline (no spec)	16,201	16,744	5,959
MRV1 + CPU ngram (ngram)	8,778	5,212	3,262
MRV1 + GPU ngram (ngram_gpu)	7,661	4,393	2,944
MRV2 baseline (no spec)	16,625	16,748	5,973
MRV2 + ngram_gpu, pad fallback	14,621	8,547	4,441
MRV2 + ngram_gpu, GPU trim	19,230	10,041	6,460

• GPU trim vs pad fallback: +32% / +17% / +45%.
• MRV2+trim is the only configuration that beats its own no-spec baseline
  (+16% repeat, +8% longctx).

cc @benchislett @TheEpicDolphin @WoosukKwon

[PatchouliTIS]

@njhill LGTM. The truncation of n-gram draft tokens now runs completely on the GPU side, eliminating the data transfer overhead that my implementation introduced. Would you mind sharing the benchmark scripts and commands you used? I'd love to run them locally to reproduce the results—the MRV1 ngram_gpu numbers were a bit lower than what I expected based on my previous runs, so I’d like to double check.

Also, please feel free to push directly to this branch, or let me know if you'd prefer me to pull your changes into this PR and help finish the remaining cleanup in model_runner.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.

[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.