[Spec] Fix Dspark and Dflash state divergence across TP rank - #33614
Conversation
hnyls2002
left a comment
There was a problem hiding this comment.
Direction is right, and the fault-injection demo is a good way to pin the failure mode. Two comments inline.
Separately, the perf claim needs an A/B. "Comparing with pre-fix TP baseline is impossible because it cannot survive under such load" doesn't hold on its own terms: you note 2.28.9 wedges this workload with or without the patch, and every run here is on 2.30.7. An unpatched baseline on 2.30.7 is exactly the comparison that's available. gamma+3 broadcasts per decode step sit on the critical path and can't overlap with compute; that needs a number next to it.
| self._tp_group = tp_group | ||
| self._enabled = tp_group.world_size > 1 | ||
|
|
||
| def sync(self, tensor: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
dsa_indexer.py already solves this: _broadcast_indexer_topk_from_rank0 (:169-181) and _broadcast_indexer_topk_from_rank0_impl (:153-166) -- broadcast from rank 0 on a path that must also work inside a captured graph. Diffing the two is the fastest way to see what's missing here.
One fact that makes it concrete: under --enable-dp-attention with attn_tp_size < tp_size, attn_tp_group is built with use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP or enable_symm_mem (parallel_state.py:2422), and both are off by default.
Can this reuse that shape instead of adding a second one?
There was a problem hiding this comment.
Just updated it. DSpark now reuses the same graph-safe broadcast path as DSA. It's available in 43b609b and dspark_tp.py in fca0998
GroupCoordinator.broadcast_capture_safe in parallel_state.py handles the transport choice: eager calls use the normal process-group broadcast, while captured calls require PyNCCL and fail if it is unavailable. Both DSA top-k and DsparkTpSync.sync use this helper now. Also, the split attn_tp_group case is fixed. When attn_tp_size < tp_size, PyNCCL is now provisioned for the captured-broadcast users that need it. When attn_tp == tp, it still aliases the normal TP group, which builds PyNCCL by default.
Thanks for your helpful comments. I wrote the sentence "Comparing with pre-fix TP baseline is impossible because it cannot survive under such load" before I found NCCL issue. And I forgot to change it when updating the rest of the description. For the A/B test, I will update it ASAP. |
Here're the paried A/B results on the same 2 DGX sparks, using the same NCCL 2.30.7, server config, 24 fixed prompts (1 sentence question such as "Explain why the sky appears blue during the day in a concise paragraph"), and output length=256 token. Each pair is an adjacent unpatched-patched run.
Since acceptance length varies between different run. So for this fixed-shape, batch_size=1 workload, I use The patched vs. unpatched difference are +0.155, +0.536, and +0.784 ms, averaging +0.492 ms (0.61%). With gamma=4 (7 added broadcasts), this is equivalent to about This is just a derived e2e estimate, and I will provide the final A/B once all changes have been made based on your comments. |
|
Hi @hnyls2002, thanks so much for your review. I have responded and updated the code. The final accuracy and A/B testing are still running, and it will take some time (10+ hours). I will update these final testing results ASAP. |
|
I have also encountered the same problem, please work hard to fix it as soon as possible |
|
This issue still exists in the latest version 0.5.17. Please work hard to fix it as soon as possible. |
|
@JackZeng0208 Boss, come on! We're waiting for you |
|
Finally, for the accuracy testing, I test the 200 GSM8k questions, 5-shot, temperature 0, top-p=1, max output=512, one evaluator thread:
|
|
/rerun-test registered/core/test_basic_sanity_dspark.py registered/core/test_basic_sanity_dflash.py registered/spec/dflash/test_dflash.py registered/spec/test_gemma4_dflash_31b_extra.py registered/dcp/test_kimi_linear_dcp_dspark4.py registered/spec/dspark/test_dspark_draft_path_default.py registered/spec/dspark/test_dspark_scheduler.py |
|
Results for 🚀 🚀 🚀 🚀 🚀 |
|
/tag-and-rerun-ci |
|
@hnyls2002 Thanks so much for the modification and review! |
|
@JackZeng0208 Can you help to use this config to do the bisection for this hanging issue? sglang/python/sglang/srt/speculative/spec_tp_sync.py Lines 21 to 38 in 9a489f8 |
No problem, I will do it now |
…ject#33614) Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
|
@hnyls2002 Testing on 2 DGX Sparks with the same config as 9a489f8. The hang points to the token syncs, not the group-min memory probe. All configs that keep the token syncs enabled run clean: The hang happends only when Dflash bisections are still testing. It may need some extra time. |
…ject#33614) Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
…ject#33614) Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
…ject#33614) Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
… rank 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 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). Route both kernels to deterministic variants by default: flashinfer's integer-histogram AIR (is_deterministic=True) for top-p, and the single-CTA kernel already compiled into sgl_kernel for top-k. 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 (both fail on the previous defaults). 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>
…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>
…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>
Motivation
Related to: #33289
To fix #33289 bug, upgrade NCCL to the newest version (2.30.7).
When TP > 1, Dspark makes serveral sampling decisions on reach rank:
correct_len/bonus/cap_trim_lensnext_token_idsIn order to maintain the speedup, SGLang skips cross-rank sync of sampled tokens by default (I found the comment here: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/sampler.py#L496). So nothing forces these values to agree. Once one rank commits a different token or accept length, sequence lengths and KV state drift apart and a later collective deadlocks (demonstrated by fault injection in the Accuracy Tests section below).
Modifications
Broadcast rank 0's sampling decisions before they affect state. The idea follows the existing
verify_lenshandling indspark_planner.pyand the EAGLE proposals in #29003 and #31478:correct_len,bonus, andcap_trim_lensin the eager verify path and the captured verify epilogue, before finalization, token output, or KV commit. Separateall_reduce(MIN)calls could combine fields from different ranks.next_token_idsbefore entering the draft path.DsparkTpSyncindspark_tp.py(newly added file) uses the TP group's PyNCCL communicator. Process-group collectives cannot be captured in CUDA graphs, and using the model's existing TP communicator preserves collective ordering across eager and captured execution. With DP attention, it usesattn_tp_group, matchingverify_lens_broadcast_groupin originaldspark_planner.py. TP=1 is a no-op.The reason why I don't match RNG seeds is because ranks can still diverge if they consume different numbers of random values, and it does not cover non-RNG differences.
In addition, for TP > 1, this add
gamma + 3broadcasts of a[bs]tensor per decode step, all within the captured graphs.Accuracy Tests
All tests below on following configs:
SGLANG_NCCL_SO_PATH(torch's bundled 2.28.9 wedges this workload's graph/eager mix with or without this patch; a separate problem from this fix)For accuracy test, we run
python3 -m sglang.test.few_shot_gsm8k --num-questions 200(temp=0): accuracy 0.965, invalid 0.000. This is expected since there's no mathematical changes.To demonstrate the failure mode this patch addresses, we fault-injected a rank-local accept divergence into both builds: rank 1 lowers every accept length by one (
correct_len = torch.clamp(correct_len - 1, min=0), NCCL 2.30.7 in both _runs).Note: the fault reproduction code has been removed before submitting PR.
Without the fix, the server deadlocks before finishing startup warmup. With the fix, the same injection is harmless: rank 0's broadcast overwrites the divergent values, the server boots, and a temp=0 request returns the expected output.
As a stability check on the fix itself (its broadcasts run inside the captured decode graphs), I sent 42 sequential chat requests with
ignore_eos=trueandmax_tokens=512(30 at temp=1, then 12 at temp=0) and probed/health_generateafter every request with a 180s client timeout, so a stuck collective surfaces as a failed request instead of a silent hang. As the result, all 42 returned HTTP 200 withfinish_reason=length. 21504 tokens generated in total and mean acc_len during the temp=1 run was 2.86 (block_size=4). Same for temp=0, neither of them occur any error.Speed Tests and Profiling
Test on 24 natural prompts, 256 output tokens, concurrency=1, temp=1, same setup as above:
Comparing with pre-fix TP baseline is impossible because it cannot survive under such load.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ⏳ Run #33282599713
Latest PR Test (Extra): ❌ Run #33282599449
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33282599743