Repository navigation
[FlyDSL] pa_decode: zero V rows past the context before PV - #7092
Merged
Merged
Conversation
Pages past a sequence are pinned to block 0 and the unwritten tail of its last page is read as is. Their probabilities are exactly zero, but 0 * NaN is NaN in the PV MFMA, so a NaN in any such row (a recycled block 0, a checkpoint image, a freed page) turned the whole output row into NaN. Per-token V scales had the same hole through P * scale. Only a sequence's last 256-token tile holds such rows. On that tile each lane builds 64-bit byte masks from the context length once, and every head chunk ANDs its V operands with them right before its PV, behind a uniform branch, so full tiles run unchanged and the loads stay pipelined. Per-token K/V scales at or past the context are staged as zero. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
Contributor
|
the pr reduce pa performance and framework should handle the issue |
This was referenced Oct 29, 2026
Contributor
|
Add a option to the api for enable zero V feature. There is a regression here |
mask_unowned_kv (default True) keeps the boundary-tile V/scale masking introduced in this PR. False compiles it out, restoring the previous kernel and the caller contract that unowned V bytes and scales are finite. The flag is part of the schedule cache key; gfx1250 ignores it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fsx950223
approved these changes
Oct 29, 2026
READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
coderfeli
approved these changes
Oct 30, 2026
5 of 11 tasks
This was referenced Oct 30, 2026
Open
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
* [FlyDSL] pa_decode: zero V rows past the context before PV Pages past a sequence are pinned to block 0 and the unwritten tail of its last page is read as is. Their probabilities are exactly zero, but 0 * NaN is NaN in the PV MFMA, so a NaN in any such row (a recycled block 0, a checkpoint image, a freed page) turned the whole output row into NaN. Per-token V scales had the same hole through P * scale. Only a sequence's last 256-token tile holds such rows. On that tile each lane builds 64-bit byte masks from the context length once, and every head chunk ANDs its V operands with them right before its PV, behind a uniform branch, so full tiles run unchanged and the loads stay pipelined. Per-token K/V scales at or past the context are staged as zero. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] pa_decode: add mask_unowned_kv option to toggle V masking mask_unowned_kv (default True) keeps the boundary-tile V/scale masking introduced in this PR. False compiles it out, restoring the previous kernel and the caller contract that unowned V bytes and scales are finite. The flag is part of the schedule cache key; gfx1250 ignores it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
ZhiweiYan-96
pushed a commit
that referenced
this pull request
Nov 2, 2026
…64 query rows) MTP-3 verify at GQA 16 (Qwen3.5-397B at TP2: 16 query heads per KV head) is 4 x 16 = 64 rows, one M-tile past the 48-row limit, so it fell back to the generic multi-M-tile path. The pipe's tile body (single_query_tile / pack_single) is written for any M_TILES, and the #7092 V-tail masks are applied once per head chunk before the M loop; the M_TILES >= 4 V-scale reload exists only in the generic pack_multi. Raise the limit to 4 * MFMA_MNK. Tests: GQA16 q4 and GQA8 q7/q8 pipe cases (64 and 56 rows) with page/tile-edge lengths, a GQA16 q4 np81 and one-hot-token case, the row gate at QL4/6/8 x GQA8 and QL4 x GQA16 (on) vs QL9 x GQA8 and QL5 x GQA16 (off), and a 64-row head-256 case in the unowned-KV test.
ZhiweiYan-96
pushed a commit
that referenced
this pull request
Nov 2, 2026
…64 query rows) MTP-3 verify at GQA 16 (Qwen3.5-397B at TP2: 16 query heads per KV head) is 4 x 16 = 64 rows, one M-tile past the 48-row limit, so it fell back to the generic multi-M-tile path. The pipe's tile body (single_query_tile / pack_single) is written for any M_TILES, and the #7092 V-tail masks are applied once per head chunk before the M loop; the M_TILES >= 4 V-scale reload exists only in the generic pack_multi. Raise the limit to 4 * MFMA_MNK. Tests: GQA16 q4 and GQA8 q7/q8 pipe cases (64 and 56 rows) with page/tile-edge lengths, a GQA16 q4 np81 and one-hot-token case, the row gate at QL4/6/8 x GQA8 and QL4 x GQA16 (on) vs QL9 x GQA8 and QL5 x GQA16 (off), and a 64-row head-256 case in the unowned-KV test.
rebklee
added a commit
to rebklee/aiter
that referenced
this pull request
Nov 3, 2026
Conflicts in traits.py, pa_decode_kernel.py and pa_decode.py: main's ROCm#7092 added mask_unowned_kv where this PR adds causal (the select, compile_pa_decode_tile and pa_decode signatures and the calls between them). Kept both, mask_unowned_kv first; all are keyword-only. The schedule cache key holds both, in the same order where it is built and unpacked. ROCm#7197 (peeled last KV tile) and ROCm#6767 (gfx1250-strict) also touch these files and merged cleanly; neither needs a change for causal. Signed-off-by: Rebecca Lee <rebecca.lee@amd.com>
5 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Symptom
ATOM #2635 found that decode paged attention lets KV rows no sequence owns reach the output. Since ROCm/ATOM#2519, Qwen3.5 GDN prefix checkpoints are copied into MHA KV blocks as PAGE images, and at fp8 those image bytes are often NaN encodings. ROCm/ATOM#2635 works around it in ATOM: pools holding state images never hand out block 0, and released units are zeroed before reuse. As discussed there, this PR fixes it at the kernel level for FlyDSL
pa_decode.On aiter main, Qwen3.5-35B-A3B-FP8 with
ATOM_PA_FLYDSL=1(ATOM main, without ROCm/ATOM#2635) scores 0.06 on GSM8K. The same server with this PR scores 0.97 (numbers below).Root cause
ceil(context_len / 256)tiles per sequence, so only the last tile of a sequence holds tokens at or pastcontext_len. In that tile, pages past the sequence are pinned to block 0, and the unwritten tail of the last page is read as is.-inf.0 * NaN = NaNin the PV MFMA, so a single NaN byte in such a row turns the whole output row into NaN.P * v_scale.pa_decode_kernel.pydocstring stated this as a caller contract: "Masked V bytes must remain finite ... callers must leave the unwritten tail of the last owned page finite."Fix
The change is Python/FlyDSL only, in
kernels/pa_decode/{op_lds,pipeline}.py.tile_valid < 256is wave-uniform. On that tile, each lane builds 64-bit byte masks for its V operands once, from the context length (tail_masks). Each head chunk ANDs its V operands with them right before its own PV (mask_v_chunk).BUFFER_KV(MTP4 fused with a KV cache larger than 2 GiB) goes through the samemask_v_chunk, but the probe does not build a cache that large.Why the overhead is small, and what was tried
Kernel time, fixed vs main, on the matrix described below (513 shapes, 3 alternating rounds, median per shape). All of these produce correct output except the E1 diagnostic.
0xFFFFFFFF >> clamp(32(d+1) - valid_bits, 0, 32): v_add, v_med3, v_lshrrev_b64)vmcnt(0)for the whole tile's V before its first PV, so PV no longer overlaps the V loads. Per-chunk masking keeps the load pipeline intact.vmcntchain as the V loads.q35b-tp1 fp8tok bs16): +0.66% geomean / +2.9% worst at qlen 1, and +1.19% / +2.5% with MTP-3 (qlen 4). Per-token scales are zeroed with a few selects per lane on every tile, outside the boundary branch.Tests
New op test
op_tests/test_flydsl_pa_decode.py::test_pa_decode_unowned_kv_rows, 72 cases:Lengths are {1, 17, 255, 257, 1000}, block 0 is unowned, and the block table has 2 padded columns. Every row no sequence owns is filled (and its per-token scale too, when there is one). The test asserts the output is finite and bit-identical to the zero-filled run.
Existing tests:
test_flydsl_pa_decode.pypasses 124/124 on main and 196/196 here (124 + 72).triton_tests/test_pa_decode_gluon.pypasses 43/43 on both.Wider probe: 468 FlyDSL shapes. On main all 468 fail; with this PR all 468 pass. It covers:
It also hits D256_PIPE, MTP4_FUSED, MTP4_PREFETCH_V, PAGE16_VPIPE, prefetch_v and M_TILES 1–4.
Kernel performance
gfx950 (MI355X), one GPU, main and this PR alternating in the same window, 3 rounds, median per shape. Same-tree noise (main round 1 vs the other two) was -0.29% geomean, +5.1% max.
Shapes:
ctx - U[0, 255], so most sequences end mid-tile.Geomean -0.58%, worst +5.0%; 16 of 513 shapes above +3%. An earlier independent run of the same code gave -0.83% / +4.7% / 13.
The bf16 page 64/128 groups get 3–5% faster. That is a scheduling side effect in kernels that sit at the register limit (fewer
s_waitcnt vmcntin the ISA), not something this change is meant to buy.All 513 shapes
End to end
Setup: Qwen3.5-35B-A3B-FP8 on MI355X.
--state-checkpoint-interval-tokens 256,--gpu-memory-utilization 0.22,ATOM_PA_FLYDSL=1.773eb3625(without fix(qwen3.5): keep PAGE checkpoint images out of the KV rows decode reads ATOM#2635), and fix(qwen3.5): keep PAGE checkpoint images out of the KV rows decode reads ATOM#2635 headcfa9b4c17for comparison.bbef5a6f2or from this PR.Server logs show only
pa_decode -> flydsl[...] (... head_dim=256 torch.float8_e4m3fn)and no gluon fallback, so every full-attention decode call ran FlyDSL.No injection. The NaN comes from Qwen3.5 checkpoint images in the KV pool. Strict-match:
On aiter main, answers start fine and then degenerate into repeated words (
personal personal personal ...), or come back empty. With this PR, ATOM main matches the ROCm/ATOM#2635-guarded result exactly on the full set.Stress run with injection. A local, unpushed ATOM hook fills every MHA K/V cache byte with
0xFF(e4m3fn NaN) and the per-token scales with NaN right after allocation, so every row a decode reads without owning it is NaN:For reference, ATOM's default gluon path (
ATOM_PA_FLYDSL=0) on this 35B TP1 recipe scores 0.970 / 0.970 without ROCm/ATOM#2635. Gluon returns finite wrong values rather than NaN at fp8, so this recipe does not show it; ROCm/ATOM#2635 measured the gluon drop on 397B TP8.What ATOM can simplify
ROCm/ATOM#2635 has three parts. On the FlyDSL path, each one is covered by this PR:
>= context_len, whose V is masked to 0, whose per-token scales are staged as 0, and whose K score is-inf>= context_lenmask on the boundary tilecfa9b4c1)ATOM should keep ROCm/ATOM#2635 as is for now:
ATOM_PA_FLYDSL=1, calls outside FlyDSL's domain fall back to gluon, e.g. bf16 KV or sliding window.mla_decode_fwd) relies on the null block and is not touched here.Once every MHA decode call runs a kernel that masks unowned rows, the three MHA-pool parts can go.
Gluon
pa_decode_gluon(not changed here)gl.where(token < context_len, v, 0)costs +13.4% geomean, +66% worst.cc @zhuyuhua-v @fsx950223 @zhiding512
PR Self-Review
Reviewed revision
795de302Behavior and intent
pa_decodereads KV rows no sequence owns on a sequence's last 256-token tile (padded block-table lanes pinned to block 0, and the unwritten tail of the last page). K is masked with-inf, but V and per-token K/V scales are not, so a NaN byte there turns the output row into NaN (0 * NaNin PV). The kernel pushed this onto callers ("unowned V bytes must remain finite"), and ATOM had to work around it in fix(qwen3.5): keep PAGE checkpoint images out of the KV rows decode reads ATOM#2635; without that workaround Qwen3.5-35B-A3B-FP8 scores 0.06 on GSM8K.e895c9c0): on gfx942/gfx950, the boundary tile ANDs V operands with per-lane byte masks right before each head chunk's PV, and per-token K/V scales at or pastcontext_lenare staged as 0. Full tiles are unchanged; clean-input output is bit-identical to main.795de302(@fsx950223): addsmask_unowned_kv(defaultTrue= the behavior above) topa_decode/compile_pa_decode_tile.Falsecompiles the masking out, restoring the previous kernel and the old caller contract. The flag is part of the schedule cache key. gfx1250 pipelines ignore it and still require finite unowned bytes.Validation and unverified scope
Validated on gfx950 (MI35x), all on
e895c9c0unless noted:pytest op_tests/test_flydsl_pa_decode.pybbef5a6f124 passed;e895c9c0196 passed (124 + 72 new)pytest op_tests/triton_tests/test_pa_decode_gluon.pye895c9c0468 / 468 pass; zero-fill output bit-identical to main 468 / 468ATOM_PA_FLYDSL=1, ATOM main without ROCm/ATOM#2635795de302(MI35X, shard 2)op_tests/test_flydsl_pa_decode.py: 196 passed, including the newmask_unowned_kv=Falseclean-input assertionUnverified:
795de302, I did not re-run the probe, the kernel benchmark or the e2e runs; all numbers above except the CI row are frome895c9c0. The default path is meant to compile to the same kernel, but that is not checked (no ISA diff or timing on795de302). Themask_unowned_kv=Falsepath has no timing data from me, and its only coverage is the CI clean-input equality check.BUFFER_KV(MTP4 fused with a KV cache > 2 GiB): samemask_v_chunk, but no test builds a cache that large.pa_decode_gluonhas the same bug and is not changed here (see above).PR-specific question
test_pa_decode_unowned_kv_rowsfills every row no sequence owns (padded block-table lanes on block 0, the tail of each last page, and their per-token scales) with NaN or Inf, and asserts the output is finite and bit-identical to the zero-filled run. Before this PR it fails with non-finite output on all 36 NaN cases and on the bf16 / per-token-scale Inf cases (54 / 72). The existing canonical testtest_pa_decodefills the whole KV cache, padding included, with finite values (uniform_(-0.5, 0.5)or fixed patterns), so unowned rows contribute0 * finite = 0and the bug cannot show up there; its only NaN use is pre-filling the output buffer to catch unwritten outputs. No other op test exercises FlyDSLpa_decodewith non-finite KV. Themask_unowned_kv=Falseassertion (795de302) checks that the opt-out path still matches on clean input.Changed op tests
op_tests/test_flydsl_pa_decode.pytest_pa_decode_unowned_kv_rows(72 cases: fp8 per-tensor trans_v / plain V, fp8 per-token, bf16; page 16/64/128; 3 head configs incl. qlen 4; NaN / Inf fills; lengths {1, 17, 255, 257, 1000}, block 0 unowned, 2 padded table columns). Asserts unowned rows have no effect on output, and thatmask_unowned_kv=Falsematches the default on clean inputtest_pa_decodecases only use finite KV everywhere, so they cannot detect unowned rows leaking into PV; this is the only test of the NaN-in-unowned-rows contract and of themask_unowned_kvflag. Existing tests are unchanged