[Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs by Juntian777 · Pull Request #59532 · vllm-project/vllm

Juntian777 · 2026-10-01T03:09:16Z

Purpose

Run the decoder-side SWA bounded replay layers (#58132) in CUDA graphs of their own, and trim them in PIECEWISE graph steps too. #58132 trims layers 21-39 to each prefill's last 128 tokens only on eager steps (above the model graph's capture range) and runs them eagerly there; those steps are often launch-bound (#58132 measured ~175 ms of kernel launches with the GPU waiting on a single 4K-8K prefill), and with a large max_cudagraph_capture_size nothing trims at all.

DecoderReplayCudaGraphManager (nvidia/decoder_replay_cudagraph.py), a CudaGraphManager owned by DeepseekV41ModelState, captures breakable PIECEWISE graphs of the replay layers, keyed by the replay batch's padded token count and reading the layer inputs from static buffers. Attention runs in their eager breaks: the replay batch mixes prefill tails and decodes, and the SWA and sparse MLA builders support uniform batches only.
PIECEWISE model graphs of at least compilation_config.decoder_replay_trim_threshold padded tokens (768 by default) break out to the replay layers, which run on a replay batch (trimmed, or whole when nothing trims) and write their outputs to fixed buffers the graph's next segment reads. FULL graphs and smaller PIECEWISE ones keep the layers on the whole batch, as before.
_prepare_replay_batch pads the replay batch to the smallest graph that fits the largest DP rank's replay batch and sets it on DecoderReplayLayers as one ReplayBatch (rows, forward context, graph runner). A replay batch above max_cudagraph_capture_size, possible only in eager steps, stays eager.
Captures run the runtime path on a dummy batch, so the graphs read the buffers a real step's replay batch fills; on MRV2 the SWA KV insert stays inside the graph and reads the slot mapping by address.
ModelState.capture_inner_cudagraphs is a new hook the runner calls after the model's capture; it is a no-op by default.
On by default with breakable piecewise graphs (the default for DeepSeek-V4.1); compilation_config.cudagraph_decoder_replay=false turns it off. Also off under LoRA, prompt embeddings or on a non-first PP rank. Off keeps [Model] Decoder-side SWA bounded replay for DeepSeek-V4.1 #58132's behavior.

No open PR runs the replay layers in graphs. #56752 (closed) captured them nested inside the model graph's capture; this captures them on their own and has the model graph break out to them.

Benchmark

Prefill-only TTFT of DeepSeek-V4.1-Flash on GB200 at f21ea3c947 (the last commit only adds the switch and renames), default CUDA graph settings, max_num_batched_tokens=16384, prefix caching off, max_tokens=1. With the default capture range every 1K-128K prefill step is eager, so this measures the eager path, not PIECEWISE trimming. Median of 3 runs in ms; c1 is one request at a time, c8 is 8 concurrent. Off is --no-swa-bounded-replay; parent is df8fd42 (main with #58132) with bounded replay on. Each topology ran all three on one node.

TP4 (repeats within ±2%; this PR's 1K c1 runs were 102/103/131 ms):

ISL	conc	Off	Parent	This PR	vs off	vs parent
1K	c1	160.9	165.9	103.3	1.56x	1.61x
1K	c8	324.3	333.7	269.1	1.21x	1.24x
4K	c1	165.4	172.5	107.0	1.55x	1.61x
4K	c8	505.3	505.2	442.7	1.14x	1.14x
8K	c1	170.0	177.8	112.5	1.51x	1.58x
8K	c8	963.5	683.9	562.9	1.71x	1.21x
32K	c1	540.1	376.0	327.8	1.65x	1.15x
32K	c8	3984.8	2714.7	2292.4	1.74x	1.18x
128K	c1	2267.1	1522.3	1392.1	1.63x	1.09x
128K	c8	9847.5	6553.1	5933.2	1.66x	1.10x

DEP2 (one fixed prompt per length; each ratio compares runs on the same GPU pair, since the two pairs of a node differ by 3-10%; c1 repeats within ±3%; c8 below 128K is omitted, it varies ±10-30% with how the two API servers split the 8 requests across DP ranks):

ISL	conc	Off	This PR	vs off	Parent	This PR	vs parent
1K	c1	103.8	72.5	1.43x	116.2	79.9	1.45x
4K	c1	135.3	92.3	1.47x	117.4	94.7	1.24x
8K	c1	246.9	159.4	1.55x	164.1	163.9	1.00x
32K	c1	983.1	609.8	1.61x	617.8	627.8	0.98x
128K	c1	3950.5	2460.3	1.61x	2518.0	2567.8	0.98x
128K	c8	9114.6	5716.4	1.59x	5467.9	5508.2	0.99x

Why TP4 gains at long prompts and DEP2 does not

A profile of one 128K prefill shows the trimmed layers need only a few ms of GPU per chunk, but on TP4 their eager launches take several times longer and layers 0-20 do not hide them, so the GPU sits idle through most of layers 21-39; in graphs that idle time is gone. A DEP2 rank runs layers 0-20 on the whole chunk without TP splitting, so the CPU is far ahead of the GPU and the eager launches were already hidden; this PR only helps DEP2 where the step is short (1K-4K).

🤖 Generated with Claude Code

…IECEWISE graphs Run the decoder-side SWA bounded replay layers on the replay batch in breakable PIECEWISE graphs of their own, captured by a CudaGraphManager owned by DeepseekV41ModelState. PIECEWISE model graphs above VLLM_DSV41_DECODER_REPLAY_TRIM_THRESHOLD tokens (1024 by default) break out to those layers, so they trim too instead of only eager steps. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

…tokens The replay batch builds attention metadata only for the groups the replay layers read, skipping the kv sources' compressor and their other compressed-KV and indexer groups (about 0.55 ms of CPU per trimming step). PIECEWISE graphs of at least 768 padded tokens now trim: the 384 and 576 graphs lost ~3 ms per mixed decode step to the replay overhead, while 768 and 1020 gained 18-25%. Trim the tests to the regression checks. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

… bound; clarify replay batch naming Add compilation_config.cudagraph_decoder_replay (default on) to turn the replay graphs and PIECEWISE trimming off, and reject a decoder_replay_trim_threshold at or below the replay window, where no graph has anything to trim. Keep the threshold in the model state only: the replay layers break out of a PIECEWISE graph whenever a replay batch is set, which the model state does at the threshold or above. Build the replay batch's DP metadata once, after its padding is known. Rename replays/trims to force_replay_batch/needs_replay_batch, replay_graphs to replay_cudagraphs, and tighten the comments. Trim the tests to the regression checks. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

A replay batch that dispatches a replay graph holds at least `window` rows: eager steps replay only when a prefill trims to its last `window` rows, and PIECEWISE steps at the trim threshold hold more than `window` tokens. Sizes below the window were never replayed, so start there and step by the window, widening to 1/8 of the power-of-two interval past 16 * window. With the GB200 default bound of 1024 this captures 8 replay graphs instead of 64. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

Juntian777 · 2026-10-06T00:22:38Z

/ci run

github-actions · 2026-10-06T00:24:02Z

✅ Triggered Buildkite CI #92990 for commit 0a9ae8147a5b.

Juntian777 · 2026-10-06T01:56:18Z

/ci retry

github-actions · 2026-10-06T01:56:36Z

✅ Queued 1 failed job(s) for retry in Buildkite CI #92990.

prepare_inputs_to_capture handed CUDAGraphMode.NONE to ModelState.prepare_attn for every capture, while PIECEWISE steps at runtime hand PIECEWISE. A model state that prepares what the graph runs by its mode then captured something else: DeepSeek-V4.1 builds the decoder replay batch for PIECEWISE graphs at the trim threshold or above (vllm-project#59532), so its graphs were captured without the replay break and ran the replay layers on the whole batch, though every step still prepared a replay batch. PIECEWISE captures now pass PIECEWISE. The other model states branch on FULL only, so they build the same metadata as before. DSV4.1-Flash DEP4 on B200, CUDA graphs up to 8192 tokens: 8k-token prefills 142.7k -> 236.9k total tok/s, 32k in / 256 out 772 -> 1,281 output tok/s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>

…IECEWISE graphs (vllm-project#59532) Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> (cherry picked from commit eb2901b)

…IECEWISE graphs (vllm-project#59532) Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

上游来源：vllm-project/vllm#59532 (eb2901b735)，适配点： - DecoderReplayLayers 换 ReplayBatch 语义，新增 graph break 与 trim_threshold； - 新增 decoder_replay_cudagraph.py（按 replay batch padded size 的可断图）； - CompilationConfig 增加 cudagraph_decoder_replay / decoder_replay_cudagraph_capture_sizes / decoder_replay_trim_threshold； - 补 forward_context.in_piecewise_cudagraph（本仓缺失）； - ModelState 新增 capture_inner_cudagraphs 钩子，model_runner 在 capture 后调用； - model_state 按上游整文件合入，保留本仓 InputBatch 无 fast_prefill、以及 model_config.enable_return_routed_experts 的适配； - 上游两个测试文件同步更新（12 passed）。 Co-authored-by: DSH Agent <dsh-agent@local> Signed-off-by: DSH Agent <dsh-agent@local> Signed-off-by: reputationly <197039020@qq.com>

上游来源：vllm-project/vllm#59532 (eb2901b735) 的配套测试改动。 Co-authored-by: DSH Agent <dsh-agent@local> Signed-off-by: DSH Agent <dsh-agent@local> Signed-off-by: reputationly <197039020@qq.com>

mergify Bot added deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models nvidia mrv2 Model Runner V2 specific labels Oct 1, 2026

github-project-automation Bot added this to NVIDIA Oct 1, 2026

Juntian777 changed the title ~~[Perf][DSv4.1] Run decoder replay layers in CUDA graphs on eager steps~~ [Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs Oct 1, 2026

Juntian777 force-pushed the feat/dsv41-replay-cudagraph branch from 1fa8c15 to c834870 Compare October 1, 2026 04:19

add decoder cudagraph configs

df6c92b

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

mergify Bot added the torch.compile label Oct 1, 2026

github-project-automation Bot added this to torch.compile integration Oct 1, 2026

github-project-automation Bot moved this to To triage in torch.compile integration Oct 1, 2026

Juntian777 force-pushed the feat/dsv41-replay-cudagraph branch from 556b658 to c7220f0 Compare October 1, 2026 07:57

ZJY0516 added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 1, 2026

[Perf][DSv4.1] Default replay graph sizes to 8 per power-of-two interval

f21ea3c

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com>

Juntian777 marked this pull request as ready for review October 2, 2026 22:33

Juntian777 requested review from AndreasKaratzas, DarkLight1337, WoosukKwon, mgoin, njhill, robertgshaw2-redhat, yewentao256, youkaichao and ywang96 as code owners October 2, 2026 22:33

Juntian777 and others added 2 commits October 5, 2026 23:50

Merge branch 'main' into feat/dsv41-replay-cudagraph

0a9ae81

ivanium merged commit eb2901b into vllm-project:main Oct 6, 2026
195 checks passed

simondanielsson mentioned this pull request Oct 6, 2026

[ROCm][Perf][DSv4.1] Enable decoder SWA bounded replay #59039

Merged

ivanium mentioned this pull request Oct 7, 2026

[Blog] DeepSeek-V4.1-Flash optimizations vllm-project/vllm-project.github.io#354

Merged

stefankoncarevic mentioned this pull request Oct 7, 2026

[CI][ROCm] Re-enable the DSv4.1 decoder replay graph test on ROCm #60392

Merged

6 tasks

cocobronie2333-dot mentioned this pull request Oct 20, 2026

[Feature][Model] Support DeepSeek V4.1 decoder replay PIECEWISE ACL graphs vllm-project/vllm-ascend#19286

Open

4 of 5 tasks

This was referenced Oct 21, 2026

[Feature]: Capture the DeepSeek-V4.1 decoder replay graphs on the model graphs' stream so they share its graph pool #62965

Closed

[Perf][DSv4.1] Capture the decoder replay graphs on the model graphs' stream #63155

Merged

This was referenced Oct 22, 2026

[Audit support] Separate dependency closures for confidence verification and SWA replay zkawfanx/vllm#16

Open

[Start here] Classify the pinned DSv4.1 timeline and prepare chronological backport batches zkawfanx/vllm#18

Open

Klaud-Cold mentioned this pull request Oct 26, 2026

[Klaud Cold] Update DeepSeek-V4.1-Flash FP4 B200 vLLM AgentX DSpark image to v0.32.0 / [Klaud Cold] 将 DeepSeek-V4.1-Flash FP4 B200 vLLM AgentX DSpark 镜像更新至 v0.32.0 SemiAnalysisAI/InferenceX#4523

Merged

Juntian777 mentioned this pull request Oct 29, 2026

[Perf][DSv4.1] Apply the output head to the decoder replay rows only #64466

Draft

6 tasks done

Navigation Menu

Repository navigation

[Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs - #59532

[Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs#59532
ivanium merged 10 commits into
vllm-project:mainfrom
Juntian777:feat/dsv41-replay-cudagraph

Juntian777 commented Oct 1, 2026 •

edited

Loading

Uh oh!

Juntian777 commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

Juntian777 commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

Juntian777 commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Benchmark

Why TP4 gains at long prompts and DEP2 does not

Uh oh!

Juntian777 commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

Juntian777 commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Juntian777 commented Oct 1, 2026 •

edited

Loading