[Speculative] Seed rejection-sampling draft proposals for deterministic inference by ajit283 · Pull Request #33395 · sgl-project/sglang

ajit283 · 2026-08-03T14:05:17Z

Motivation

Classic EAGLE/NEXTN rejection sampling draws each draft proposal with fast_sample(). Although target sampling and verify-side rejection coins support request-local seeds, this proposal draw used the global PyTorch RNG (exponential_()), so --enable-deterministic-inference still changed outputs with batch shape and CUDA graph replay.

This is the draft-worker follow-up documented in #30822. Before this change, Qwen3.5-9B NEXTN rejection sampling produced 12 unique outputs in 12 changing-batch trials with a fixed sampling seed.

Modifications

Add an optional deterministic path to fast_sample() keyed by:
- request sampling seed
- absolute sequence position
- draft step / MTP layer
- vocabulary entry
Domain-separate draft proposal randomness from target sampling and verify-side rejection coins.
Convert MurmurHash32 output to open-interval uniforms with midpoint binning, then use exponential-race categorical sampling in float64.
Preserve the existing RNG-backed fast path when deterministic inference is disabled.
Plumb request seeds and positions through eager and CUDA-graph EAGLE/NEXTN workers.
Support both ordinary and multi-layer EAGLE, including the single-CUDA-graph multi-layer draft-extend path.
Keep deterministic rejection sampling guarded on non-CUDA-compatible devices because this implementation uses a Triton hash kernel.
Add sampler tests for replay, batch permutation/splitting, seed/position/step separation, hash endpoints, CUDA graph updates, and categorical distribution sanity.
Add dedicated deterministic E2E smoke tests without changing the existing unseeded rejection-sampling and chain-MTP test configurations.
Add a reproducible seeded/unseeded sampling microbenchmark.

Accuracy Tests

Regular EAGLE/NEXTN

Qwen3.5-9B, TP2, NEXTN, rejection sampling, temperature 1.0, sampling seed 42:

Version	Changing-batch trials	Unique outputs
Before	12	12
After	12	1

The deterministic harness also passed changing batch sizes, prefix batching, bitwise logprob comparison, and radix-cache consistency.

Post-rebase validation on B200:

7 passed, 16 warnings, 2 subtests passed
Total samples: 6, Unique samples: 1
OK

The unit suite includes 100,000 distinct seeded draws from [0.1, 0.2, 0.7] and verifies that observed frequencies are within 0.01 absolute tolerance.

Multi-layer EAGLE

stepfun-ai/Step-3.5-Flash, TP8 on 8x B200, FA4, three MTP layers, single-CUDA-graph draft extend:

Total samples: 6, Unique samples: 1
MULTI_LAYER_EAGLE_REJECTION_SAMPLING_IS_DETERMINISTIC

Speed Tests and Profiling

B200, vocabulary size 163,840, CUDA-graph replay:

Batch	Unseeded	Seeded	Slowdown
1	31.9 us	47.2 us	1.48x
8	38.0 us	63.5 us	1.67x
32	54.5 us	131.9 us	2.42x
128	138.9 us	544.1 us	3.92x

Peak allocated memory for batch 128 and vocabulary 163,840:

Input probabilities:       80.00 MiB
Baseline allocated:        80.00 MiB
Peak allocated:           321.25 MiB
Incremental seeded peak:  241.25 MiB

The deterministic mode is opt-in. A fused hash/exponential/argmax kernel can be considered separately if end-to-end profiling shows that the materialized hash and score matrices materially affect serving throughput.

Checklist

Format code and pass Ruff / git diff --check.
Add unit and registered E2E tests.
Preserve existing unseeded rejection-sampling and chain-MTP coverage.
Provide accuracy, latency, and peak-memory measurements.
Follow the SGLang code style guidance.
No documentation update is needed; this fixes an existing flag combination without adding user-facing options.

CI States

Latest PR Test (Base): ✅ Run #37454576286
Latest PR Test (Extra): ❌ Run #37454575883
Latest PR Test (AMD ROCm 10): ❌ Run #37454576246

gemini-code-assist · 2026-08-03T14:05:21Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

ajit283 · 2026-08-03T14:11:23Z

@zyzshishui This PR follows up on the deterministic draft-proposal sampling gap documented in #30822. I would appreciate your review, particularly of the RNG stream separation and EAGLE/NEXTN seed plumbing.

nvpohanh · 2026-08-12T05:31:26Z

/tag-and-rerun-ci

nvpohanh · 2026-08-12T08:38:35Z

@ajit283 could you rebase/merge with main?

nvpohanh · 2026-08-13T05:58:58Z

/rerun-failed-ci

nvpohanh · 2026-08-13T13:01:53Z

All NV pipelines have passed

nvpohanh · 2026-08-18T11:21:44Z

@Fridge003 do you know how is in charge of reviewing deterministic features? thanks!

nvpohanh · 2026-08-24T03:49:17Z

@Qiaolin-Yu Could you help us to review this? Thanks!

nvpohanh · 2026-08-27T12:33:26Z

@ajit283 could you fix the conflicts? thanks!

…-sampling

nvpohanh · 2026-08-27T14:20:48Z

[by Codex] The required PR Test Base failure is caused by this PR, not an intermittent CI issue. test_chain_read_ratchet.py flags two new reads of the resolved enable_deterministic_inference setting through borrowed server_args objects:

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py:205
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py:684

Please read the resolved config bag instead, using get_exec().deterministic.enable_deterministic_inference (and import get_exec from sglang.srt.runtime_context) at both sites.

Failed job: https://github.com/sgl-project/sglang/actions/runs/33074073313/job/98523815165

…-sampling

nvpohanh · 2026-09-01T05:37:00Z

[by Codex] @ajit283, the relevant CI failure appears to be caused by this PR: https://github.com/sgl-project/sglang/actions/runs/33450499018/job/99686771962. The PR-added test/registered/spec/eagle/test_eagle_reject_sampling_deterministic.py failed its deterministic token-count expectation: expected [2] but got [1], and the retry failed the same way. Could you please fix it? Thanks!

nvpohanh

[by Claude Code] Review of the deterministic draft-proposal sampling change. Inline comments below; each is labeled with severity and confidence.

nvpohanh · 2026-10-01T06:43:48Z

@ajit283 could you address the comments? thanks

nvpohanh

[by Claude Code] Re-review at f004e95. The latest commit fixes three of my four earlier threads (draft-extend position key, step/seed aliasing, and the double draft_probs clone), so I am resolving those threads. The get_spec() thread on eagle_draft_cuda_graph_runner.py line 209 is still open. One new inline comment below.

…unner Gate the draft CUDA graph sampling_seed buffer on the resolved get_spec().speculative_use_rejection_sampling, matching the draft_probs buffer above it, instead of the raw server_args field. Drop test_murmur_hash32_cpu_distinguishes_draft_step_column_domains: it only checks murmur_hash32 and never calls fast_sample, so it cannot catch a regression in the draft-step column offset. The GPU test test_seed_position_and_draft_step_select_separate_streams covers it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-02T12:41:04Z

/tag-and-rerun-ci

nvpohanh · 2026-10-02T13:55:01Z

[by Claude Code] @ajit283, the relevant CI failure appears to be caused by this PR: https://github.com/sgl-project/sglang/actions/runs/37007979535/job/110841646756. The PR-added test/registered/unit/spec/test_eagle_draft_sampling.py::TestEagleDraftSamplingPosition::test_draft_extend_uses_post_verify_position fails with AttributeError: 'types.SimpleNamespace' object has no attribute 'prepared_draft_extend_inputs' at eagle_worker_v2.py:1041 (_draft_extend_for_decode), and the retry fails the same way. After the latest main merge (#41175), _draft_extend_for_decode reads batch_result.prepared_draft_extend_inputs, but the test's mocked batch_result does not define it. Adding prepared_draft_extend_inputs=None to that SimpleNamespace should fix it. Could you please fix it? Thanks!

sgl-project#41175 made `_draft_extend_for_decode` read `batch_result.prepared_draft_extend_inputs`, so the mocked batch result in test_eagle_draft_sampling.py raised AttributeError after merging main. Set it to None so the test takes the CPU fallback path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-02T14:08:07Z

[by Claude Code] @ajit283, following up on my comment above: I pushed c1b3d4c to your branch. It adds prepared_draft_extend_inputs=None to the mocked batch_result in test/registered/unit/spec/test_eagle_draft_sampling.py, so no action is needed from you for that failure. Please pull before pushing again. CI is rerunning: https://github.com/sgl-project/sglang/actions/runs/37017522731. I also resolved the two remaining review threads, since commit 3db4208 already addressed them. Thanks!

nvpohanh · 2026-10-02T15:32:07Z

/rerun-failed-ci

Resolve the sample_draft_proposal conflict by keeping main's sampling_softmax while passing the PR's seeded draft-sampling arguments to fast_sample. main switched fast_sample's single-sample argmax to the FP32-only row_argmax kernel; keep torch argmax for the seeded path's float64 scores. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-05T05:33:23Z

/rerun-failed-ci

nvpohanh · 2026-10-05T06:43:53Z

/rerun-failed-ci

nvpohanh · 2026-10-06T00:20:11Z

/rerun-failed-ci

Picks up the latest main for a fresh CI run after a flaky DSV4-Flash H200 accept-length failure (https://github.com/sgl-project/sglang/actions/runs/37251609204/job/111751683459). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-06T06:32:54Z

[by Claude Code] @ajit283, the relevant CI failure appears to be caused by this PR: https://github.com/sgl-project/sglang/actions/runs/37420455068/job/112128680788. After merging the latest main, test/registered/unit/spec/test_block_verification_args.py::test_unsupported_configurations fails with ValueError not raised for enable_deterministic_inference=True. That test came from #42297 (block verification), which forces rejection sampling on and expects deterministic inference to be rejected. This PR removes the "--speculative-use-rejection-sampling is incompatible with --enable-deterministic-inference" check that the test relied on. Could you either keep block verification incompatible with deterministic inference or update that test? Thanks!

Block verification forces rejection sampling, which this PR makes deterministic via seeded draft and verify draws, so --enable-deterministic-inference is no longer rejected on CUDA. Update the test from sgl-project#42297 accordingly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-06T07:43:58Z

[by Claude Code] Follow-up: I pushed f0e205e to address the failure above with a test-only change. Block verification reuses this PR's seeded draft and verify draws, so test_block_verification_args.py now expects --enable-deterministic-inference to be accepted on CUDA instead of rejected. Please revert it if you'd rather keep block verification incompatible with deterministic inference.

nvpohanh · 2026-10-06T08:17:58Z

[by Claude Code] The relevant CI failure (https://github.com/sgl-project/sglang/actions/runs/37431477801/job/112163452489) is not caused by this PR. It is a known breakage on main that #42722 will fix. CI for this PR is blocked until that PR merges.

Picks up sgl-project#42750, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/37431477801/job/112163452489 (test_dsa_metadata_replay TARGET_VERIFY mismatch). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-06T16:45:26Z

/rerun-failed-ci

nvpohanh · 2026-10-06T18:15:24Z

/rerun-failed-ci

nvpohanh · 2026-10-06T21:15:25Z

/rerun-failed-ci

nvpohanh · 2026-10-06T22:15:34Z

[by Claude Code] All NV pipelines have passed.

nvpohanh · 2026-10-07T01:46:07Z

@Qiaolin-Yu could you help to review this? Thanks!

ajit283 requested review from Qiaolin-Yu, Ying1123, hnyls2002 and merrymercy as code owners August 3, 2026 14:05

github-actions Bot added the speculative-decoding label Aug 3, 2026

ajit283 force-pushed the deterministic-draft-sampling branch from cf0c551 to 68b6a65 Compare August 3, 2026 14:16

ajit283 force-pushed the deterministic-draft-sampling branch from 68b6a65 to 791f50f Compare August 11, 2026 13:22

fix(spec): seed rejection-sampling draft proposals

66bae0b

ajit283 force-pushed the deterministic-draft-sampling branch from 791f50f to 66bae0b Compare August 11, 2026 13:22

nvpohanh added deterministic Issues on deterministic inference/kernels bug Something isn't working labels Aug 12, 2026

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Aug 12, 2026

Merge branch 'main' into deterministic-draft-sampling

64a7c56

kpham-sgl assigned Qiaolin-Yu Aug 19, 2026

Merge remote-tracking branch 'upstream/main' into deterministic-draft…

880575b

…-sampling

Ajit Mistry added 2 commits September 1, 2026 01:22

fix(spec): read deterministic setting from runtime config

ae0b97e

Merge remote-tracking branch 'upstream/main' into deterministic-draft…

60e62ae

…-sampling

Ajit Mistry added 2 commits September 4, 2026 00:07

fix(spec): preserve draft probabilities across graph replay

77dc4c7

Merge remote-tracking branch 'upstream/main' into pr-33395-token-fix

2b6b56c

nvpohanh reviewed Oct 1, 2026

View reviewed changes

fix(spec): separate deterministic draft streams

f004e95

nvpohanh reviewed Oct 2, 2026

View reviewed changes

Comment thread test/registered/cpu/test_sampling.py Outdated

nvpohanh and others added 2 commits October 2, 2026 05:31

Merge branch 'main' into deterministic-draft-sampling

67fb19a

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 2, 2026

nvpohanh added the bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) label Oct 6, 2026

nvpohanh mentioned this pull request Oct 6, 2026

[CI] Flaky tests and CI infrastructure failures seen while babysitting PRs #42752

Open

apex-mochen mentioned this pull request Oct 8, 2026

[HiCache] Avoid full-store scans for uncached file prefix lookups #42077

Open

5 tasks done

Conversation

ajit283 commented Aug 3, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Regular EAGLE/NEXTN

Multi-layer EAGLE

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Aug 3, 2026

Uh oh!

ajit283 commented Aug 3, 2026

Uh oh!

nvpohanh commented Aug 12, 2026

Uh oh!

nvpohanh commented Aug 12, 2026

Uh oh!

nvpohanh commented Aug 13, 2026

Uh oh!

nvpohanh commented Aug 13, 2026

Uh oh!

nvpohanh commented Aug 18, 2026

Uh oh!

nvpohanh commented Aug 24, 2026

Uh oh!

nvpohanh commented Aug 27, 2026

Uh oh!

nvpohanh commented Aug 27, 2026

Uh oh!

nvpohanh commented Sep 1, 2026

Uh oh!

nvpohanh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nvpohanh commented Oct 1, 2026

Uh oh!

nvpohanh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nvpohanh commented Oct 2, 2026

Uh oh!

nvpohanh commented Oct 2, 2026

Uh oh!

nvpohanh commented Oct 2, 2026

Uh oh!

nvpohanh commented Oct 2, 2026

Uh oh!

nvpohanh commented Oct 5, 2026

Uh oh!

nvpohanh commented Oct 5, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 6, 2026

Uh oh!

nvpohanh commented Oct 7, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

ajit283 commented Aug 3, 2026 •

edited by github-actions Bot

Loading