[kpham-sgl]

Changes

Follow-up to #43812 and #43736. Together they admit decoder SWA bounded replay with DP attention under the breakable prefill CUDA graph, but the DP-geometry swap from #43812 skipped itself while the stream was capturing and graph replays never re-run the late-layer loop, so inside the graph the late-layer MoE gather kept the full-shape counts while the hidden states were the tail block (only consistent when the shape fits within max_running_requests x 128).

• late_layer_dp_counts takes in_graph: prefill graphs pad every DP rank to one shape, so the tail block is rank-uniform and the counts come from the captured shape with no collective. The capture bakes the geometry in; replays inherit it.
• enter_late_layer_dp uses device ops only (full_like for the count tensor, clamp on the runner's static non-padded scalar) so it is capture-safe and rebounds per replay.
• The model loop decides tail_in_graph before computing the counts.

Accuracy

DeepSeek-V4.1-Flash on 2xGB300, TP2/DP2/EP2, MegaMoE, DSpark, breakable prefill graph, bounded replay on, evaluated with the GB300 CI test's evaluators (max_running_requests 32 per rank, so the in-graph tail block of 4096 rows is smaller than the captured 4096-8192 shapes).

Eval	Score	CI threshold
GSM8K (zero-shot chat, 200)	0.980	0.95
MMLU (500)	0.926	0.90

DSpark accept length 3.4-4.3. A 10-shot completion GSM8K smoke check on the same box gave 0.925 with the graph vs 0.910 eager, both with replay and DP2.

---

CI States

Latest PR Test (Base): ❌ Run #38500425131
Latest PR Test (Extra): ❌ Run #38500425143
Latest PR Test (AMD ROCm 10): ❌ Run #38500425067

[kpham-sgl]

/rerun-test test/registered/e2e/models/test_deepseek_v41_flash_gb300.py

[github-actions]

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

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300.py

[kpham-sgl]

Validation, done on 2xGB300 (khoa-pham-gb300-2gpu) by running the GB300 CI test's mixin locally at TP2/DP2/EP2 with MegaMoE, DSpark, --cuda-graph-backend-prefill breakable, --enable-decoder-swa-bounded-replay.

• Prefill graph captured shapes up to 8192 under a shared MAX_LEN bucket; with 32 running requests per rank the in-graph tail block is 4096 rows, so the 4096-8192 shapes ran the late layers on a block smaller than the shape, which is the regime this PR fixes.
• GSM8K (zero-shot chat, 200): 0.98 (threshold 0.95). MMLU (500): 0.926 (threshold 0.90). DSpark accept length 3.4-4.3.
• A 10-shot completion GSM8K smoke check on the same box: 0.925 with the graph vs 0.910 eager, both with replay and DP2.

At TP2 the model needs the Engram host table and --mem-fraction-static 0.86.

[kpham-sgl]

/rerun-test test/registered/e2e/models/test_deepseek_v41_flash_gb300.py

[github-actions]

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

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300.py