perf(sampling): opt-in fast verify path for topk=1 chain spec

[cicirori]

Summary

Opt-in optimization for chain-spec (eagle_topk = 1) only. In this configuration the draft is always argmax-greedy, so the verify pass can collapse the sequential renorm chain into a single fused-sample kernel followed by a constant-time token-equality verify.

Today the spec verify path goes through

softmax → top_k_renorm_prob → top_p_renorm_prob → chain_speculative_sampling_target_only

which scans full vocab twice at the renorm steps. For vocab≈200K (M2.5/MiniMax) this dominates verify sampling cost.

This PR adds an opt-in branch in verify() under TOKENSPEED_SPEC_FAST_CHAIN_SAMPLING=1 that, in the topk=1 chain-spec case:

1. Samples one target token per chain position via top_k_top_p_sampling_from_logits(filter_apply_order="top_k_first") — the same flashinfer API and filter mode that stock TRTLLM 1.3.0rc14 uses for verify
2. Verifies by token-equality against the (already argmax-greedy) draft candidates using the existing verify_chain_greedy kernel (same kernel the greedy branch already uses)
3. Bonus token at first mismatch = the sampled target there

Scope and applicability

Only applies when all of the following hold:

• speculative_eagle_topk == 1 (chain spec) — tree spec topk>1 keeps the generic rejection kernel
• non-greedy sampling — greedy path is already on the fast verify_chain_greedy kernel
• no grammar mask — per-position bitmasks need the renorm chain (auto-fallback)
• request param hint matches the narrow scope this backend supports (no min_p, no penalties, no logit_bias)

For any other configuration, verify() falls through to the existing sequential path with no behavior change.

Measured impact (M2.5plus-fp4 NVFP4, TP=2 B200, vocab=200054, c=32 / 200 reqs)

Sampling kernel level (per verify, inside captured CUDA graph)

Captured via nsys profile --cuda-graph-trace=node at typical full batch (gridX=128):

Path	breakdown	total µs/verify
Sequential (baseline)	OnlineSoftmax 89 + RadixTopKRenormProb 165 + AirTopPRenormRadix×3 410 + AirTopPRenormApply 191 + AirTopPInit 2 + ChainSpecSampling 84	~941
Fast (this PR)	RadixTopKMaskLogits + TopPSamplingFromProb + VerifyChainGreedy (fused dispatch)	~55–90
Speedup		~10–17×

The fast-path range covers small-batch (~55 µs, well-sampled at gridX=4) up to extrapolated full-batch (~90 µs). nsys event buffers overflow aggressively at full batch with --cuda-graph-trace=node, so the precise full-batch number for top_k_first is bracketed by the joint-mode capture at the same gridX (~85 µs sample + ~2 µs verify) and the +6–19% micro-bench delta between the two filter modes.

End-to-end throughput (3 reps each, no nsys overhead)

Config	avg tok/s	stddev	Δ
Sequential (default)	1667	31 (1.9%)	—
Fast (this PR)	1707	32 (1.9%)	+2.4% (within 1σ)

The ≥10× kernel speedup does not fully translate to E2E because sampling now occupies <1% of decode iter wall in this backend — MoE + attention + allreduce dominate. The PR is most useful when:

• sampling shows up as a hotspot in profiling (e.g. larger batch or larger vocab models)
• you want to eliminate the renorm chain to simplify reasoning about decode hot path
• a future configuration shifts the bottleneck

Numerical semantics

filter_apply_order="top_k_first" is bit-equivalent in filter semantics to the sequential top_k_renorm_prob → top_p_renorm_prob chain it replaces — both apply top-p on the top-k-renormalized distribution. Same prompt + sampling params produce statistically equivalent samples whether spec is on or off (empirical TV vs sequential analytic reference: 0.0040 at vocab=1024 / top_k=50 / top_p=0.95, identical noise floor to the sequential path's own 0.0035).

An earlier revision of this PR used filter_apply_order="joint" (top-k AND top-p applied simultaneously on the original distribution), which introduced a ~5% TV semantic shift relative to the sequential path. Profiling showed joint was at most ~19% faster than top_k_first at full-batch kernel level but the difference disappeared in E2E noise, so top_k_first was selected for the cleaner semantics and TRTLLM alignment.

Prior art

The chain-spec + token-equality verify shape, the flashinfer API choice, and the filter_apply_order="top_k_first" selection all mirror stock TensorRT-LLM 1.3.0rc14's verify sampling (tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py).

CUDA graph compatibility

The fast path uses kernels (top_k_top_p_sampling_from_logits, verify_chain_greedy) that are already in-graph in this backend's sample() and greedy verify() paths. Confirmed via nsys --cuda-graph-trace=node: fast-path verify kernels in the benchmark trace are reported as in_graph (alongside eager calls from prefill sample()).

Persistent buffers (_predict_buf, _accept_index_buf, _accept_length_buf) are reused as in the sequential path — no per-step allocation in the hot path.

Tests

tokenspeed-kernel/test/ops/test_fast_chain_sampling.py — Monte Carlo on 262,144 trials, two checks:

1. Fast path's empirical output matches the sequential-filter analytic reference within noise (TV < 0.02, measured 0.0040)
2. Sequential path's empirical output also matches the sequential-filter analytic reference (TV < 0.02, measured 0.0035) — validates the test infrastructure

Both pass.

Opt-in / default

Default off. Read once at import time. Spawn-launched TP workers inherit os.environ via Python's cached dict (independent of setproctitle's clobber on /proc/<pid>/environ), so setting TOKENSPEED_SPEC_FAST_CHAIN_SAMPLING=1 before launching propagates to all workers.

Test plan

• Statistical correctness test (2 Monte Carlo checks, both passing at TV ≈ 0.004)
• E2E benchmark, 3 reps × baseline+fast at c=32
• nsys profile --cuda-graph-trace=node confirms in-graph kernel composition and per-verify kernel speedup
• Pre-commit clean (isort, black, etc.)

🤖 Generated with Claude Code