[RFC]: Suffix decoding on GPU (`suffix_gpu`) — GPU-resident suffix decoding compatible with async scheduling

[zip95297]

Summary

Suffix decoding gets the best acceptance lengths among model-free drafters on repetitive and agentic workloads, but the in-tree implementation (method="suffix", #25784, from Arctic Inference) is CPU-side and rejected by the async-scheduling whitelist. This RFC proposes suffix_gpu: a GPU-resident suffix drafter that follows the same device-state drafter contract as ngram_gpu (#29184), so suffix decoding composes with async scheduling — no host sync on the draft path, drafts come from the previous step's device sampled ids.

The in-tree diff is small (config + a proposer wrapper + runner wiring + unit tests, ~500 lines). The drafter core (suffix automaton over device tensors, cross-request global index with background rebuilds, fused Triton kernels, CUDA-graph capture) lives in a standalone package (SuffixGPU), imported lazily — vLLM works without it installed, same pattern as the CPU suffix method's arctic-inference dependency.

Motivation

• [RFC]: Can we implement n-gram and suffix speculative decoding in model_runner_v2? #38069 asked for n-gram/suffix speculative decoding in the new runner under the condition that it stays compatible with async scheduling; it went stale, but the gap is still there.
• Async scheduling is on the way to being the default. method="suffix" forces --no-async-scheduling today (config/vllm.py whitelist: "async scheduling is only supported with EAGLE/MTP/Draft Model/NGram GPU/DSpark"), giving up CPU/GPU overlap exactly where spec decode helps most: high-concurrency TPOT.
• The reason is structural, not incidental: the CPU suffix tree needs the committed host token ids each step, and under async scheduling the scheduler runs one step ahead of the worker, so those ids do not exist on the host at schedule time. A GPU-state drafter sidesteps this by drafting from the previous step's sampled ids that are already on device — the mechanism [Core] NGram GPU Implementation compatible with Async Scheduler #29184 built for ngram_gpu.

So: keep suffix decoding's draft quality, get async scheduling's overlap, by moving the drafter's state and compute onto the GPU.

Design

suffix_gpu plugs into the existing ngram_gpu machinery; no scheduler changes.

• Config (config/speculative.py): method="suffix_gpu", reusing the suffix_decoding_* knobs (max_tree_depth, max_spec_factor, min_token_prob, max_cached_requests) plus GPU-specific ones (suffix_gpu_global_capacity, suffix_gpu_delta_capacity, suffix_gpu_max_occurrences, suffix_gpu_use_cuda_graph, suffix_gpu_ingest_chunk). A use_gpu_state_drafter() helper groups ngram_gpu + suffix_gpu for the shared device-state paths; the async whitelist admits suffix_gpu.
• Proposer (v1/spec_decode/suffix_proposer_gpu.py): wraps the SuffixGPU drafter behind the ngram_gpu contract — same update_token_ids_ngram device bookkeeping (borrowed verbatim), same propose(k, num_tokens_no_spec, token_ids, sampled, counts) -> (drafts [B,k], valid counts [B]), same async D2H of per-request valid draft counts feeding the existing worker-side trim (update_scheduler_for_invalid_drafts). Two suffix-specific additions:
  • a per-request local matcher plus a cross-request global suffix index (the analogue of the CPU method's global tree / max_cached_requests), fed off the critical path: in-flight responses are ingested on a side stream, finished requests are final-flushed before their persistent batch rows are reused;
  • the whole draft chain is captured into CUDA graphs, bucketized by batch size, pre-captured at engine warmup so the first serving step pays no Triton JIT / capture latency (eager Triton fallback if capture fails or suffix_gpu_use_cuda_graph=0).
• Runner (v1/worker/gpu_model_runner.py): the ngram_gpu-only gates (resident token buffers, optimistic-accept bookkeeping + rejection correction, valid-count D2H, scheduler_output shallow copy) become use_gpu_state_drafter() gates; suffix adds the two ingest hooks and graph pre-capture. No new scheduler interplay.

Packaging follows the suffix / arctic-inference precedent: external optional dependency, lazy import, config validation fails with a clear install hint. (Open question below.)

Correctness

All numbers below: vLLM branch commit f646af65c, 1x NVIDIA L20 48GB, torch 2.13.0+cu130, drafter package suffix-gpu 0.1.1 (PyPI, tag v0.1.1; runs used the equivalent package code at repo commit 1ac9238, installed editable — https://github.com/zip95297/SuffixGPU).

• Unit tests (tests/v1/spec_decode/test_suffix_gpu.py): draft correctness on repetitive history, CUDA-graph vs eager draft agreement, graph pre-capture at warmup, JIT warmup with graphs disabled, global index cross-request drafting. Skips cleanly without CUDA or the package. Result: 6/6 passed (L20).
• E2E GSM8K (in-tree pattern, same as ngram_gpu's async test): test_suffix_gpu_with_async_scheduling — Llama-3.1-8B-Instruct, suffix_gpu k=16 under async scheduling, 1319 questions 5-shot: accuracy 0.762 (threshold 0.70), passed.
• Acceptance e2e (tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py -k acceptance): test_suffix_gpu_acceptance mirrors the CPU suffix acceptance test — 10 warm-up rounds over the same 100 prompts, suffix_gpu k=16 + async, pinned suffix_gpu_ingest_chunk=1, num_backoff=8, max_occurrences=128, prefix caching off, V1 runner. Warm-up series: round-0 rate 0.428 / AL 2.03 -> round-9 rate 0.857 / AL 5.70; the end rate clears the CPU test's 0.80 floor (0.79-0.86 across reruns — the final-round delta is ~1.4k drafted tokens, greedy near-ties move it a few points). Both tests pass together (2 passed); both — including the pre-existing CPU one — pin enable_prefix_caching=False, since warm-round prefix-cache hits change decode batch composition and depress the measured rate below the floor for the CPU method too.
• Greedy consistency: spec on/off must produce token-identical greedy outputs. Run under VLLM_BATCH_INVARIANT=1 + enforce_eager (a bitwise gate is meaningless otherwise: spec decode changes batch composition, and batch-dependent FP reduction order flips near-tie argmax even between two no-spec runs — measured on this box before enabling the invariant mode). Legs: async no-spec vs async suffix_gpu, plus suffix_gpu-eager / ngram_gpu / sync suffix CPU discriminators, 16 prompts x 256 tokens. Result: all legs token-identical (0/16 mismatches each; graph replay output == eager output).
• Drafter-level equivalence and latency vs the CPU suffix tree (no engine, deterministic Spec-Bench replay): 221 tests pass including fuzz-equivalence against arctic on unambiguous corpora; replay tokens/step within -7%..0% of the CPU tree cold and parity to +3% warm; <= 532 MB reserved VRAM in the largest tested configuration. On drafting cost, the CPU tree wins small batches but its sequential per-request walk grows linearly with batch size, while graph-mode drafting is one flat batched launch: crossover at B64-128, and beyond it CUDA-graph mode is 2.1-2.7x faster per step even in the CPU's best case (B=256: 1.85 vs 4.98 ms; B=512: 5.15 vs 10.68 ms); under replay conditions that include per-step tree updates the GPU drafter is already faster from B32 (CPU 1.0-3.4 ms at B=58-80 vs a flat ~0.5 ms). This is what makes the drafter viable exactly in the high-concurrency regime async scheduling targets. Full tables and repro commands: https://github.com/zip95297/SuffixGPU/blob/main/RESULTS.md

Performance

Setup mirrors #25784 (Spec-Bench, k x concurrency grid, TPOT + drafted / accepted tokens) and #29184 (async-vs-sync same-method comparison): meta-llama/Llama-3.1-8B-Instruct, bf16, TP=1, 1x NVIDIA L20 48GB, vllm bench serve --dataset-name spec_bench --spec-bench-output-len 256 --no-oversample, concurrency 1/4/16/64/128/192/256, k in {5,16}, prefix caching off, all variants on model runner v1. Every variant gets --max-num-batched-tokens 8192 --max-num-seqs 320 (spec decode schedules (k+1) padded token slots per running request before invalid ones are trimmed, so the serve-context defaults silently cap effective concurrency at the top grid points) and an identical 64-prompt warmup pass before measurement (cold CUDA graphs / Triton JIT / suffix index otherwise land in the first grid point; drafted/accepted numbers are the per-run deltas vllm bench serve records in each result json, so warmup traffic is excluded by construction). Reproduction script: benchmarks/bench_pr_matrix.sh (attached to the PR).

Variants: async no-spec baseline; suffix + sync scheduling (what you must run today); ngram_gpu + async (in-tree GPU drafter baseline); suffix_gpu + async (this RFC).

k=5 — output tok/s (mean TPOT ms)

variant	c1	c4	c16	c64	c128	c192	c256
async_nospec	47 (21.0)	172 (22.7)	578 (26.4)	1286 (47.0)	1566 (75.6)	1659 (109.6)	1654 (147.1)
suffix_sync	58 (17.2)	245 (15.4)	754 (19.8)	1297 (45.5)	1343 (86.0)	1375 (127.2)	1373 (169.5)
ngram_gpu_async	46 (21.6)	168 (23.3)	515 (30.3)	1126 (52.7)	1341 (89.2)	1389 (128.5)	1396 (172.3)
suffix_gpu_async	57 (17.6)	239 (15.6)	744 (19.3)	1316 (43.2)	1410 (80.5)	1461 (117.6)	1467 (154.5)

k=16 — output tok/s (mean TPOT ms)

variant	c1	c4	c16	c64	c128	c192	c256
async_nospec	47 (21.0)	172 (22.7)	578 (26.4)	1286 (47.0)	1566 (75.6)	1659 (109.6)	1654 (147.1)
suffix_sync	55 (18.2)	239 (15.9)	755 (19.5)	1273 (46.3)	1322 (86.0)	1357 (125.8)	1363 (167.2)
ngram_gpu_async	43 (23.0)	161 (24.3)	519 (29.9)	1156 (52.2)	1339 (89.6)	1387 (131.0)	1405 (170.6)
suffix_gpu_async*	52 (19.2)	230 (16.2)	736 (19.5)	1276 (44.8)	1349 (84.3)	1402 (121.5)	1417 (154.6)

* suffix_gpu k=16 ran at gpu_memory_utilization=0.85 (other rows 0.9): at 0.9 the device-resident drafter state plus the (k+1)-slot verification batch OOMed the EngineCore at c256 on the 46GB L20 (see Known limits).

Drafted / accepted tokens per point (AL, acceptance rate)

variant	k	c1	c16	c64	c256
suffix_sync	5	79832/22918 (1.54, 28.7%)	86405/33544 (1.94, 38.8%)	86516/32697 (1.89, 37.8%)	86240/31708 (1.85, 36.8%)
suffix_gpu_async	5	79742/22332 (1.52, 28.0%)	87458/33453 (1.93, 38.3%)	88350/35551 (2.04, 40.2%)	88444/38485 (2.23, 43.5%)
ngram_gpu_async	5	33913/9467 (2.40, 27.9%)	34863/9738 (2.40, 27.9%)	34527/9332 (2.35, 27.0%)	34358/9667 (2.41, 28.1%)
suffix_sync	16	83232/22867 (1.54, 27.5%)	95070/34453 (1.98, 36.2%)	93340/33505 (1.94, 35.9%)	92211/32536 (1.89, 35.3%)
suffix_gpu_async	16	82552/22582 (1.53, 27.4%)	95486/34166 (2.00, 35.8%)	98323/36571 (2.13, 37.2%)	105420/40513 (2.33, 38.4%)
ngram_gpu_async	16	35468/9978 (2.41, 28.1%)	34844/9586 (2.38, 27.5%)	33685/9643 (2.43, 28.6%)	34316/9778 (2.42, 28.5%)

Per-position acceptance rate — suffix CPU vs suffix_gpu, c=1 only

Measured at concurrency 1 (the c1 grid point, after the identical warmup pass), where scheduling noise is minimal, so this is the cleanest drafter-quality comparison. Each row also lists that run's drafts / draft tokens / accepted tokens / AL / overall acceptance rate.

k=5, c=1:

variant	drafts	draft toks	accepted toks	AL	rate	p0	p1	p2	p3	p4
suffix_sync	42793	79832	22918	1.54	28.7%	0.336	0.120	0.049	0.019	0.012
suffix_gpu_async	42945	79742	22332	1.52	28.0%	0.330	0.114	0.046	0.019	0.012

k=16, c=1:

variant	drafts	draft toks	accepted toks	AL	rate	p0	p1	p2	p3	p4	p5	p6	p7	p8	p9	p10	p11	p12	p13	p14	p15
suffix_sync	42677	83232	22867	1.54	27.5%	0.329	0.115	0.046	0.018	0.010	0.005	0.004	0.002	0.002	0.002	0.001	0.001	0.000	0.000	0.000	0.000
suffix_gpu_async	42293	82552	22582	1.53	27.4%	0.328	0.113	0.045	0.018	0.011	0.005	0.004	0.002	0.002	0.001	0.001	0.001	0.001	0.001	0.001	0.000

The two per-position curves are essentially identical at both k values: the GPU drafter reproduces the CPU suffix tree's draft quality request-locally. Its aggregate advantage appears only as concurrency grows (AL 1.53 -> 2.33 from c1 to c256 at k=16, vs 1.98 -> 1.89 for the CPU tree), i.e. it comes from the cross-request global index compounding under concurrent traffic, not from a different per-draft profile.

What the numbers show:

1. suffix_gpu + async beats suffix + sync from c64 up (+7% at k=5 saturation, 1467 vs 1373 tok/s; +4% at k=16, 1417 vs 1363; lower TPOT), and matches it at c1-c16 — resolving the suffix x async conflict is a real win once batches grow, exactly where async overlap matters.
2. suffix_gpu + async beats ngram_gpu + async at every grid point (up to +44%, k=5 c16: 744 vs 515 tok/s) — it is the stronger GPU drafter on this workload.
3. Against async no-spec, spec decoding wins through c64 (TPOT -31% at c4, -27% at c16, -8% at c64) but loses beyond c128 at saturation on Spec-Bench's mixed categories (1467 vs 1654 tok/s at k=5 c256) — true for all three drafters here (suffix_sync 1373, ngram_gpu 1396), see Known limits.

Known limits

• Saturation: beyond ~c128 on Spec-Bench's mixed categories, async no-spec out-throughputs every drafter here (k=5 c256: no-spec 1654 vs suffix_gpu 1467, suffix_sync 1373, ngram_gpu 1396 tok/s) — once the GPU is compute-bound, verification FLOPs for rejected drafts cost more than acceptance saves. On repetition-heavy traffic the crossover moves right; measured, not hidden. (Contrary to the usual GPU-drafter expectation, c=1 does not regress on this workload: 57-58 tok/s vs 47 no-spec, the warm suffix index pays for the drafter cost.)
• Memory headroom at k=16: the device-resident drafter state plus the (k+1)-slot verification batch OOMed the EngineCore at c256 with the default gpu_memory_utilization=0.9 on a 46GB L20; the k=16 row was measured at 0.85. Large-k + high-concurrency deployments need to budget for the drafter's device state.
• Spec-decode acceptance metrics for GPU-state drafters counted scheduler-padded slots in the denominator (affects ngram_gpu equally). This PR includes the worker-side fix — per-request trim counts ride on ModelRunnerOutput and are subtracted in make_spec_decoding_stats — with method-neutral naming; mechanism-wise it matches the stale open ngram-only PR fix(ngram): match async ngram_gpu acceptance rate to CPU #44056 (happy to rebase or split if maintainers prefer; [Bugfix][Core][Spec Decode] Exclude scheduler padding from draft metrics #50518 covers the scheduler-side view for uniform batches).

Proposed change

Land suffix_gpu as a speculative decoding method: config surface, SuffixProposerGPU, runner wiring behind use_gpu_state_drafter(), unit tests. Branch: zip95297/suffix_gpu on https://github.com/zip95297/vllm-dev.

Docs plan (same shape as #25784's docs addition): a suffix_gpu subsection in docs/features/spec_decode.md — when to pick it over suffix/ngram_gpu, the suffix_gpu_* knobs and their defaults, the async-scheduling compatibility note, the memory-headroom guidance for large k, and the package install hint. Lands with this PR or as an immediate follow-up once the packaging question below is settled (in-tree vs external changes the install section).

Open question

Keep the drafter core as an external optional dependency (arctic-inference precedent; small in-tree diff, kernels iterate on their own cadence) or port it in-tree under v1/spec_decode/ (~1.5 kLOC pure Python + Triton, torch-only deps; one source of truth, CI coverage)? I lean in-tree for a decode-hot-path feature but either works; the wiring in this PR is the same either way.

Future work

Variable draft-length scheduling for GPU-state drafters (Dynamic SD #32374 integration and per-request draft lengths) is deliberately out of scope; it applies to ngram_gpu and suffix_gpu equally and deserves its own design discussion.

[zip95297]

Once suffix_gpu lands, I'll apply Dynamic SD (#32374) to both GPU-state drafters — ngram_gpu and suffix_gpu. The machinery is already in tree and method-agnostic: num_speculative_tokens_per_batch_size expands to a batch-size → K lookup, schedule() reads it into num_spec_tokens_to_schedule, the async placeholder path already uses that value, _maybe_override_dynamic_sd_cudagraph_mode auto-switches cudagraph to PIECEWISE, and the worker trim (update_scheduler_for_invalid_drafts) already clamps valid_k to the scheduled K. The only thing blocking it today is that both drafters assert num_speculative_tokens == self.k inside propose() — Dynamic SD was built for CPU drafters whose propose loops accept a variable K. Relaxing that assert (keep drafting self.k tokens into the fixed [B, k] buffer / CUDA graph, let the worker trim down to the scheduled K) is a one-line change per drafter.

Why this first: it directly lifts the fixed-k saturation ceiling. Every decode request costs 1+k against the token budget, so k=32 caps in-flight decodes at max_num_batched_tokens / 33; shrinking K at high batch size lets more requests advance per step, and it is adaptive — small batches keep a big K for AL/TPOT, large batches shrink K for throughput.

What it does not do: it is global per step, not per request, so a request with a long suffix match still gets truncated to the step's K; the schedule is static (configured), not learned from observed acceptance; and c1 is unaffected (batch=1 maps to max K).

[zip95297]

Implementation PR is up: #52097

Quick orientation for reviewers:

• This directly addresses [RFC]: Can we implement n-gram and suffix speculative decoding in model_runner_v2? #38069, where @WoosukKwon noted n-gram/suffix support would be welcome "as long as the implementation is compatible with async scheduling" — suffix_gpu is exactly that.
• suffix_gpu follows the ngram_gpu device-state drafter contract ([Core] NGram GPU Implementation compatible with Async Scheduler #29184), so suffix decoding composes with async scheduling; no scheduler changes.
• ~590 lines outside tests (+~300 tests). The drafter core is an external optional dependency (SuffixGPU), same pattern as suffix / arctic-inference — in-tree vs. external packaging is the main open question I'd appreciate feedback on.
• Correctness: unit tests 6/6; GSM8K under async scheduling 0.763 (threshold 0.70); acceptance-rate e2e mirroring the CPU suffix test reaches 0.857 / AL 5.70; greedy outputs token-identical with spec on/off under batch-invariant mode.
• The padded-slot metrics fix included here also affects ngram_gpu (mechanism matches fix(ngram): match async ngram_gpu acceptance rate to CPU #44056);

cc @benchislett @WoosukKwon @sfc-gh-aqiao @PatchouliTIS @njhill

[zip95297]

Update: MRv2 integration status + a fixed-k verification cost question (#40704)

I've been porting suffix_gpu to ModelRunner V2 behind the BaseSpeculator interface (drafting directly from the runner-maintained RequestState.all_token_ids / total_len, no token mirror, no draft-id round-trips). Correctness is in good shape: GSM8K accuracy matches V1, and per-point accepted/drafted counts are within ~3% of V1 at every grid point. But throughput regresses vs V1 on the same commit, same configs, with only the runner flag flipped (Spec-Bench, Llama-3.1-8B-Instruct bf16 TP=1, 1x L20, async scheduling, 480 reqs/point, output-len 256):

V2 vs V1 (output tok/s)	k=5	k=16
c = 1–16	−21% … −28%	−28% … −44%
c = 64–256	−17% … −23%	−53% … −56%

async_nospec V1 vs V2 is flat (±1.6%) across the whole sweep, so the runner itself is not the issue. The dominant cost is fixed-k verification: async scheduling reserves k placeholder slots per running request before the draft lengths are known. The V1 runner trims the invalid slots pre-forward (update_scheduler_for_invalid_drafts), so padding costs ~0 FLOPs; on V2 every decode request pays the full (k+1)-token forward regardless of the drafter's actual emission (k=16 with ~3.5 valid tokens ⇒ ~3.8x verify tokens per step), and acceptance metrics additionally need a worker-side num_invalid_spec_tokens correction so padding is not counted as real drafts. A secondary cost at low concurrency is the drafter running over the full slot space instead of the active batch.

My current branch carries a host-side version of the V1 trim: a shallow copy of SchedulerOutput shrunk at the start of execute_model, driven by per-slot num_valid counts D2H-copied on a side stream. It is functionally fine, but it needs an event sync at step start to learn the previous step's valid counts on the host — which puts a CPU↔GPU sync back on the async-scheduling critical path.

#40704 looks like it solves exactly this for ngram/ngram_gpu on MRv2 — VariableDraftTrimmer clamping scheduled draft slots on device via the num_valid_drafts_for_trim contract (the drafter publishes per-slot valid counts on GPU; the runner clamps cu_num_logits / query_start_loc with the CPU keeping only upper bounds, no CPU sync, reusing the DSpark adaptive-verification layout machinery, falling back to pad-and-verify when the attention backend can't support varlen trimming), plus idx_mapping-indexed n-gram kernels over resident token rows.

Three questions:

1. Is VariableDraftTrimmer intended to be drafter-agnostic — i.e., should a second GPU-state drafter (suffix) simply publish num_valid_drafts_for_trim and drop any drafter-specific trimming? Anything that would block reusing it as-is?
2. Would you recommend stacking the suffix_gpu MRv2 integration on top of [ModelRunner V2] Speculative Decoding NGram GPU Implementations #40704 , following the same RequestState constructor-injection and idx_mapping row-indexed kernel patterns?
3. With device-side trimming, never-issued slots surface as ordinary rejections in the spec-decode stats. For suffix I'd like to keep reporting them as invalid rather than rejected (async D2H of the same per-slot buffer → num_invalid_spec_tokens), so acceptance stays comparable with the V1 runner — would that be acceptable alongside the trimmer?

Happy to rework onto the num_valid_drafts_for_trim contract if that's the preferred direction. cc @PatchouliTIS @njhill