[vincentzed]

Motivation

The DeepSeek-V4.1 cookbook launch uses --enable-decoder-swa-bounded-replay. Today the server refuses to start when that flag is combined with a prefill CUDA graph, so V4.1 prefill always runs eagerly.

Eager prefill costs about 85 ms per batch on 4x B300, no matter how short the prompt is. A 590-token prompt pays it, and so does every follow-up turn of a long conversation.

Roadmap #42170 lists "Enable decoder bounded replay" with no owner. The draft #42078 fixes one piece of stale state and says it does not enable the combination. This PR covers that case as well.

Modifications

With decoder bounded replay, the layers after the last KV source layer (the "late layers") only run over the last 128 tokens of each request. The number of those rows depends on how many requests are in the batch, not on how many tokens it has. A prefill CUDA graph has a fixed shape per token count, so the late layers could not be captured in it.

1. The late layers now have their own CUDA graphs, one per number of tail rows (128, 192, 256, 384, ... up to one window per request slot). The prefill graph runs the early layers and then replays the matching late-layer graph from a graph break. A batch is padded to the next size up, which is less than 1.5x its real tail.
2. All late-layer graphs are captured once at startup, during the first prefill graph warmup. No request triggers a capture.
3. The attention metadata for the late layers is captured once per tail size and refreshed in place on every replay.
4. When decoder bounded replay is on, the low-ratio compressor and indexer run eagerly at a graph break instead of inside the graph. The eager indexer handles any context length, so the graph is no longer limited to 16K tokens of context.
5. Capturing the graphs needs much less memory than before:
   • All graphs share one scratch buffer for padded query heads. Each graph used to keep its own (1 GB at 16K tokens).
   • Graph metadata is sized by the batch's own context instead of the model's maximum context, when nothing inside the graph needs the full width.
   • The candidate table from the indexer is dropped as soon as the late layers have used it.
6. Decode no longer sees late-layer metadata left behind by the previous prefill.
7. The checks that reject prompt logprobs and full hidden-state capture moved out of the captured code, because a replay skips that code.
8. Only the breakable backend is accepted. The full and tc_piecewise backends, ROCm, and prefill context parallelism are still rejected at startup.

Every batch now goes through the graph: any number of requests, any size up to the chunk size, any context length. Encoder bounded replay is not part of this PR.

Accuracy Tests

4x B300, deepseek-ai/DeepSeek-V4.1-Flash, TP4 and EP4. Both columns run this branch. "Eager" is what the cookbook launch does today.

# eager
python3 -m sglang.launch_server --trust-remote-code --model-path deepseek-ai/DeepSeek-V4.1-Flash --tp 4 --ep-size 4 --mem-fraction-static 0.8 --cuda-graph-max-bs-decode 64 --max-running-requests 16 --enable-decoder-swa-bounded-replay --disable-prefill-cuda-graph
# graph
python3 -m sglang.launch_server --trust-remote-code --model-path deepseek-ai/DeepSeek-V4.1-Flash --tp 4 --ep-size 4 --mem-fraction-static 0.8 --cuda-graph-max-bs-decode 64 --max-running-requests 16 --enable-decoder-swa-bounded-replay --cuda-graph-backend-prefill breakable
# DSpark rows add: --speculative-algorithm DSPARK --speculative-dspark-block-size 5
python3 -m sglang.test.few_shot_gsm8k --num-questions 1319 --parallel 32

GSM8K, 1319 questions	eager	graph
no speculation	0.906	0.905
DSpark, block size 5	0.904	0.913

Each cell is one run. Repeated runs of one configuration moved between 0.900 and 0.915 while I worked on this, so these are within noise. The DSpark accept length is 3.39 in both columns.

Speed Tests and Profiling

Time to first token, median of 5 requests sent one at a time. Prompts are sent as token ids, so tokenization is not included.

Request, no speculation	eager	graph	speedup
590 new tokens	89.3 ms	30.5 ms	2.9x
2,060 new tokens	88.7 ms	43.4 ms	2.0x
6,318 new tokens	102.1 ms	75.8 ms	1.35x
14,795 new tokens	160.6 ms	141.3 ms	1.14x
417 new tokens after 13.3K cached tokens	93.4 ms	29.8 ms	3.1x
554 new tokens after 59.9K cached tokens	104.1 ms	41.3 ms	2.5x

Request, DSpark	eager	graph	speedup
590 new tokens	89.1 ms	28.5 ms	3.1x
2,060 new tokens	89.2 ms	41.0 ms	2.2x
6,318 new tokens	100.0 ms	72.4 ms	1.38x
14,795 new tokens	150.6 ms	141.7 ms	1.06x
417 new tokens after 13.3K cached tokens	88.1 ms	31.7 ms	2.8x
554 new tokens after 59.9K cached tokens	98.2 ms	41.4 ms	2.4x

GSM8K wall time, 32 parallel	eager	graph	speedup
no speculation	143.1 s	73.7 s	1.94x
DSpark, block size 5	65.8 s	41.8 s	1.57x

• The one case that only ties is the largest one: 48 requests of about 7.4K tokens at 8 concurrent run at 105.2K input tokens/s eager and 105.6K with the graph (107.9K and 106.9K with DSpark). Those batches fill the 16K chunk and the GPU is the limit there.
• Capturing all graphs takes about 17 GB per GPU. Before the memory changes in this PR, capturing the same set ran out of memory.
• In the GSM8K runs, all but 1 of about 1,480 prefill batches used the graph.

Rerun after merging main at 28621739fc

• Main was merged twice on 2026-10-12. The conflicts came from three PRs:
  • [Dsv4.1] Mega-mhc integration #42808 added a mega_mhc_seam_open argument to the per-layer call. This PR had moved that call into a helper, so the helper now passes the argument too.
  • [AMD][V4.1][*/N] Fix DP attention + TP MoE decode NaN and support bounded replay #43077 added a ROCm-only path for DP attention in the same layer loop. It is kept as it is on main, in front of the helper call.
  • [BCG][DSV4-Flash][CP] Support DeepSeek V4 prefill with interleave CP and breakable CUDA graphs #38603 added the prefill graph under context parallelism. Its condition and this PR's condition are both kept in validate_deepseek_v41_features and in low_ratio_prefill_graph.
• Same launch commands and probe as above, one run per cell.

GSM8K, 1319 questions	eager	graph
no speculation, accuracy	0.908	0.905
no speculation, wall time	135.5 s	72.9 s
DSpark, accuracy	not rerun	0.910
DSpark, wall time	not rerun	41.9 s

TTFT, no speculation	eager	graph
590 new tokens	77.5 ms	27.0 ms
2,060 new tokens	80.7 ms	40.7 ms
6,318 new tokens	94.6 ms	72.5 ms
14,795 new tokens	132.8 ms	138.4 ms
417 new tokens after 13.3K cached tokens	81.8 ms	32.2 ms
554 new tokens after 59.9K cached tokens	92.7 ms	43.5 ms

• Eager prefill on main is faster than it was when the tables above were measured. At 14,795 new tokens the two now tie, and the graph is 4% slower in this run. The graph numbers are about the same as before.
• No scheduler exceptions. In the GSM8K runs all but 1 prefill batch used the graph. The 26 unit tests in the three touched test files pass.

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.

---

CI States

Latest PR Test (Base): ❌ Run #38269076379
Latest PR Test (Extra): ❌ Run #38269076201
Latest PR Test (AMD ROCm 10): ❌ Run #38269076416

[vincentzed]

@Oasis-Git this is for the "Enable decoder bounded replay" item on #42170. It lets the cookbook V4.1 launch keep --enable-decoder-swa-bounded-replay and still use the breakable prefill CUDA graph. Could you take a look, or point me to the right reviewer?

Two related things:

• [Fix] Stop DeepSeek-V4.1 encoder bounded replay from crashing on the first prefill without speculation #42902 is a small separate fix. On main, --enable-encoder-swa-bounded-replay without speculation crashes on the first prefill.
• I also have encoder bounded replay working under the graph on an earlier version of this branch. I will port it to this design and open it as a follow-up once this one has had a first review.

@harmya this also covers the stale tail state from your draft #42078. The replay path now sets or clears the tail for every batch.

[vincentzed]

@Oasis-Git a gentle reminder. This still needs a first review and the run-ci label, and it has no conflicts with main. #42902 is the small crash fix that goes with it, and it is also waiting for a review.

[vincentzed]

I merged main into this branch (now at 28621739fc). The conflicts came from #42808, #43077 and #38603, which all changed the same layer loop and the same two conditions. The PR description lists what was kept from each, under "Rerun after merging main".

After the merge, GSM8K is 0.905 with the graph and 0.908 eager, and the wall time is 72.9 s against 135.5 s. Time to first token for 590 new tokens is 27.0 ms against 77.5 ms. At 14,795 new tokens the two now tie, because eager prefill on main got faster.

@Oasis-Git a gentle reminder: this still needs a first review and the run-ci label. The files it touches change often on main, so an early look would help keep it mergeable.

[vincentzed]

Closing this. #43736 has merged and covers it: decoder bounded replay now runs under the breakable prefill CUDA graph, with the tail rows in the same graph.