Skip to content

[DSv4.1] Run encoder SWA bounded replay inside the breakable prefill CUDA graph - #44815

Closed
kpham-sgl wants to merge 1 commit into
mainfrom
kpham/dsv41-encoder-replay-bcg
Closed

kpham-sgl wants to merge 1 commit into
mainfrom
kpham/dsv41-encoder-replay-bcg

Conversation

@kpham-sgl

@kpham-sgl kpham-sgl commented Oct 17, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

--enable-encoder-swa-bounded-replay rejected every prefill CUDA graph backend, so DSv4.1 with encoder replay (admitted with DP attention since #43813) ran eager prefill. This admits the breakable backend, the same way #43736 did for the decoder replay.

Modifications

  • Hook: the encoder gate rejects only prefill graph backends other than disabled/breakable.
  • Prefill graph runner: the static batch carries the four encoder_swa_* fold fields.
  • DSv4 backend: BCG capture/replay prep set the fold's row floor and activate the window layout; refresh_for_breakable_cuda_graph_replay_ copies request_window_layout in place; under the graph the layout uses a fixed group count (req_to_token_pool.size) and the workspace is reserved at the first capture.
  • window_layout(pad=...): padding rows (DP fill, graph bucket fill; out_loc == 0) take a spare group and never commit. Grouped with the last request they shortened its commit range and committed garbage at its ring slot 0; this was latent in the eager DP path too.
  • In-graph ratio-1/2 compressors skip the folded replay rows via a bucket-sized DSV4Metadata.low_ratio_replay_rows mask.
  • DP prefill-graph vote counts the folded replay rows so every rank votes on the rows the runner replays.

Accuracy / Performance

Not yet run on GPU; CPU unit tests only (the new cases fail on main). Validation on DSv4.1 Flash with --enable-encoder-swa-bounded-replay --cuda-graph-backend-prefill breakable --attn-dp-size N is pending.

Checklist


CI States

Latest PR Test (Base): ❌ Run #38985250998
Latest PR Test (Extra): ❌ Run #38985250955
Latest PR Test (AMD ROCm 10): ❌ Run #38985251059

…CUDA graph

Admit --enable-encoder-swa-bounded-replay with the breakable prefill backend
(and attention DP). The fold's row floor, replay-row compressor skip and the
request-window layout become bucket-sized metadata refreshed in place per
replay; padding rows take a spare window group and never commit; the DP
prefill-graph vote counts the folded replay rows.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant