[YAMY1234]

Motivation

python3 -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V4.1-Flash --trust-remote-code \
  --tp 4 --dp 4 --enable-dp-attention --ep-size 4 \
  --cuda-graph-backend-prefill disabled \
  --max-running-requests 64 --chunked-prefill-size 32768 \
  --enable-encoder-swa-bounded-replay --enable-decoder-swa-bounded-replay

Startup rejects both replay flags under DP attention, including when either is used alone. Removing the checks is insufficient: encoder replay adds up to 128 rows per prefix hit after the scheduler has gathered the token counts, and decoder replay removes earlier rows at the late-layer boundary. The MoE gather must use the correct count from every rank at both boundaries, including ranks that decode or are idle.

This PR depends on #43166, which folds encoder replay into the extend forward; until it merges, please review only the last two commits. It adds CUDA DP-attention support for both replay flags. #43077 (merged) added the ROCm decoder path; this PR reuses its late-layer enter/exit helpers on CUDA and leaves its ROCm behaviour unchanged.

Modifications

• Include each rank's folded encoder replay rows in the existing scheduler token-count gather. The folded batch validates its own gathered slot instead of replacing counts after the gather.
• Carry each rank's decoder trim count as one extra column of that same gather. At the late-layer boundary, every rank uses full rows − trimmed rows for the MoE gather through [AMD][V4.1][*/N] Fix DP attention + TP MoE decode NaN and support bounded replay #43077's _enter_late_layer_dp / _exit_late_layer_dp, then restores its original counts after the last layer. The helper's count copy is made non-blocking, so no additional eager-forward collective, device-to-host count read or stream sync is needed.
• Late layers do not use the full batch's gathered token ids; check at startup that they begin after all hash-routed layers.
• Support CUDA DP attention with attention TP 1 and the default MoE backend (none), next to [AMD][V4.1][*/N] Fix DP attention + TP MoE decode NaN and support bounded replay #43077's ROCm acceptance; retain the checks for attention TP > 1, MoE all-to-all backends and prefill graphs.

Accuracy Tests

The table below was collected on the earlier prototype; the rebased head was re-gated afterwards (last paragraph). DeepSeek-V4.1-Flash, 4×GB300, TP4, DP attention 4; cold prompts followed by full prefix hits, all 200 hit prompts cached:

configuration	flags off, GSM8K-200 cold / hit	flags on, cold / hit
encoder, no speculation	197 / 194	196 / 197
encoder, DSpark block 5	196 / 195	196 / 195
encoder + decoder, no speculation, round 1	195 / 196	197 / 197
encoder + decoder, no speculation, round 2	196 / 194	197 / 196

Each arm finds all 30 needles at 32K and 64K on cold, first-hit and second-hit passes. First tokens agree on all 90 comparisons. Encoder-only first-token KL max is 6.5e-5 without speculation; the two combined rounds have maxima 6.0e-5 and 1.0e-4.

Rebased head (current main + #43166), same setup: both flags GSM8K-200 198 / 197, needles 30/30 on all three passes, first tokens unchanged; encoder + DSpark block 5 197 / 197. Unit tests: 67 passed, 146 subtests (this PR's files plus #43077's). They cover folded rows in the real scheduler gather, the backend's decoder-tail layout with and without encoder replay, zero trim for decode/idle/verify, a trimming peer with non-trimming ranks, restoring the full DP counts, and launch validation for the flags together or separately. Unsupported platform/layout combinations remain rejected.

Speed Tests and Profiling

DeepSeek-V4.1-Flash on 4×GB300, both replay flags enabled, speculative decoding disabled. Two same-node rounds in reversed order, on 24c93754fa69 plus #42715 (81a98aa493a, a3db1d2df79) and #42921 (95efd58ec3c); the optimizations are not included in this PR. DP uses TP4 / DP attention 4 / EP4, chunk 32768 (8192/rank); TP4 uses EP4, chunk 8192. Both use mem fraction 0.8 and max-running 256 (DP 64/rank), above every measured concurrency.

Values average the two rounds. Cold: K input tok/s / TTFT p50 ms; hit: TTFT p50 ms; decode: K output tok/s / per-user tok/s / step p50 ms.

static workload	DP attention 4	TP4
cold 8K, c8	119.1 / 584	43.8 / 1493
cold 8K, c32	163.5 / 1557	44.0 / 5935
cold 32K, c8	158.7 / 1607	44.0 / 5938
cold 32K, c32	166.2 / 6244	44.0 / 23795
64K prefix + 1K new, c1	209	200
64K prefix + 1K new, c16	454	788
decode 8K input, c32	3.191 / 99.7 / 10.02	3.999 / 125.0 / 7.99
decode 8K input, c64	5.641 / 88.1 / 11.33	5.988 / 93.6 / 10.68
decode 8K input, c128	9.330 / 72.9 / 13.71	10.586 / 82.7 / 12.01
KV pool total tokens per rank	61,464,576	64,036,352

Cold prompts are flushed; every timed prefix hit caches all 65,536 prefix tokens. Decode rates use the first timed wave's common token 129–896 steady interval, excluding prefill and ramp/drain. c32/c64 traces generate 4096 tokens, c128 1024; every streamed token is counted. Prefill graphs are disabled; c128 captures its active decode shape (DP 32/rank, TP 128).

DP gains 2.7–3.8× cold prefill throughput and ~42% lower c16 hit TTFT with both flags. TP4 has ~5% lower c1 hit TTFT and higher decode throughput at all three tested concurrencies, including ~13% more at c128. Hit TTFT crosses between c1 and c16; no decode crossover is observed through c128 at 8K input.

The flags' own cost under DP attention (flags on vs off, earlier prototype of this change): cold 8K c8 input +12.8% / +8.9%; hit TTFT c1 +10 / +19 ms, c16 +27 / +23 ms; decode c32 step +0.31 / +0.35 ms.

Checklist

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

---

CI States

Latest PR Test (Base): ❌ Run #40266004374
Latest PR Test (Extra): ❌ Run #40266004379
Latest PR Test (AMD ROCm 10): ❌ Run #40266004458