[Sampling] Deterministic top-p / top-k renorm by default (fixes TP rank divergence in DFlash/DSpark) - #38565
[Sampling] Deterministic top-p / top-k renorm by default (fixes TP rank divergence in DFlash/DSpark)#38565gilfordting wants to merge 1 commit into
Conversation
451ff7b to
bad0bb3
Compare
|
/rerun-test registered/kernels/ops/moe/test_renorm.py registered/sampling/test_sampling_mask.py registered/core/test_basic_sanity_dflash.py registered/core/test_basic_sanity_dspark.py registered/core/test_basic_sanity_eagle3.py registered/spec/dflash/test_dflash.py registered/spec/eagle/test_eagle_reject_sampling.py registered/spec/eagle/test_spec_eagle.py registered/spec/dspark/test_dspark_kernel_parity.py |
|
Results for 🚀 🚀 🚀 |
gilfordting
changed the title
bad0bb3 to
f063d4c
Compare
|
/rerun-group unit/mem_cache |
|
/rerun-group unit/observability |
|
/rerun-test test_metrics.py test_radix_cache_hit.py test_unified_radix_cache_kl_full.py test_hicache_variants.py test_hicache_storage.py |
|
Results for 🚀 🚀 🚀 🚀 ⛔ |
|
Results for 🚀 |
|
Results for 🚀 🚀 🚀 |
|
/rerun-test registered/kernels/ops/moe/test_renorm.py registered/sampling/test_sampling_mask.py registered/core/test_basic_sanity_dflash.py registered/core/test_basic_sanity_dspark.py registered/core/test_basic_sanity_eagle3.py registered/spec/dflash/test_dflash.py registered/spec/eagle/test_eagle_reject_sampling.py registered/spec/eagle/test_spec_eagle.py registered/spec/dspark/test_dspark_kernel_parity.py |
|
Results for 🚀 🚀 🚀 |
…nk divergence in DFlash/DSpark) flashinfer's default top_p_renorm_probs (AIR radix, >= 0.6.7) and top_k_renorm_probs (radix multi-CTA) accumulate float sums with atomicAdd, so two calls on byte-identical input return probabilities that differ in the last bits. Every TP rank runs them independently on the same logits in speculative verification (DFlash, DSpark, EAGLE) and in the sampler's min_p path; a last-bit gap flips a rejection-sampling coin or the bonus token on one rank only, the per-rank radix/KV caches drift, and a later prefix match deadlocks an NCCL collective (sgl-project#33549, sgl-project#33289; sgl-project#33614 is the broadcast workaround). Add sglang.srt.layers.sampling_renorm with top_p_renorm_prob / top_k_renorm_prob that default to deterministic kernels: flashinfer's integer-histogram AIR (is_deterministic=True) for top-p, and the single-CTA kernel already compiled into sgl_kernel for top-k. Route the three call sites (sampler, dflash_utils, eagle_utils) through it. New `deterministic` kwarg and SGLANG_RENORM_DETERMINISTIC env var opt back into the faster kernels. Add regression tests asserting bit-identical output across repeated calls. The change lives in sglang rather than the sgl_kernel wrappers because sglang-kernel ships as a pinned prebuilt wheel; a wrapper change would not be testable in CI or reach users until the next kernel release. Measured: TP=2 DFlash with top_p=0.9 and 32 streams wedged within 4 min on H100 and B300 three times out of three; with deterministic renorm it ran 30 min, 21k requests, zero cross-rank divergence. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
f063d4c to
e8ee82d
Compare
|
/rerun-test registered/kernels/ops/moe/test_renorm.py registered/sampling/test_sampling_mask.py registered/core/test_basic_sanity_dflash.py registered/core/test_basic_sanity_dspark.py registered/core/test_basic_sanity_eagle3.py registered/spec/dflash/test_dflash.py registered/spec/eagle/test_eagle_reject_sampling.py registered/spec/eagle/test_spec_eagle.py registered/spec/dspark/test_dspark_kernel_parity.py |
|
Results for 🚀 🚀 🚀 |
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
1 participant
Motivation
Fixes the root cause behind the DFlash/DSpark TP>1 deadlocks in #33549 and #33289 (the ones #33614 works around by broadcasting rank 0's decisions).
Root cause.
sgl_kernel.top_p_renorm_probsandsgl_kernel.top_k_renorm_probsforward to flashinfer's default kernels, and both of those are non-deterministic run to run on byte-identical input:top_p_renorm_probs(flashinfer >= 0.6.7, [feat] Add air top-p algorithm flashinfer-ai/flashinfer#2752) dispatches rows longer than 2048 entries to the AIR radix kernel. Its defaultis_deterministic=Falsepath accumulates the bucket histogram with floatatomicAdd(include/flashinfer/air_top_p.cuh). sglang never passesis_deterministic.top_k_renorm_probsdispatches toRadixTopKRenormProbMultiCTA, which accumulates the kept-mass sum with floatatomicAdd(include/flashinfer/topk.cuh). flashinfer has no deterministic option for it.Every TP rank runs these kernels independently on the same logits. In DFlash/DSpark non-greedy verification (
build_dflash_verify_target_probs->tree_speculative_sampling_target_only) the output is compared against a uniform coin per draft position, and the bonus token is sampled fromrelu(q - p)over the same tensor. When a coin lands inside the few-ULP gap between two ranks'target_probs, one rank alone accepts/rejects a draft token or picks a different bonus token. Its radix/KV state silently drifts from the other ranks; a later turn on that conversation prefix-matches to different lengths per rank and the ranks wedge in an NCCL collective. EAGLE calls the same kernels but broadcastspredict/accept_indexfrom rank 0, which is why only DFlash/DSpark hang. Themin_ppath of the main sampler has the same exposure for committed tokens (see the comment inSampler._sync_token_ids_across_tpabout "rare instances where they are not deterministic": this is one of them).Measurements (Modal,
lmsysorg/sglang:v0.5.19-cu130, flashinfer 0.6.18, also reproduced on the v0.5.16 + flashinfer 0.6.15 stack; identical on H100, B200 and B300).Kernel-level, 200 repeated calls on the same input (256 rows x 163840 vocab, top_p 0.95 / top_k 40):
top_p_renorm_probsdefault (AIR)top_p_renorm_probs(is_deterministic=True)top_k_renorm_probsdefault (radix multi-CTA)top_k_renorm_probstree_speculative_sampling_target_only, flashinfertop_k_top_p_sampling_from_probs,sampling_from_probs,min_p_sampling_from_probs(seeded)The top-p noise is entropy-gated (bit-stable on 4,096 peaky real Llama-3.1-8B decode rows, fires on flat rows); the top-k noise is not (87 to 99 of 99 calls differ even on peaky rows).
End to end, TP=2, DFlash, Llama-3.1-8B +
z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat,temperature=1.0,top_p=0.9, 32 concurrent multi-turn streams, fp8 KV, with a per-verify-step all-gather of fingerprints of{rng state, prefix_lens, draft_tokens, target_logits, target_probs, accept_len, bonus, commit_lens, out_tokens}so the first differing quantity is named:target_probsdiffers (logits, prefix_lens, RNG identical) ->bonus/out_tokensdiffer on one rank -> next step'sdraft_tokensdiffer -> wedge. 0.4 to 0.7 % oftop_p_renorm_probscalls returned different bytes when called twice on the same input on the same ranktop_p=1.0top_p=0.9,top_k=40[rows, 40]uses the single-CTA kernelThe earlier failure to reproduce this synthetically was because test traffic used
top_p=1.0, which skips the renorm kernel entirely (need_top_p_samplingis False).Modifications
python/sglang/srt/layers/sampling_renorm.pyexposingtop_p_renorm_prob/top_k_renorm_probwith adeterministic: Optional[bool] = Noneargument that defaults to the env varSGLANG_RENORM_DETERMINISTIC(default on):top_p_renorm_probs(..., is_deterministic=True)(integer-histogram AIR; the flag exists since 0.6.7 and sglang pins newer). Without flashinfer (MUSA) the single-CTA kernel compiled intosgl_kernel, which is deterministic.top_k_renorm_probsalready compiled intosgl_kernel(flashinferrenorm.cuat the pinned commit; fixed-order block reductions). flashinfer has no deterministic option for its radix top-k.deterministic=False(orSGLANG_RENORM_DETERMINISTIC=0) keeps today's kernels.sampler.py,speculative/dflash_utils.py,speculative/eagle_utils.py: import the renorm functions from that module instead ofsgl_kernel. These are the only three call sites.environ.py:SGLANG_RENORM_DETERMINISTICdocumented next to the other determinism knobs.test/registered/kernels/ops/moe/test_renorm.py: existing correctness tests now run through the new module, plus two regression tests that call each renorm 30 times on a flat 256 x 128256 distribution and assert bit-identical output (the same check against the default flashinfer kernels fails 30 times out of 30).The fix lives in
sglangrather than in thesgl_kernelwrappers on purpose:sglang-kernelships as a pinned prebuilt wheel, so a wrapper change there would neither be testable in CI nor reach users until the next kernel release (a first version of this PR did that and CI ran the tests against the 0.4.6.post1 wheel). Making the same default insgl_kernellater is a natural follow-up, as is anis_deterministicoption for flashinfer's radix top-k so the fast kernel can be used again.This is complementary to #33614, not a replacement. #33614 makes replicated state agree whatever the source of disagreement; this PR removes the source, so rank 0's own output and every other consumer of
target_probs(logprobs, sampling masks) become reproducible as well.Accuracy Tests
Numerically the deterministic variants compute the same renormalization (the existing
test_renorm.pycorrectness tests pass unchanged; deterministic mode differs from the default only in accumulation order). The 30-minute TP=2 DFlash run above produced 21,294 requests with zero errors and the same accept length as the unfixed runs before they wedged.Speed Tests and Profiling
Microbenchmark, B200, flashinfer 0.6.18, 50 iterations after warmup, batch x vocab:
is_deterministic=True(new default)Cost per decode step. The renorm runs once per speculative verify step over
batch_size x draft_tokensrows, and in the plain sampler only on themin_ppath (the fusedtop_k_top_p_sampling_from_probskernel, already deterministic, is untouched). For a representative production shape (32 requests x 8 draft tokens = 256 rows, 163840 vocab, Kimi K2.6 NVFP4 on 4 GPUs, 25 to 40 ms per step):top_p < 1, notop_k(the common case; dense path, top-p only)top_kand non-top_krequests (dense path, top-k then top-p)top_k(sparsetorch.topkpath, renorm on[rows, k])top_p = 1.0Today's whole sampling stage (softmax, renorm, accept) is roughly 1 to 2 ms of such a step, so this is a change within that slice, not to the model forward. End to end on the TP=2 DFlash run above (H100, 32 streams,
top_p=0.9) throughput was ~710 requests/min with the deterministic kernels versus ~715 requests/min for the default kernels in the minutes before they wedged, i.e. unchanged within noise. Deployments that want the faster kernels and do not run TP>1 speculative decoding can setSGLANG_RENORM_DETERMINISTIC=0.Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #34291766245
Latest PR Test (Extra): ❌ Run #34291765978
Latest PR Test (AMD ROCm 10): ❌ Run #34291766176