[Feat][Model] DeepSeek-V4.1 decoder-side SWA bounded replay

[ivanium]

Purpose

Stacked on #56227 (base branch feat/dsv41-swa-bounded-replay, so the diff shows only the two commits of this PR). It completes the DeepSeek-V4.1 SWA bounded replay design with its decoder side.

Layers past the last KV-source layer (21-39) read the layer-20 compressed KV and indexer through kv_source and own only their 128-token sliding-window KV, of which decode reads just the trailing window. In prefill they now run on each request's last 128 tokens (its replay window) instead of the whole prompt: after layer 20 the batch is compacted to those rows, the mHC states are gathered, the replay layers run under attention metadata rebuilt for the compacted batch, and the results are scattered back so sampling and the drafter see full-batch rows. A trimmed request's replay_start rises to its window start, so its window is floored there exactly like a replayed prefix-cache hit. The top-k indices and candidate blocks that layer 20's indexer publishes for the layers after it are keyed by batch row, so they are realigned to the compacted rows as well.

The replay metadata is built by private clones of the runner's builders (AttentionMetadataBuilder.clone): the runner's builders back the full batch's metadata with persistent buffers, and that metadata outlives the forward (the runner hands it to the speculator). The three DeepSeek-V4.1 metadata classes carry common and builder back-references for this, and CommonAttentionMetadata.replace_tokens produces the copy over the compacted token set.

CUDA graphs: DeepSeek-V4.1 uses breakable piecewise graphs by default, so the replay is an eager break of the model graph that re-plans from each step's batch, and the replay layers run their own breakable graphs keyed by the padded replay size, captured while the outer capture is paused. Decode batches and full graphs never trim. The compacted slot mappings live in fixed buffers because the window KV insert inside the replay graph reads them by address.

Per-step plan data (rows, query boundaries, window starts) is laid out on the CPU and read by the GPU through the runner's UvaBufferPool, so no pinned memory is allocated and no H2D copy is issued per step. The new per-step GPU work is the gathers and scatters, one torch.maximum, and the rebuilt replay-layer metadata.

Second commit: the DSpark drafter's layers are all sliding-window attention, so it can only read the last window context positions of a request and its context-KV precompute now runs on those rows only (the DFlash precompute gains a row-selection hook that keeps every row by default). Under decoder trimming the other rows carry no real hidden state, so this also stops the drafter from inserting KV derived from them. It rides the same switch.

Switch: part of --swa-bounded-replay (on by default), since replay-layer window KV is unwritten outside each request's replay window and a prefix-cache hit must therefore replay the window, which is what the encoder side does. The decoder side stays off with one warning when the model cannot trim safely: a pipeline stage that holds the last KV-source layer but not every layer after it, sequence/data/prefill-context parallelism (the per-rank batch would shrink), an Engram layer after the cut, or a drafter without a sliding window, with a wider one than the target's, or with mixed layer types. Adaptive verification is supported: it splits only the leading verification requests on the GPU, which are never trimmed, so the row selection from the CPU query boundaries holds and the replay batch takes its boundaries from the device query_start_loc minus the rows trimmed before them.

Numerics: this changes cold-prefill outputs for every prompt longer than 128 tokens. Within a request's trailing window the replay layers' SWA attention is floored at the window start, so the first rows of the window attend fewer keys than in a full prefill, and the final token's replay-layer KV inputs come from those rows. The two-pass evaluation below covers both the cold path (pass 1) and the prefix-hit path (pass 2).

Known gaps: prompt logprobs of rows outside the trailing window are computed from scattered zeros, so requests that need prompt_logprobs should run with --no-swa-bounded-replay for now. A DSpark draft config without a sliding_window field turns the decoder side off with the warning above.

Third commit: the DFlash graph-capture path built the draft attention metadata without the dummy batch's replay_start, so a DSpark drafter (whose sliding-window caches replay like the target's) failed the SWA builder's assertion at capture under --swa-bounded-replay. It now passes the zeroed replay_start, as the runner's own capture and the drafter's per-step build already do. This is a gap of #56227 with DSpark and could land there instead.

Not a duplicate: no open PR trims the DeepSeek-V4.1 replay layers; #56227 is the encoder side this builds on.

Usage

Nothing to configure: with #56227, DeepSeek-V4.1 on model runner V2 replays on both sides by default. --no-swa-bounded-replay turns both off.

Test Plan

pytest tests/models/test_deepseek_v4_decoder_replay_layers.py tests/v1/spec_decode/test_dspark_context_rows.py \
  tests/v1/cudagraph/test_breakable_cudagraph.py tests/v1/attention/test_deepseek_v4_swa_visible.py

Accuracy: deepseek-ai/DeepSeek-V4.1-Flash on TP4 GB200 with vllm serve, gsm8k (5-shot) run twice against the same server; pass 2 re-sends every prompt, so each request is a full-prompt prefix hit. Both arms on the same tree. Plus a local cold-prefill comparison, replay on vs off, on 12 prompts of 73 to 5415 tokens (the long ones span several prefill chunks), greedy, 32 output tokens; run once plain and once with the checkpoint's built-in DSpark drafter ({"method": "dspark", "num_speculative_tokens": 5, "enable_adaptive_verification": true}), which exercises the drafter's window-only context KV and adaptive verification's GPU-side query boundaries.

Test Result

50 tests pass; pre-commit and mypy are clean.

gsm8k, n=1319 (pass 1 cold / pass 2 all hits)	replay on (this PR + #56227)	replay off (--no-swa-bounded-replay)
exact_match (flexible-extract)	0.9280 / 0.9333	0.9318 / 0.9242

All four numbers are within run-to-run noise of each other. The plain cold-prefill comparison produced coherent, correct completions on both sides; 10 of 12 first tokens were identical (the two exceptions are a greedy near-tie and the replay-off side emitting EOS), and the multi-chunk prompts agreed on all 32 tokens. With the DSpark drafter and adaptive verification, all 12 first tokens were identical and every completion coherent and correct; 8 of 12 agreed on all 32 tokens, the rest diverged after a greedy near-tie.

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

[mergify]

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

[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.