[Bug] Fix DP-attention state-capturer crash on cuda_graph_batch=None (#30712)

[aryanyadav0402]

Motivation

Fixes #30712.

On the prefill / piecewise CUDA-graph path with DP attention and routed-expert return enabled, the DP-attention state capturers crash with:

TypeError: unsupported operand type(s) for *: 'int' and 'NoneType'

Root cause. The routed-expert and indexer state capturers in ModelRunner.forward() were passed:

cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None)

regardless of which runner actually executed the forward pass. On the prefill/piecewise CUDA-graph path (or whenever the decode runner was not the one that ran), that resolves to None, so get_dp_local_slice_cpu computes dp_rank * None and raises for dp_rank > 0. get_dp_local_slice_cpu itself was already correct — the defect was the caller feeding it the wrong runner's extent.

Modifications

• Add ModelRunnerOutput.cuda_graph_padded_extent: Optional[int], set at each _forward_raw exit to the padded per-rank extent of the runner that actually executed:
  • decode graph → decode_cuda_graph_runner.bs
  • prefill graph → prefill_cuda_graph_runner.padded_num_tokens (new property exposing the padded static token count of the last replay)
  • eager / split-prefill → None
• Both capturer call sites now read output.cuda_graph_padded_extent instead of the hard-coded getattr(self.decode_cuda_graph_runner, "bs", None).

This makes the invariant can_run_graph == True ⇒ extent is not None hold by construction. I deliberately did not add a None-guard that falls back to the eager prefix-sum layout: on the graph path that would silently read the wrong slice for dp_rank > 0 (a correctness bug) instead of the loud TypeError. The change also removes an over-defensive getattr, per the no-getattr-defensive convention.

The decode path is byte-identical (it already passed .bs; it now passes the same value via the new field).

Accuracy Tests

New CPU regression test test/registered/unit/layers/test_dp_attention_local_slice.py:

• TestModelRunnerOutputCudaGraphExtent — exercises the fix wiring: ModelRunnerOutput exposes cuda_graph_padded_extent, it defaults to None on the eager path, and the graph-path extent feeds a valid rank-padded slice (not dp_rank * None).
• TestGetDpLocalSliceCpu — pins the underlying slicing invariant (graph → dp_rank * padded_extent; eager → prefix sum; the two must diverge for dp_rank > 0; None on the graph path raises).

Verified in the lmsysorg/sglang:latest image:

• Stock tree (fix reverted, test present): 3 failed, 4 passed — the three wiring tests fail because the field does not exist.
• Fixed tree: 7 passed.

On reproducing the live server crash — full transparency: I attempted an end-to-end server repro on 4× H100 NVL (driver 595 / CUDA 13.2), but could not trigger it on this build for a reason unrelated to this bug: prefill CUDA-graph capture fails during warmup for every MoE/DP model I tried (DeepSeek-V2-Lite: CUBLAS_STATUS_EXECUTION_FAILED / MLA shape errors; Qwen1.5-MoE and OLMoE fail elsewhere in capture), on both stock and patched trees, before any request. With prefill-graph capture unavailable, the runtime dp_rank * None path is unreachable here. The evidence for this PR is therefore the root-cause analysis plus the deterministic regression test above, which fails on the unpatched tree and passes with the fix. Happy to add a live before/after trace if a maintainer can point me to a build/config where prefill-graph capture with DP attention succeeds.

(Separately, the Qwen capture failure surfaced inside the routed-experts capturer itself — an off-by-dp_size shape mismatch — which looks like a distinct defect; I'll file it as its own issue.)

Speed Tests and Profiling

No performance impact — the change only threads an already-computed integer through ModelRunnerOutput; no new work on any hot path.

Checklist

• Format code with pre-commit (ruff/black/isort clean).
• Add unit tests (CPU regression test above).
• Update documentation — n/a (bug fix, no API/user-facing surface change).
• Accuracy/speed benchmarks — n/a (no model-output or perf change; see above).
• Follow the SGLang code style guidance.

---

Disclosure: developed with AI assistance (Claude). The change is authored and owned by me, and I can defend it in review.

---

CI States

Latest PR Test (Base): ❌ Run #29306267159
Latest PR Test (Extra): ❌ Run #29306267055

[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!

[aryanyadav0402]

@misaAle thanks again for pointing me at this one.

The fix records the padded per-rank extent of the runner that actually executed on ModelRunnerOutput.cuda_graph_padded_extent (decode → bs, prefill graph → padded_num_tokens, eager → None) and has the capturers read it, so can_run_graph=True ⇒ extent is not None holds by construction rather than via a None-guard. There's a CPU regression test that fails on the unpatched tree and passes with the fix.

One thing I want to be transparent about: I couldn't reproduce the original runtime crash end-to-end on my hardware. On this build prefill CUDA-graph capture fails during warmup for every MoE/DP model I tried (unrelated to this bug — it happens on stock too, before any request), so the dp_rank * None path is unreachable here. Since you have the working repro, would you be able to confirm the fix resolves the crash on your setup? Happy to iterate if anything looks off.

Separately, while testing I hit what looks like a distinct defect in the same subsystem: during prefill-graph capture with DP attention, state_capturer/base.py writes the DP-gathered topk_indices into a per-rank-sized buffer and fails with a shape mismatch off by exactly the dp_size factor (16384 = dp_size × 8192). That's not fixed by this PR; I'll write it up as its own issue.

[misaAle]

@aryanyadav0402 yeah actually I'm not able to reproduce the issue on the standard sglang docker image. I was using a SlimeRL image: slimerl/slime:nightly-dev-20260618a, and it's currently the only way I can reproduce the issue. I've tried reproducing the issue using the standard docker images, but I would just run into the same issues you have:

RuntimeError: The expanded size of the tensor (8192) must match the existing size (16384)
  ... state_capturer/base.py:41 -> self.buffer[:batch, layer_id, :] = topk_indices

Regardless, seems like this combination has just not been tested and the prefill path is disabled by intent. Would recommend following up on the issue you filed. Thank you for investigating!

[aryanyadav0402]

Thanks @misaAle — that's really helpful, and it lines things up nicely.

The expanded size (8192) must match (16384) you hit on the standard docker image is exactly the separate sizing bug I filed as #31116 — the routed-experts device buffer omits the prefill graph's max capture size, so the DP-concatenated topk overflows it. That's an independent repro of #31116 on stock docker, thanks for confirming it. With #31116's sizing fix applied, capture completes and you then land on #30712 at request time — I've validated both end-to-end (before/after logs), and together they give a fully clean run on the forced-prefill path.

Totally agree the combo is gated by intent today. Since the driver here is the SlimeRL routed-experts-capture path, would you be open to treating these two as the enablement work for DP-attn × prefill-graph × routed-experts? Concretely: this PR (#31100) is a safe, backward-compatible fix — it only records the actually-executed runner's padded extent and removes a banned defensive getattr, with no change to the default (disabled) path — so it's low-risk to merge even while the path stays gated. And I'm happy to open the #31116 sizing PR to finish it. If the path isn't something you want enabled yet, that's fine too — just let me know and I'll keep them scoped as latent-bug hardening.