Skip to content

[Spec] Fix Dspark and Dflash state divergence across TP rank - #33614

Merged
hnyls2002 merged 23 commits into
sgl-project:mainfrom
JackZeng0208:fix-dspark-tp-sync
Aug 30, 2026
Merged

hnyls2002 merged 23 commits into
sgl-project:mainfrom
JackZeng0208:fix-dspark-tp-sync

Conversation

@JackZeng0208

@JackZeng0208 JackZeng0208 commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

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:

  1. The draft Markov chain samples proposal tokens step by step (with in-graph philox noise introduced since [Spec] Support sampling in the DSPARK graph-folded draft proposal #33298)
  2. Target verify derives correct_len / bonus / cap_trim_lens
  3. Prefill samples next_token_ids

In 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_lens handling in dspark_planner.py and the EAGLE proposals in #29003 and #31478:

  1. Broadcast every Markov proposal token in both eager and graph-folded sampling. Sync only at the end of a draft block is too late because each step depends on the previous token.
  2. Broadcast correct_len, bonus, and cap_trim_lens in the eager verify path and the captured verify epilogue, before finalization, token output, or KV commit. Separate all_reduce(MIN) calls could combine fields from different ranks.
  3. Broadcast prefill next_token_ids before entering the draft path.

DsparkTpSync in dspark_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 uses attn_tp_group, matching verify_lens_broadcast_group in original dspark_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 + 3 broadcasts of a [bs] tensor per decode step, all within the captured graphs.

Accuracy Tests

All tests below on following configs:

  • 2 DGX Sparks
  • TP=2
  • DeepSeek-V4-Flash-0731 with 131k context
  • Decode CUDA graphs and folded sampling enabled
  • Add Fix DSPARK SM120 decode dispatch for non-instantiated topk widths #33407 code. It fixes an unrelated bug for SM12x sparse-MLA topk dispatch crash at boot. Dspark + DeepSeek V4 Flash cannot start on DGX Sparks without it.
  • NCCL 2.30.7 for SGLang's PyNCCL communicator via 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 runpython3 -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=true and max_tokens=512 (30 at temp=1, then 12 at temp=0) and probed /health_generate after 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 with finish_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:

Requests Output throughput TTFT mean/median/P90 (ms) TPOT mean/median/P90 (ms)
24 31.9 tok/s 157.6 / 153.7 / 178.6 30.8 / 31.7 / 35.5

Comparing with pre-fix TP baseline is impossible because it cannot survive under such load.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ⏳ Run #33282599713
Latest PR Test (Extra): ❌ Run #33282599449
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33282599743

@hnyls2002 hnyls2002 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread python/sglang/srt/speculative/dspark_components/dspark_verify.py Outdated
@JackZeng0208

Copy link
Copy Markdown
Contributor Author

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.

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.

ShangmingCai added a commit that referenced this pull request Aug 6, 2026
@JackZeng0208

JackZeng0208 commented Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

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.

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.

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.

pair temp build tok/s mean TPOT (ms) mean accept len mean TPOT * mean accept length
1 1.0 unpatched 32.100 30.539 2.6427 80.704
1 1.0 patched 31.922 30.728 2.6315 80.859
2 0.0 unpatched 36.260 26.945 3.0005 80.849
2 0.0 patched 37.178 26.275 3.0974 81.385
3 0.0 unpatched 36.315 26.910 2.9784 80.148
3 0.0 patched 38.231 25.524 3.1708 80.932

Since acceptance length varies between different run. So for this fixed-shape, batch_size=1 workload, I use mean TPOT * mean accept length to estimates time per verifier invocation.

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 0.492 ms / 7 = 70.2 us per broadcast on this two-node fabric.

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.

@JackZeng0208

Copy link
Copy Markdown
Contributor Author

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.

@Han-xin58

Copy link
Copy Markdown

I have also encountered the same problem, please work hard to fix it as soon as possible

@icewool

icewool commented Aug 10, 2026

Copy link
Copy Markdown

This issue still exists in the latest version 0.5.17. Please work hard to fix it as soon as possible.

@Han-xin58

Copy link
Copy Markdown

@JackZeng0208 Boss, come on! We're waiting for you

@JackZeng0208

Copy link
Copy Markdown
Contributor Author

Finally, for the accuracy testing, I test the 200 GSM8k questions, 5-shot, temperature 0, top-p=1, max output=512, one evaluator thread:

Mode Score Evaluation latency Output tokens Output throughput
target-only 0.960 838.544 s 16943 20.205 token/s
Dspark fixed 0.960 492.706 s 17078 34.662 token/s

@hnyls2002

hnyls2002 commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

/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

@sgl-project sgl-project deleted a comment from github-actions Bot Aug 29, 2026
@github-actions

github-actions Bot commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

Results for /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:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dflash.py
cd test/ && python3 registered/spec/dflash/test_dflash.py

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_dflash_31b_extra.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/dspark/test_dspark_draft_path_default.py
cd test/ && python3 registered/spec/dspark/test_dspark_scheduler.py

@hnyls2002

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@hnyls2002
hnyls2002 merged commit f60bc73 into sgl-project:main Aug 30, 2026
79 of 103 checks passed
@JackZeng0208
JackZeng0208 deleted the fix-dspark-tp-sync branch August 30, 2026 00:31
@JackZeng0208

Copy link
Copy Markdown
Contributor Author

@hnyls2002 Thanks so much for the modification and review!

@hnyls2002

Copy link
Copy Markdown
Collaborator

@JackZeng0208 Can you help to use this config to do the bisection for this hanging issue?

DSPARK_MEM = 1 # gates both graph capture and folded sampling
DSPARK_DRAFT_GREEDY = 2
DSPARK_DRAFT_SAMPLE = 3
DSPARK_DRAFT_MULTINOMIAL = 4
DSPARK_GRAPH_SAMPLE = 5 # in-graph philox, redrawn per replay
DSPARK_GRAPH_GREEDY = 6
DSPARK_PLAN = 7
DSPARK_ACCEPT_GREEDY = 8
DSPARK_ACCEPT_SAMPLE = 9
DSPARK_ACCEPT_GRAPH = 10
DSPARK_TARGET = 11
# -- DFlash --
DFLASH_MEM = 12
DFLASH_SELECTOR = 13
DFLASH_ACCEPT_SAMPLE = 14
DFLASH_ACCEPT_GREEDY = 15
DFLASH_TARGET = 16

@JackZeng0208

Copy link
Copy Markdown
Contributor Author

@JackZeng0208 Can you help to use this config to do the bisection for this hanging issue?

DSPARK_MEM = 1 # gates both graph capture and folded sampling
DSPARK_DRAFT_GREEDY = 2
DSPARK_DRAFT_SAMPLE = 3
DSPARK_DRAFT_MULTINOMIAL = 4
DSPARK_GRAPH_SAMPLE = 5 # in-graph philox, redrawn per replay
DSPARK_GRAPH_GREEDY = 6
DSPARK_PLAN = 7
DSPARK_ACCEPT_GREEDY = 8
DSPARK_ACCEPT_SAMPLE = 9
DSPARK_ACCEPT_GRAPH = 10
DSPARK_TARGET = 11
# -- DFlash --
DFLASH_MEM = 12
DFLASH_SELECTOR = 13
DFLASH_ACCEPT_SAMPLE = 14
DFLASH_ACCEPT_GREEDY = 15
DFLASH_TARGET = 16

No problem, I will do it now

kediwu0331 pushed a commit to Zhylkaaa/sglang that referenced this pull request Aug 30, 2026
…ject#33614)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
@JackZeng0208

Copy link
Copy Markdown
Contributor Author

@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. SGLANG_SPEC_TP_SYNC=off deadlocks after 2.2 min. _INIT only still deadlocks after 18.5 min. The memory sync doesn't affect any decision on this setup: both ranks stay at around 23 to 27 GB free, well above the 1 GB thresholds.

All configs that keep the token syncs enabled run clean: _ALL, _INIT | _RNG, _ALL without DSPARK_MEM, and _ALL without each of DSPARK_GRAPH_SAMPLE, DSPARK_ACCEPT_SAMPLE, or DSPARK_TARGET individually. That's 195 min total with no hangs.

The hang happends only when DSPARK_GRAPH_SAMPLE, DSPARK_ACCEPT_SAMPLE, and DSPARK_TARGET are all disabled.

Dflash bisections are still testing. It may need some extra time.

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 31, 2026
…ject#33614)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
nzr-niu pushed a commit to nzr-niu/sglang that referenced this pull request Sep 1, 2026
…ject#33614)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
RolaoDenthu pushed a commit to RolaoDenthu/sglang that referenced this pull request Sep 1, 2026
…ject#33614)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
gilfordting added a commit to gilfordting/sglang that referenced this pull request Sep 8, 2026
… 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>
gilfordting added a commit to gilfordting/sglang that referenced this pull request Sep 8, 2026
…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>
gilfordting added a commit to gilfordting/sglang that referenced this pull request Sep 8, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

9 participants