Fix(DSpark): Broadcast Verify Budget Across TP Ranks for Compact Graph Tier Alignment

[EanWang211123]

Motivation

In DSpark compact ragged verify with TP > 1, each TP rank independently computes verify_token_budget from locally resolved draft confidence (including overlap precompute via prepare_verify_budget). Per-request verify_lens are broadcast from TP0 after top-k scheduling, but graph_num_tokens (compact CUDA graph key / M tier) is derived from each rank’s local budget in schedule_layout() when _dynamic_graph_tier is enabled.

Draft confidence can differ slightly across TP ranks (floating-point / sharded draft paths). Ranks then pick different budgets and different graph tiers for the same step. Both ranks still participate in the same TP MoE/attention collectives with mismatched token shapes, which leads to a hang: CPU blocked in copy_done.synchronize() while GPUs stay busy.

Reproduction logs

1. Pre-broadcast CPS debug: local verify_lens diverge on the same request

SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER=1 logs before verify_lens broadcast; TP0/TP1 can disagree on the same req:

[2026-07-14 12:35:34 TP0] [DSPARK-CPS]   req=37 prefix=555 verify_len=4 sort_survival=[0.811, 0.720, 0.416, 0.311, 0.172, 0.061, 0.021]
[2026-07-14 12:35:34 TP1] [DSPARK-CPS]   req=37 prefix=555 verify_len=3 sort_survival=[0.810, 0.718, 0.417, 0.311, 0.173, 0.062, 0.021]

Similar divergence on other requests in the same step (e.g. req=34: TP0 verify_len=8, TP1 verify_len=6).

2. Graph tier mismatch on the same step (root cause for hang)

With SGLANG_LOG_DECODE_GRAPH_KEY=1, the same decode step replays different target verify graphs on TP0 vs TP1 (raw_bs=32 on both):

[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=29 prefix=561 verify_len=3 sort_survival=[0.741, 0.640, 0.341, 0.097, 0.019, 0.010, 0.002]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=30 prefix=756 verify_len=2 sort_survival=[0.467, 0.324, 0.293, 0.192, 0.133, 0.072, 0.030]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=31 prefix=804 verify_len=6 sort_survival=[0.984, 0.910, 0.905, 0.846, 0.487, 0.351, 0.214]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=32 prefix=974 verify_len=4 sort_survival=[0.896, 0.737, 0.437, 0.225, 0.095, 0.038, 0.007]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=33 prefix=629 verify_len=3 sort_survival=[0.647, 0.458, 0.311, 0.247, 0.213, 0.116, 0.045]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=34 prefix=584 verify_len=5 sort_survival=[0.748, 0.720, 0.600, 0.558, 0.303, 0.053, 0.021]
[2026-07-14 12:48:14 TP1] [DSPARK-CPS]   req=35 prefix=374 verify_len=2 sort_survival=[0.668, 0.325, 0.180, 0.089, 0.031, 0.008, 0.002]
[2026-07-14 12:48:14 TP0] Decode graph replay: worker=target key_size=160 (num_tokens) mode=TARGET_VERIFY raw_bs=32 slots=160
[2026-07-14 12:48:14 TP1] Decode graph replay: worker=target key_size=144 (num_tokens) mode=TARGET_VERIFY raw_bs=32 slots=144
[2026-07-14 12:48:14 TP0] Decode graph replay: worker=draft key_size=32 (bs) mode=TARGET_VERIFY raw_bs=32
[2026-07-14 12:48:14 TP1] Decode graph replay: worker=draft key_size=32 (bs) mode=TARGET_VERIFY raw_bs=32

TP0 target verify uses M=160; TP1 uses M=144. Draft graphs match (key_size=32), but target verify does not.

3. Hang symptom: symmetric CPU wait, GPUs still busy

py-spy on both scheduler ranks (TP0 pid 115005, TP1 pid 115007):

Process 115005: sglang::scheduler_TP0
Thread 115005 (active): "MainThread"
    cuEventSynchronize (libcuda.so.575.57.08)
    ...
    synchronize (torch/cuda/streams.py:254)
    process_batch_result_decode (scheduler_components/batch_result_processor.py:654)
    process_batch_result (scheduler.py:3492)
    ...
    event_loop_overlap (scheduler.py:1639)

Process 115007: sglang::scheduler_TP1
Thread 115007 (active): "MainThread"
    cuEventSynchronize (libcuda.so.575.57.08)
    ...
    synchronize (torch/cuda/streams.py:254)
    process_batch_result_decode (scheduler_components/batch_result_processor.py:654)
    process_batch_result (scheduler.py:3492)
    ...
    event_loop_overlap (scheduler.py:1639)

Both ranks stuck in result.copy_done.synchronize() (waiting for prior forward’s D2H on copy_stream).

nvidia-smi pmon at hang time:

# gpu         pid   type     sm    mem    enc    dec    jpg    ofa     fb   ccpm    command
# Idx           #    C/G      %      %      %      %      %      %     MB     MB    name
    4     115005     C     99      0      -      -      -      - 127086      0    sglang::schedul
    5     115007     C     99      0      -      -      -      - 127086      0    sglang::schedul

GPUs not idle; forward likely blocked inside mismatched TP collective / kernel path.

4. Typical server config when reproducing

SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER=1 \
SGLANG_LOG_DECODE_GRAPH_KEY=1 \
SGLANG_RAGGED_VERIFY_MODE=compact \
sglang serve \
  --speculative-algorithm DSPARK \
  --speculative-dspark-block-size 7 \
  --speculative-dspark-sps-table-path /workspace/sps-profiles/tp2_dspark_sps_2d.json \
  --speculative-dspark-align-verify-tokens-to-graph-tier \
  --tp 2 \
  --moe-runner-backend flashinfer_mxfp4 \
  ...

Modifications

• Add SpecTpSync.sync_cpu() to broadcast CPU decision tensors through the existing TP CPU process group while respecting the configured SpecTpSyncSite.
• Add _sync_verify_token_budget() to DSparkVerifyPlanner, using SpecTpSyncSite.DSPARK_PLAN to broadcast TP0’s scalar budget; retain -1 as the wire sentinel for None.
• Synchronize and write back verify_token_budget in both overlap and non-overlap paths before schedule_layout(), ensuring all TP ranks use the same budget for verify-lens scheduling, graph-tier alignment, and graph_num_tokens selection.
• Preserve the mainline DSPARK_PLAN synchronization of verify_lens, providing a second consistency boundary after budget-based scheduling.
• Add unit coverage verifying that non-zero TP ranks adopt TP0’s budget and that the None sentinel is preserved correctly.

Expected behavior after fix: for the same step, logs should show matching target graph keys, e.g.:

[TP0] Decode graph replay: worker=target key_size=160 (num_tokens) mode=TARGET_VERIFY raw_bs=32 ...
[TP1] Decode graph replay: worker=target key_size=160 (num_tokens) mode=TARGET_VERIFY raw_bs=32 ...

Cost: one CPU-side scalar broadcast per decode step when tp_size > 1 and budget scheduling is active.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

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 #33462766027
Latest PR Test (Extra): ❌ Run #33462765865
Latest PR Test (AMD ROCm 7.2): ❌ Run #33462766032

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Phoenix3334]

Independent B300 TP8 validation: O2 complete; strict eager boundary confirmed

I independently backported this PR's TP0 verify-budget broadcast logic onto SGLang baseline 692c5f7d and completed the full stability matrix on 8 × NVIDIA B30Z / B300, TP=8.

Stability results

Scheduling	Lifecycle	Result
no-overlap	same server	30/30
no-overlap	restart per workload	30/30
overlap	same server	30/30
overlap	restart per workload	30/30
Total		120/120

Test points were batch 64 / 96 / 128, 10 workloads per point, fixed prompt, output length 1024, compact ragged verify, CUDA Graph enabled, and Radix Cache disabled.

TP consistency audit

verify_token_budget / sum(verify_lens) divergence across TP0-TP7: 0
target CUDA Graph key divergence across TP0-TP7:               0
illegal memory access:                                         0
AcceleratorError / scheduler exception / client truncation:    0
real ProcessGroupNCCL watchdog failure:                         0

This strongly validates the PR's intended effect: broadcasting TP0's final verify budget keeps compact target-verify Graph-tier selection identical across TP ranks and removes the original observed cross-rank mismatch in this B300 TP8 matrix.

Residual issue does not invalidate this PR

A separate compact SPS transition failure remains on machine A:

batch_size_per_rank=80
fraction: 0.2 -> 0.4
M: 160 -> 240
compact target verify

With CUDA Graph enabled, the transition crashes with both overlap enabled and disabled. Before the crash, all TP ranks agree on:

budget=160
sum(verify_lens)=240
target key_size=240
raw_bs=80
slots=128
cross-rank divergence=0

A fresh server running only M=240 completed 1/1. More importantly, an isolated strict eager control using only the profiler's official run_one_round(..., frac=...) flow completed both transition rounds:

frac=0.2, M=160: 4.195949799844881 steps/s, 49 steady steps, match_fraction=1.0
frac=0.4, M=240: 4.242118847887335 steps/s, 47 steady steps, match_fraction=1.0
CUDA / NCCL / scheduler errors: 0

This residual failure is therefore being tracked in #31023 as a CUDA Graph capture/replay or Graph-associated metadata/buffer lifetime issue, not as a failure of this PR's TP budget synchronization.

Cross-machine status

A separate machine B replay used ten independent fresh servers and the same A-machine serve.sh sps-profile + gen_sps.sh chain:

pass:  10
crash: 0
error: 0

Every run completed frac=0.2 and frac=0.4 with final target key_size=240 and zero TP divergence. The residual Graph transition issue is therefore currently nondeterministic and/or environment-sensitive across machines.

Current assessment

PR's cross-TP budget / Graph-tier fix: validated on B300 TP8
original workload matrix:               120/120 pass
strict eager M=160 -> M=240 control:     pass
TP-consistent CUDA Graph transition:     unresolved in #31023
speed cost of one CPU scalar broadcast:  still pending dedicated profiling

An earlier CUDA_LAUNCH_BLOCKING=1 attempt is not counted because the launched service was eager and the SPS CLI rejected it before running the transition workload. A new isolated CUDA-Graph-enabled blocking reproducer is being handled separately.

[Phoenix3334]

Follow-up after the later Bug-2 isolation: the B300/TP8 validation for this PR still stands independently.

I backported the TP0 verify-budget broadcast onto baseline 692c5f7d and ran the full compact-ragged TP8 matrix:

scheduling	lifecycle	result
no-overlap	same server	30/30
no-overlap	restart per workload	30/30
overlap	same server	30/30
overlap	restart per workload	30/30
total		120/120

Across TP0-TP7 in those runs:

verify_token_budget divergence: 0
sum(verify_lens) divergence:     0
target CUDA-Graph key divergence: 0
illegal-memory / scheduler / NCCL-watchdog failures: 0

The residual CUDA-Graph failure that originally remained after this synchronization has since been isolated separately: on B300, #32467 can directly produce/repair the invalid ragged plan (2592/6000 bad unpatched -> 0/30000 patched), and #32467-only full-model startup/replay paths were clean in the tested matrix.

So the current evidence separates the two layers cleanly:

#31195 -> distributed control-plane consistency
#32467 -> GPU plan/index producer correctness

The later Bug-2 findings therefore reinforce rather than weaken the conclusion here: the TP budget/Graph-tier mismatch fixed by this PR was a real and independently validated Bug 1.