Perf: assemble SWA decode window in L1, drop the GM gather - #882
Perf: assemble SWA decode window in L1, drop the GM gather#882Hzfengsy wants to merge 2 commits into
Conversation
- Assemble the decode SWA window directly in L1 instead of gathering it into GM scratch first. The window is WIN consecutive positions and WIN == BLOCK_SIZE, so it is always the tail of one paged block plus the head of the next; two pl.gather_row calls carry those runs into one pl.create_l1 tile, with the runtime split expressed through valid_shape so the cube still sees a single ATTN_K_TILE-column matmul. - Delete the swa_gather_kv scope, its GATHER_* tiling constants and asserts, and the T*WIN GM staging tensor it fed. - Fill window rows past the visible prefix from a live pool row: they are allocated but not yet written at cold start, and an unwritten cache line reaching the QK matmul would propagate through row_max before sparse_bias could mask its score. - Pass swa_lens into sparse_attn_swa to bound that prefix. qk_pv completes 36.2 -> 32.0 us and merge_norm starts 43.7 -> 40.7 us with the 9.5 us swa_gather_kv scope removed; the end-to-end device wall is unchanged within run-to-run spread. Measured on decode_sparse_attn_swa.py -p a2a3 -d 0, Ascend 910B1, median of 5 runs.
📝 WalkthroughWalkthroughThe SWA decode attention path now passes ChangesSWA window flow
Estimated code review effort: 4 (Complex) | ~30 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4e355fcb0d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| qk_kv = pl.gather_row(qk_kv, ori_kv_flat, [qk_len, 0], [0, 0], | ||
| [ATTN_K_TILE, HEAD_DIM], valid_shape=[qk_tail, HEAD_DIM]) |
There was a problem hiding this comment.
When swa_lens[qk_t] < WIN, this does not replicate the live pool row at source offset zero: gather_row copies qk_tail consecutive rows beginning there. For example, at cold start with qk_len == 1, it reads source rows 0 through 126, most of which may themselves be unwritten cache storage. Non-finite values from those rows can therefore reach the QK matmul and propagate through row_max despite the additive mask, defeating the stated purpose of this overwrite; initialize the tail explicitly or source every tail row from storage known to be initialized.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 14f3863.
The call passed src_offset=[0, 0] with a runtime valid_shape, so it
copied pool rows [0, qk_tail) — a consecutive run, not a replicated
row, exactly as described. Those rows are no better initialized than the
ones they replaced, so the guarantee the comment claimed was not
delivered. That was also a regression against the previous GM gather,
which zero-filled invalid slots explicitly.
The tail is now filled by replicating qk_slot0, the window's oldest
visible slot, which is written by definition since it is a position the
request has already produced. Costs a pl.range(qk_tail) loop of
single-row gathers, taken only when the window is shorter than WIN.
Not caught by the tests because init_ori_kv fills the whole pool with
torch.rand, so every source row is finite in the fixture regardless.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/deepseek/v4-flash/decode_sparse_attn_swa.py (1)
659-679: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test fixture that exercises the block-crossing gather path.
init_swa_indicesalways setsintra = w % BLOCK_SIZEforwstarting at 0, soindices[t, 0]always has intra-block offset 0 regardless ofblk. This meansqk_ris always 0 in this fixture, and theif qk_r > 0:branch insparse_attn_swa(the new two-run assembly this PR introduces) is never exercised by the standalone test.Add a fixture variant where the window start is not block-aligned, so the second
pl.gather_rowcall is actually tested.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/deepseek/v4-flash/decode_sparse_attn_swa.py` around lines 659 - 679, Add a non-block-aligned SWA index fixture alongside init_swa_indices, using a window start offset within BLOCK_SIZE so the generated indices cross a block boundary and produce qk_r > 0. Update the standalone sparse_attn_swa test to use this variant and exercise both pl.gather_row calls, while preserving the existing aligned fixture coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@models/deepseek/v4-flash/decode_sparse_attn_swa.py`:
- Around line 659-679: Add a non-block-aligned SWA index fixture alongside
init_swa_indices, using a window start offset within BLOCK_SIZE so the generated
indices cross a block boundary and produce qk_r > 0. Update the standalone
sparse_attn_swa test to use this variant and exercise both pl.gather_row calls,
while preserving the existing aligned fixture coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 575d7ec6-7a08-4a46-8a3b-ac8b84956f7d
📒 Files selected for processing (2)
models/deepseek/v4-flash/decode_attention_swa.pymodels/deepseek/v4-flash/decode_sparse_attn_swa.py
The tail fill sourced qk_tail consecutive rows starting at pool row 0, so it replaced unwritten rows of the current block with an equally unwritten run elsewhere in the pool rather than with known-good data. Replicate the window's oldest visible slot instead: that row is written by definition, being a position the request has already produced. sparse_bias masks these rows either way, but a non-finite value reaches the QK matmul and propagates through row_max before the additive mask can apply, so the fill must come from initialized storage.
1 participant
into GM scratch first. The window is WIN consecutive positions and
WIN == BLOCK_SIZE, so it is always the tail of one paged block plus
the head of the next; two pl.gather_row calls carry those runs into
one pl.create_l1 tile, with the runtime split expressed through
valid_shape so the cube still sees a single ATTN_K_TILE-column
matmul.
asserts, and the T*WIN GM staging tensor it fed.
visible prefix. Those rows are allocated but not yet written at cold
start, and a non-finite value there reaches the QK matmul and
propagates through row_max before sparse_bias can mask its score.
qk_pv completes 36.2 -> 32.0 us and merge_norm starts 43.7 -> 40.7 us
with the 9.5 us swa_gather_kv scope removed; the end-to-end device wall
is unchanged within run-to-run spread. Windows shorter than WIN pay the
replication loop instead, moving qk_pv 38.2 -> 40.2 us. Measured on
decode_sparse_attn_swa.py -p a2a3 -d 0, plain and
--short-window-fixture, Ascend 910B1, median of 5 runs.