Repository navigation
[qwen 3.8 next] Fuse exact greedy chain verification with argmax finalization - #41172
Open
Qiaolin-Yu wants to merge 1 commit into
Open
Qiaolin-Yu wants to merge 1 commit into
Qiaolin-Yu wants to merge 1 commit into
Conversation
Qiaolin-Yu
marked this pull request as ready for review
September 24, 2026 21:32
Qiaolin-Yu
requested review from
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
Ying1123,
celve,
hnyls2002,
merrymercy and
yuan-luo
as code owners
September 24, 2026 21:32
This was referenced Sep 24, 2026
Qiaolin-Yu
changed the title
7 of 10 tasks
nvrohanv
added a commit
to nvrohanv/sglang
that referenced
this pull request
Oct 25, 2026
…aths for EAGLE v2 chains Opt-in with one flag, outputs bitwise identical: --enable-speculative-step-cuda-graph captures each EAGLE v2 decode step as one CUDA graph, built on splicing full-backend graphs into an enclosing capture, and turns on six pieces: one TP broadcast for the verify result, a lean sampled verify for chains, fused relay stores with single-request relay views, a fused topk-1 draft-extend tail, one launch for the step graph's input copies, and chain verify sampling in two launches. Each piece has a SGLANG_SPEC_* environment variable as a test / bisect override (0 turns it off under the flag). With the flag off and no override set each changed file runs its existing code. A weight update that asks for a CUDA graph recapture is refused while the step graph is on. The fused relay carries the change proposed in sgl-project#41170; the draft-extend tail uses the two-stage argmax form proposed in sgl-project#41172. At startup, when the flag is on and a step graph runner is active, the HTTP server sends one batch-size-1 request per DP rank with the sampling a chat request without sampling fields resolves to (the chat API's own to_sampling_params over the model's generation config), so the step graph captures that key before serving; it is skipped with --skip-server-warmup, and with SGLANG_RUST_SERVER one line says so. A failed capture memory check is logged once per key on every rank with the free, needed, reserved and allocated memory, and /get_server_info reports memory_refusals, last_memory_refusal and the allocator's num_alloc_retries / num_ooms. The draft-extend tail is registered in the kernel inventory (Triton, CUDA only). Also kernel tests for chain verify sampling, the draft-extend tail and the relay scatter.
nvrohanv
added a commit
to nvrohanv/sglang
that referenced
this pull request
Oct 26, 2026
…aths for EAGLE v2 chains Opt-in with one flag, outputs bitwise identical: --enable-speculative-step-cuda-graph captures each EAGLE v2 decode step as one CUDA graph, built on splicing full-backend graphs into an enclosing capture, and turns on six pieces: one TP broadcast for the verify result, a lean sampled verify for chains, fused relay stores with single-request relay views, a fused topk-1 draft-extend tail, one launch for the step graph's input copies, and chain verify sampling in two launches. Each piece has a SGLANG_SPEC_* environment variable as a test / bisect override (0 turns it off under the flag). With the flag off and no override set each changed file runs its existing code. A weight update that asks for a CUDA graph recapture is refused while the step graph is on. The fused relay carries the change proposed in sgl-project#41170; the draft-extend tail uses the two-stage argmax form proposed in sgl-project#41172. At startup, when the flag is on and a step graph runner is active, the HTTP server runs, for each batch size k from 1 to N, one batch of k requests per DP rank with the sampling a chat request without sampling fields resolves to (the chat API's own to_sampling_params over the model's generation config). N is the smaller of the runner's largest batch size (at most 8, and at most the largest of every per-stage graph) and the maximum number of running requests. A batch of token ids reaches the scheduler as one message, so its requests decode as one batch of k, and the step graph captures each of those keys before serving; a batch size still without a key gets one more batch, then a warning. The log lists each captured key with its memory and the time the warmup took. The warmup is skipped with --skip-server-warmup, and with SGLANG_RUST_SERVER one line says so. A failed capture memory check is logged once per key on every rank with the free, needed, reserved and allocated memory. When PyTorch caches enough unused memory, the cache is released once for that key (at most 32 times per process) and the check repeats. /get_server_info reports memory_refusals, last_memory_refusal, cache_releases, captured_keys, max_batch_size and the allocator's num_alloc_retries / num_ooms. The draft-extend tail is registered in the kernel inventory (Triton, CUDA only). Also kernel tests for chain verify sampling, the draft-extend tail and the relay scatter.
4 of 5 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Top-k-one speculative verification performs argmax, output initialization, and greedy chain verification in separate operations. Fuse final argmax reduction with chain verification for eligible large-vocabulary CUDA logits.
Modifications
Preserve lowest-index tie breaking and NaN behavior, support draft argmax with unit probabilities, and retain other sampling/backend paths.
This PR is independently based on
main. QSA KV preparation is tracked separately in #40972.Accuracy Tests
Added focused tests for this change:
Independent branch test rerun: 36 passed on NVIDIA B200. Each PR was tested from its own
main-based source snapshot.Combined implementation validation with official
sgl-eval0.1.0, temperature 1, top-p 1, maximum 131072 tokens, 8 repeats over 30 questions, concurrency 128, thinking enabled, no acceptance simulation:Real MTP acceptance length was approximately 2.146, including the bonus token, estimated by weighting the complete rounded log intervals by request count.
Audited all 8 prediction files per mode against native metrics, generation settings and token totals. These are integration results with the optimization series enabled, not an isolated accuracy A/B for this PR.
Speed Tests and Profiling
Benchmark scope: the E2E results below were measured with the full optimization series combined. Reproducing that implementation requires all PRs listed below, including this PR. They are not standalone performance results for this PR.
Full optimization series used for integration measurements
Each split PR is independently based on
main; the list above describes the combined benchmark source, not a required merge order. The combined AIME26 results in Accuracy Tests also use this full series (with acceptance simulation disabled).Evidence specific to this change
An isolated batch-size-1 greedy verify microbenchmark decreased from approximately 8.38 to 4.50 microseconds. The corresponding cumulative E2E stage did not establish a stable throughput improvement.
The combined E2E table below must not be used to infer this PR's individual throughput gain.
Full-series E2E performance (all PRs above combined)
Hardware and workload: NVIDIA B200 x4 / TP4,
nvidia/Qwen3.8-Flash-Next-NVFP4, 8192 input / 1024 output tokens, 5 requests at concurrency 1, plan stream disabled:Forward occupancy was collected with SGLang's device timer. Per-run values were 98.90%, 98.91%, and 98.89% for normal decode, and 97.75%, 97.79%, and 97.80% for NEXTN. These are full-series integration measurements with plan stream disabled; the NEXTN runs use simulated acceptance length 3.3, as in the performance table.
Every run completed 5 requests and generated 5120 output tokens. Decode throughput is 1000 / mean TPOT in milliseconds; the table reports the median of three runs. NEXTN performance uses simulated acceptance length 3.3. These measurements do not establish an isolated main-versus-this-PR speedup.
Checklist
CI States
Latest PR Test (Base): ❌ Run #36056870097
Latest PR Test (Extra): ❌ Run #36056869658
Latest PR Test (AMD ROCm 10): ❌ Run #36056869969