[HZY-Wade]

Motivation

A PD-disaggregated nvidia/GLM-5.2-NVFP4 deployment can hang or crash after decode requests when DP attention and EAGLE MTP are enabled together. The equivalent non-disaggregated deployment completes the same requests.

This is a follow-up to #30839. A minimal reproduction uses the following relevant arguments. Mooncake and RDMA arguments may need adjustment for the local network.

Prefill:

python -m sglang.launch_server \
  --model-path nvidia/GLM-5.2-NVFP4 \
  --trust-remote-code \
  --quantization modelopt_fp4 \
  --moe-runner-backend flashinfer_trtllm \
  --disaggregation-mode prefill \
  --disaggregation-bootstrap-port 8998 \
  --disaggregation-transfer-backend mooncake \
  --tp 8 \
  --dsa-prefill-backend trtllm \
  --host 0.0.0.0 \
  --port 30000

Decode:

python -m sglang.launch_server \
  --model-path nvidia/GLM-5.2-NVFP4 \
  --trust-remote-code \
  --quantization modelopt_fp4 \
  --moe-runner-backend flashinfer_trtllm \
  --disaggregation-mode decode \
  --disaggregation-bootstrap-port 8998 \
  --disaggregation-transfer-backend mooncake \
  --tp 8 \
  --dp-size 8 \
  --enable-dp-attention \
  --dsa-decode-backend trtllm \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --host 0.0.0.0 \
  --port 30001

Five related failure modes or follow-up issues were observed.

First, the rank receiving the real request can enter the eager MTP draft path while idle DP ranks replay the target verification CUDA graph. The ranks then reach different collective phases and the scheduler eventually triggers the watchdog.

For GLM-5.2 IndexShare, a PD decode round can arrive without dsa_topk_indices. The old fallback disabled graph replay only on the active rank. Idle ranks remained graph-eligible even though all ranks participate in the same DP and EP forward.

Second, overlap scheduling can make the graph vote observe stale DSA seed state. When a seeded running batch merges a seedless prebuilt batch, future_dsa_topk_indices_available becomes false while the materialized dsa_topk_indices tensor can remain non-empty. The future flag must be authoritative whenever future inputs exist, otherwise the active rank falls back only after the DP-wide vote.

Third, after synchronizing eager fallback, DP padding exposed an eager draft postprocessing mismatch:

draft_topk1_postprocess
  assert positions.shape[0] == next_token_logits.shape[0]

EAGLE postprocessing must consume only rank-local token rows rather than padded logits.

Fourth, idle eager ranks and active padded ranks exposed two DSA shape mismatches. Idle ranks do not have real page-table rows and must skip sparse index selection. Active ranks must trim padded query and top-k rows to the real batch described by pre-planned metadata, run attention on that prefix, then restore the physical DP shape before later MLP and EP collectives.

Fifth, the correctness-first fix originally folded seedless fallback into the generic DP graph gate. This kept ranks consistent but also forced target verification and draft-extend eager whenever only the EAGLE draft phase lacked a DSA seed. Under the reproduced workload, about 9.7% of reported decode batches lost generic CUDA graph replay, creating an avoidable performance regression.

Modifications

• Add a speculative-worker hook that reports when a batch requires rank-consistent eager draft execution.
• Detect the GLM-5.2 seedless IndexShare condition before DP MLP synchronization in disaggregated decode.
• Treat future_dsa_topk_indices_available as authoritative under overlap scheduling.
• Append a dedicated draft-graph eligibility vote to the existing DP MLP-sync all-gather. No additional collective is introduced.
• Make EAGLEDraftCudaGraphRunner consume both the generic graph vote and the draft-specific vote.
• Keep target verification and draft-extend on their normal CUDA graph path when only draft seed state is missing.
• Slice eager draft logits, hidden states, and positions to the local draft-token count.
• Short-circuit eager DSA index selection on idle DP ranks with all-invalid top-k rows.
• Trim TRT-LLM DSA attention inputs to the real metadata batch and zero-pad the output back to the physical DP shape.
• Add regression coverage for current and future seed states, overlap batch merging, independent generic and draft graph votes, active and idle DP padding, DSA idle indexing, and TRT-LLM page-table alignment.

Accuracy and Stability Tests

The changes synchronize control flow and isolate DP padding from metadata that describes only real request rows. Model math for real tokens is unchanged.

• Repository pre-commit hooks pass for all modified files.
• Focused local suite: 14 tests passed, plus 4 subtests.
• The patch was hot-applied to the reproduced PD decode deployment with TP8, DP8, DP attention, TRT-LLM DSA, NVFP4, and EAGLE MTP.
• Under mirrored and validation traffic, all 8 scheduler ranks remained alive.
• 2,019 reported decode batches completed with generic CUDA graph enabled: 2,019 true and 0 false.
• Per-rank running batch size reached 7 during observation.
• No scheduler exception, assertion, watchdog, or CUDA graph collective hang occurred.

Speed Tests and Profiling

The draft-specific gate preserves the required rank-consistent eager fallback but limits it to the EAGLE draft phase.

Before this follow-up, approximately 9.7% of reported decode batches disabled the generic graph path because a seedless draft also disabled target verification and draft-extend graphs. After this follow-up, the observed deployment reported 0 generic graph fallbacks across 2,019 decode batch samples.

This is an operational validation rather than a controlled end-to-end throughput benchmark. The expected gain is workload-dependent and proportional to the frequency of seedless draft rounds; no extra DP collective is added.

Checklist

• Format the code with repository pre-commit hooks.
• Add unit regression tests.
• Follow the SGLang code style guidance.
• Validate on the reproduced PD + DP attention + EAGLE deployment.
• Documentation update (not applicable; no user-facing API change).
• Controlled accuracy and speed benchmark tables.

---

CI States

Latest PR Test (Base): ❌ Run #30020688468
Latest PR Test (Extra): ❌ Run #30020688547

[gemini-code-assist]

Caution

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

[gemini-code-assist]

Caution

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

[HZY-Wade]

@ShangmingCai cc

[HZY-Wade]

Hi @IvanShan177 , just checking whether you are still available to review this PR when you have time. AMD-AGI/Infera#58 has now independently reproduced the same issue on ROCm and adopted the core DP draft-graph vote from this PR. Please let me know if I should ask another reviewer instead. Thanks!