[Model] Decoder-side SWA bounded replay for DeepSeek-V4.1 by ivanium · Pull Request #58132 · vllm-project/vllm

ivanium · 2026-09-22T10:07:47Z

Purpose

Decoder-side SWA bounded replay for DeepSeek-V4.1, on top of the encoder side (#56227). The layers past the last KV-source layer (21-39 of 40 on V4.1-Flash) own only sliding-window KV, so in eager prefill steps they run on each request's last 128 tokens, cold prefills included. Simplified re-submission of #56752 (closed); no other open PR implements this.

DeepseekV41ModelState.prepare_attn prepares the kept rows as a sub-batch, the way a microbatch is prepared: one Triton kernel gathers rows, positions, slot mappings and the replay-layer KV start, and the sub-batch gets attention metadata and a forward context of its own. DP ranks agree on whether any rank trims and on every rank's replay token count.
DecoderReplayLayers gathers the layer inputs by those rows, runs the replay layers under that context and scatters the outputs back.
Only eager steps trim. A step that fits a CUDA graph (up to 1024 tokens by default) keeps the replay layers on its whole batch inside the graph: such steps are launch-bound, so the graph is worth more than the trim. prepare_attn already gets the step's graph mode, which DP ranks share, so there are no nested captures, no eager breaks and no runner changes; the diff is confined to vllm/models/deepseek_v41/, a config docstring and the CI job list.
With the fused MoE finalize of [Kernel][DSV4.1] Fuse MoE finalize into the TP all-reduce + mHC boundary #58586, a decode step's cut layer hands its unfinalized MoEOutput to the replay layers unchanged, and layer 21's first boundary finalizes it. A trimming step holds a prefill longer than the window, more tokens than any step that defers its finalize, so the trim path only sees tensors.
Stays off, with a warning, when the PP stage does not hold every layer after the cut, under sequence-parallel MoE, PCP or microbatching, with an Engram layer after the cut, or with a drafter window wider than the target's. --no-swa-bounded-replay turns off both sides. A no-op on ROCm.

Test Plan

pytest tests/models/test_deepseek_v41_replay_batch.py tests/models/test_deepseek_v41_decoder_replay_layers.py

Accuracy: gsm8k 5-shot, n=1319, two passes against one server (pass 2 is all full-prompt prefix hits).
Performance: vllm-bench, random token-id prompts, prefix cache reset before every point.
Both on V4.1-Flash, TP4 and DP4 x TP1 x EP4 (DEP4), decoder side on vs off with the encoder side on in both (off via a local patch, not in this PR).

Test Result

Unit tests cover the replay rows and sub-batch metadata (adaptive verification, prompt-logprobs requests, DP agreement), graph steps getting no replay batch, and the gather/scatter with the indexer buffers. pre-commit and mypy are clean.

Accuracy (GB200, V4.1-Flash)

arm	pass 1 (cold)	pass 2 (prefix hits)
TP4, decoder replay on	0.9249	0.9234
TP4, decoder replay off	0.9303	0.9249
DEP4, decoder replay on	0.9333	0.9287
DEP4, decoder replay off	0.9257	0.9181

exact_match/flexible-extract; one standard error is 0.0072. 0 request errors. On TP4 the fused MoE finalize of #58586 is active, so every decode step passes the cut layer's unfinalized MoE output through the replay layers.

Performance (GB200, V4.1-Flash)

Saturated prefill throughput with decoder replay stays at 83K-86K tok/s on TP4 and 199K-228K on DEP4 at every prompt length, against 68K-54K and 134K-120K without. From 16K up, prefill TTFT drops 25-36% on TP4 and 37-45% on DEP4 at any concurrency, and TPOT drops up to 29% and 39% at high concurrency.

TP4:

ISL	TTFT, 1 request	max conc	prefill tok/s at max conc	OSL 512 out tok/s at max conc	OSL 512 TPOT at max conc
1K	55 / 55 ms	256	83.1K / 68.0K (+22%)	+6%	-2%
4K	206 / 185 ms	256	82.8K / 65.5K (+26%)	+7%	-7%
8K	206 / 191 ms	256	86.0K / 64.9K (+33%)	+20%	-16%
16K	211 / 285 ms	256	85.8K / 63.6K (+35%)	+25%	-20%
32K	423 / 573 ms	128	86.1K / 62.2K (+38%)	+29%	-22%
64K	831 / 1178 ms	64	85.8K / 59.5K (+44%)	+35%	-26%
128K	1667 / 2552 ms	32	85.5K / 54.4K (+57%)	+40%	-29%

DEP4:

ISL	TTFT, 1 request	max conc	prefill tok/s at max conc	OSL 512 out tok/s at max conc	OSL 512 TPOT at max conc
1K	209 / 197 ms	256	198.6K / 132.0K (+50%)	+4%	-1%
4K	212 / 200 ms	256	223.2K / 134.4K (+66%)	+19%	-13%
8K	215 / 213 ms	256	226.2K / 134.0K (+69%)	+34%	-24%
16K	240 / 383 ms	256	228.3K / 133.0K (+72%)	+31%	-22%
32K	479 / 765 ms	128	223.5K / 130.4K (+71%)	+44%	-29%
64K	945 / 1546 ms	64	221.2K / 126.4K (+75%)	+38%	-26%
128K	1946 / 3217 ms	32	213.4K / 119.8K (+78%)	+62%	-39%

Decoder replay on / off. TTFT is the prefill-only p50. Each point is one run of 4 x concurrency prompts with the default max_num_batched_tokens 16384 (per rank on DEP4) and max CUDA graph size 1024; each arm runs on its own GB200 node, in parallel.

Regressions. Only prompts of 8K or less at concurrency 64 or less; the table has every such point on either setup, with decoder replay on vs off.

A single 1K prompt runs the same code as the baseline, inside the CUDA graph: TP4 matches (55 / 55 ms); DEP4 measures within 1-6%. Trimming such steps, as the previous revision did, took TP4's TTFT from 56 to 110 ms.
Single 4K-8K prompts cost ~2.7 ms of CPU per eager step: 0.2 ms preparing the sub-batch, ~2 ms building its attention metadata with the same builders the main batch runs every step, 0.5 ms of gather/scatter. Those steps are launch-bound (~175 ms of kernel launches with the GPU waiting on the CPU), so the trim saves no time there. With both TP4 arms on one node the gap is 0-3%; across nodes in the main runs it measures up to +21% on TP4 and +7% on DEP4, with throughput within 5%.
The larger TTFT p50 swings at concurrency 4-64 change sign between points and setups, with throughput between -5% and +13%.

ISL, conc	TP4 prefill-only TTFT p50	TP4 OSL 512 TTFT p50	TP4 OSL 512 out tok/s	DEP4 prefill-only TTFT p50	DEP4 OSL 512 TTFT p50	DEP4 OSL 512 out tok/s
1K, 1	-1%	-2%	-1%	+6%	+1%	+0%
1K, 4	+5%	-33%	+2%	+5%	+9%	+4%
1K, 16	-12%	+11%	-1%	+5%	-3%	+6%
1K, 64	-18%	+11%	+9%	-13%	+8%	-4%
4K, 1	+11%	+21%	-5%	+6%	+7%	+0%
4K, 4	+6%	+14%	+4%	+12%	-20%	+5%
4K, 16	-21%	+23%	-2%	-13%	-0%	+13%
4K, 64	-22%	+5%	+7%	-41%	-38%	+7%
8K, 1	+8%	+19%	-5%	+1%	+7%	-0%
8K, 4	+9%	+22%	-5%	-22%	+18%	-0%
8K, 16	-23%	-18%	+2%	-36%	-25%	-5%

Same-node check (TP4, both arms on one GB200 node, 10 single prefill-only requests each):

ISL	TTFT p50 on / off
2K	188.6 / 183.1 ms (+3.0%)
4K	193.8 / 190.0 ms (+2.0%)
8K	197.7 / 198.3 ms (-0.3%)
16K	207.8 / 285.8 ms (-27.3%)

AMD (MI355X, ROCm 7.2.3)

A no-op: the ROCm model classes have no decoder_replay_layers, and the encoder side stays off on ROCm as before. The new tests are not in the AMD mirror of the CI job. Checked at 9a86c2c99f against its base 496c647 (the rebase onto main and the eager-only follow-up change only code ROCm does not run), V4.1-Flash TP4 with the recipes.vllm.ai MI355X command minus --no-swa-bounded-replay (DSpark, 5 speculative tokens):

	this PR	base
gsm8k pass 1 (cold)	0.9030	0.8961
gsm8k pass 2 (prefix hits)	0.9045	0.8946

tests/evals/gsm8k/gsm8k_eval.py accuracy (5-shot, n=1319, concurrency 128), within noise. Prefix-sharing prompts of up to 256K tokens served without errors on both trees, with the same prefix-cache hits. The two new test files plus test_deepseek_v41_replay_start.py and test_deepseek_v4_vl_rocm.py: 23 passed.

🤖 Developed with Claude Code; all changes reviewed and tested by the author.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

ivanium · 2026-09-25T23:44:04Z

/ci run

github-actions · 2026-09-25T23:44:22Z

✅ Triggered Buildkite CI #91349 for commit ef562a74aa5c.

mergify · 2026-09-26T19:19:26Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ivanium.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Juntian777 · 2026-09-28T22:57:21Z

+@triton.jit
+def _gather_replay_batch_kernel(
+    query_start_loc_ptr,  # [num_reqs + 1]
+    dropped_before_ptr,  # [num_reqs + 1] rows trimmed before each boundary


_gather_replay_batch_kernel JIT-compiles on the first long eager prefill while serving (jit_monitor warns). It's a bare @triton.jit not registered for JIT warmup, and the startup dummy runs never launch it: graph captures return before the trim, and dummy rows aren't prefills, so trims is False. _pad_replayed_slots_kernel from #56227 has the same issue: dummy batches carry no replay start, so replay_start_np.any() is False. Could we warm both once in the self._replay is None block (it runs during the startup dummy runs), e.g. _gather_replay_batch_kernel.warmup(..., NUM_GROUPS=len(specs), BLOCK=1024, grid=(1,)), or wrap them as VllmTritonJitKernel?

Thx for catching the issue. Fixed now

ivanium · 2026-09-29T02:20:03Z

/ci run

github-actions · 2026-09-29T02:20:22Z

✅ Triggered Buildkite CI #91717 for commit 394e46fd7190.

github-actions · 2026-09-29T07:33:00Z

✅ Triggered Buildkite CI #91784 for commit e0f1dc431dd0, running 1 failed step(s) from Buildkite CI #91717.

ROCm serves DeepSeek-V4.1 through the AMD model classes, so vllm-project#58132 never builds decoder_replay_layers. Construct it on the AMD model and expose it on the vision wrapper. ROCm leaves the prefix-cache window at 0, so the shared model state builds replay_start and warms the gather kernel when decoder replay layers exist. The pad kernel stays on the prefix-cache path. Signed-off-by: Frida Andersson <fanderss@amd.com>

### What this PR does / why we need it? Add decoder bounded replay for DeepSeek V4.1 on Ascend MRV2, following vllm-project/vllm#58132. For eligible prefill requests, layers after the last KV-source layer process only the final sliding-window tokens. Earlier layers process the full batch. Decode tokens and rows required for prompt logprobs are preserved. - Prepare replay batches and independent attention metadata through the MRV2 model-state hook. - Coordinate replay decisions and token counts across DP ranks, including idle ranks. - Gather rows, positions, query boundaries, and slot mappings with a Triton kernel. - Handle sequence-parallel states and DSA CP indexer buffers, then scatter retained outputs and auxiliary hidden states back to their original positions. - Reuse upstream Ascend SWA visibility bounds and preserve encoder-side replay boundaries for prefix-cache hits. The implementation targets MRV2 only and adds no MRV1 runner changes. ### Does this PR introduce _any_ user-facing change? Yes. Decoder replay is controlled by `swa_bounded_replay=True` in `LLM`, or `--swa-bounded-replay` in the CLI. Supported configurations require MRV2 and `NONE` or `FULL_DECODE_ONLY`. Trimming runs only in eager steps; graph steps retain the full batch. Prefix caching uses the existing encoder-side bounded replay support. Sequence parallelism requires DSA CP. Speculative decoding is limited to DSpark with a compatible sliding window. MRV1, PP, PCP, DCP, microbatching, LoRA, and configurations with Engram layers after the replay boundary disable decoder replay with a warning. ### How was this patch tested? Waiting to be tested. - vLLM main: vllm-project/vllm@c4df37d Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

…roject#17496) ### What this PR does / why we need it? Add decoder bounded replay for DeepSeek V4.1 on Ascend MRV2, following vllm-project/vllm#58132. For eligible prefill requests, layers after the last KV-source layer process only the final sliding-window tokens. Earlier layers process the full batch. Decode tokens and rows required for prompt logprobs are preserved. - Prepare replay batches and independent attention metadata through the MRV2 model-state hook. - Coordinate replay decisions and token counts across DP ranks, including idle ranks. - Gather rows, positions, query boundaries, and slot mappings with a Triton kernel. - Handle sequence-parallel states and DSA CP indexer buffers, then scatter retained outputs and auxiliary hidden states back to their original positions. - Reuse upstream Ascend SWA visibility bounds and preserve encoder-side replay boundaries for prefix-cache hits. The implementation targets MRV2 only and adds no MRV1 runner changes. ### Does this PR introduce _any_ user-facing change? Yes. Decoder replay is controlled by `swa_bounded_replay=True` in `LLM`, or `--swa-bounded-replay` in the CLI. Supported configurations require MRV2 and `NONE` or `FULL_DECODE_ONLY`. Trimming runs only in eager steps; graph steps retain the full batch. Prefix caching uses the existing encoder-side bounded replay support. Sequence parallelism requires DSA CP. Speculative decoding is limited to DSpark with a compatible sliding window. MRV1, PP, PCP, DCP, microbatching, LoRA, and configurations with Engram layers after the replay boundary disable decoder replay with a warning. ### How was this patch tested? Waiting to be tested. - vLLM main: vllm-project/vllm@c4df37d Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

…ct#58132) Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

…roject#17496) ### What this PR does / why we need it? Add decoder bounded replay for DeepSeek V4.1 on Ascend MRV2, following vllm-project/vllm#58132. For eligible prefill requests, layers after the last KV-source layer process only the final sliding-window tokens. Earlier layers process the full batch. Decode tokens and rows required for prompt logprobs are preserved. - Prepare replay batches and independent attention metadata through the MRV2 model-state hook. - Coordinate replay decisions and token counts across DP ranks, including idle ranks. - Gather rows, positions, query boundaries, and slot mappings with a Triton kernel. - Handle sequence-parallel states and DSA CP indexer buffers, then scatter retained outputs and auxiliary hidden states back to their original positions. - Reuse upstream Ascend SWA visibility bounds and preserve encoder-side replay boundaries for prefix-cache hits. The implementation targets MRV2 only and adds no MRV1 runner changes. ### Does this PR introduce _any_ user-facing change? Yes. Decoder replay is controlled by `swa_bounded_replay=True` in `LLM`, or `--swa-bounded-replay` in the CLI. Supported configurations require MRV2 and `NONE` or `FULL_DECODE_ONLY`. Trimming runs only in eager steps; graph steps retain the full batch. Prefix caching uses the existing encoder-side bounded replay support. Sequence parallelism requires DSA CP. Speculative decoding is limited to DSpark with a compatible sliding window. MRV1, PP, PCP, DCP, microbatching, LoRA, and configurations with Engram layers after the replay boundary disable decoder replay with a warning. ### How was this patch tested? Waiting to be tested. - vLLM main: vllm-project/vllm@c4df37d Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

上游来源：vllm-project/vllm#58132 (4b2e1cfa7b)，先带入自包含的 DecoderReplayLayers 与其单测；model_state/model.py 的接线在后续提交完成。本仓 ForwardContext 已含 is_padding，测试无需改写。 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#58132 (4b2e1cfa7b)，适配点： - model_state 按上游整文件合入（含 replay batch gather kernel / kept InputBatch / DP 同步）；本仓 InputBatch 没有 fast_prefill 字段，去掉该 replace 参数； - model.py 按本仓 layer 返回 5 元组 + aux 重算的 API 重写 _run_layer_range / _collapse / _run_replay_layers（上游为 previous_aux 返回式 API），forward 在 decoder_replay_start 处切分，eager prefill 交给 DecoderReplayLayers； - 补上 dp_utils 的 skip_dp_coordination/should_skip_dp_coordination（本仓缺失）； - vl_model 暴露 decoder_replay_layers； - 上游 test_deepseek_v41_replay_batch.py 带入（7 passed + decoder 层 3 passed）。 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 ci/build deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models labels Sep 22, 2026

ivanium marked this pull request as ready for review September 22, 2026 18:11

ivanium requested review from AndreasKaratzas, DarkLight1337, Harry-Chen, heheda12345, khluu and ywang96 as code owners September 22, 2026 18:11

claude Bot reviewed Sep 22, 2026

View reviewed changes

ivanium force-pushed the feat/dsv41-decoder-replay-eager branch 2 times, most recently from 1652023 to 9a86c2c Compare September 22, 2026 18:16

wangyicong52 mentioned this pull request Sep 24, 2026

[Bugfix][DSv4.1] Preserve partial external hits under bounded replay #58504

Closed

ivanium added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 25, 2026

WoosukKwon approved these changes Sep 26, 2026

View reviewed changes

mergify Bot added the needs-rebase label Sep 26, 2026

ivanium force-pushed the feat/dsv41-decoder-replay-eager branch from 6709ae4 to 69d3837 Compare September 26, 2026 23:19

mergify Bot removed the needs-rebase label Sep 26, 2026

This was referenced Sep 28, 2026

Closed ivanium/vllm#54

Closed

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

Merged

[RFC]: DeepSeek-V4.1-Flash performance on MI325X (gfx942) #59038

Open

Juntian777 reviewed Sep 28, 2026

View reviewed changes

ivanium force-pushed the feat/dsv41-decoder-replay-eager branch from e8b989d to bcac4aa Compare September 28, 2026 23:52

ivanium merged commit 4b2e1cf into vllm-project:main Sep 29, 2026
13 checks passed

ivanium deleted the feat/dsv41-decoder-replay-eager branch September 29, 2026 08:28

wangyicong52 mentioned this pull request Sep 29, 2026

[Bugfix][DSv4.1] Skip SWA bounded replay for KV loads that carry the window #59197

Merged

yiz-liu mentioned this pull request Sep 29, 2026

[Feature][Model] Support DeepSeek V4.1 decoder bounded replay vllm-project/vllm-ascend#17496

Merged

Juntian777 mentioned this pull request Oct 1, 2026

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

Merged

gcanlin mentioned this pull request Oct 3, 2026

[Perf][DSv4.1] Trim the last KV-source layer in decoder replay #59894

Merged

xijiaat mentioned this pull request Oct 14, 2026

[DSv4.1][PCP] Share CPU-offloaded Engram tables #61554

Open

6 tasks done

LopezCastroRoberto mentioned this pull request Oct 19, 2026

[Perf] Enable DeepGEMM sparse MQA logits on the FP8 indexer cache, DSv4.1 #62011

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

Navigation Menu

Repository navigation

[Model] Decoder-side SWA bounded replay for DeepSeek-V4.1 - #58132

[Model] Decoder-side SWA bounded replay for DeepSeek-V4.1#58132
ivanium merged 5 commits into
vllm-project:mainfrom
ivanium:feat/dsv41-decoder-replay-eager

ivanium commented Sep 22, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

ivanium commented Sep 25, 2026

Uh oh!

github-actions Bot commented Sep 25, 2026

Uh oh!

mergify Bot commented Sep 26, 2026

Uh oh!

Juntian777 Sep 28, 2026

Uh oh!

ivanium Sep 29, 2026

Uh oh!

ivanium commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

ivanium commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Accuracy (GB200, V4.1-Flash)

Performance (GB200, V4.1-Flash)

AMD (MI355X, ROCm 7.2.3)

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

ivanium commented Sep 25, 2026

Uh oh!

github-actions Bot commented Sep 25, 2026

Uh oh!

mergify Bot commented Sep 26, 2026

Uh oh!

Juntian777 Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

ivanium Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

ivanium commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ivanium commented Sep 22, 2026 •

edited

Loading