Repository navigation
Conversation
Three rejection tests for the preconditions the sparse ASM relies on to skip apply_mask: ragged key length, and per-batch seqlens_k on the raw and packed entry points. They match on "sorted-sparse" because _check_varlen_capable's message also contains "per-batch key lengths", so a looser match passes with the guard removed. _plain_block_scale is safe for K not because the gather stays inside the logical rows -- it does not, the K pipeline prefetches one tile past the end -- but because that gather goes through a buffer descriptor whose num_records clamps the overshoot in hardware.
downcast_to_mxfp is used by sage_quant_mxfp4 alone, but importing it at module scope runs aiter/ops/triton/moe/__init__.py, which eagerly builds the moe_wgrad autotune decorator and its tuned-config lookup. Every MHA v4 caller paid that: mha_v4 -> mha_v4_quant -> these wrappers, for five FP4 sizing constants that have nothing to do with MoE. Move the import into the one function that needs it. sage_quant_mxfp4 is unchanged; the MoE namespace is now touched only when it actually runs.
Thread an optional lse tensor through fmha_v4_fwd_sparse (cu, header, pybind, the custom op and both fakes), and share one populate_lse_kernarg between the dense and sorted-sparse entries so the two cannot drift. _check_lse_capable loses its `sparse` parameter, leaving _LSE_CAPABLE_QV as the single gate for both paths. Refresh the gfx950 fmha_v4 objects: every row now carries the LSE epilogue, and a LUT row that selects no KV tile exports -inf instead of a finite value. Measured against torch.logsumexp over the selected blocks: bf16 0.0009 nats, bf16fp8 0.03, fp8 0.05. Adds 16 sparse-LSE tests covering a per-selection reference and the empty-row case; 358 passed, 1 skipped.
Pure SGPR renumber across the eight quantized hd128 kernels. The softmax-skip threshold pair now shares the block-sparse LUT base registers -- the two features are mutually exclusive and asserted so -- which keeps the LSE descriptor clear of the LUT state and leaves a 4-aligned scalar quad that no build toggle can claim. No instruction is added, removed or reordered, so the objects are numerically unchanged: outputs are bitwise identical to the previous ones on every shape checked, 88 dense and 770 sparse validation cases pass, and the op tests report 358 passed, 1 skipped.
return_lse exports the logsumexp of the scores the kernel computed, not of the exact-math scores. Error tracks the row max at each format's score precision: ~0.001 nats for BF16 Q/K, 36 for MXFP6, 153 for MXFP4, against 1418 nats of range. Safe for merging same-recipe ring chunks, wrong as an absolute quantity or across mixed recipes. Documented in both docstrings and the .md. Also restore black/ruff cleanliness after a11f99f: three lse: Optional[Tensor] annotations needed # noqa: UP045 (the union form graph-breaks Dynamo) and two files needed reformatting.
Three distributions now span how much attention averages over V rows, which is what decides whether V quantization error survives into the output. Measured at b1/h8/s8192 with Q/K pinned MXFP6, as a fraction of V's own round-trip error: diffusion 0.35x (n_eff 191), peaked 0.82x (n_eff 13), tiebreak 0.98x (n_eff 1.6). `peaked` widens the logit spread onto ~10-20 effective tokens. Against 74 dumped heads in that n_eff band it lands at percentile 51 on V amplification (0.82 vs their median 0.81) - a typical head, not an optimistic one. It does not reach the 43% of that band which amplifies above 1x. `tiebreak` drives attention onto ~2 near-tied keys, so nothing averages and the output inherits V's error whole. The tie gaps are spread geometrically across the formats' mantissa resolution: with a single gap no format resolves the pair, each breaks the tie by whichever way its rounding falls, and the resulting ranking is arbitrary - MXFP8 Q/K measured 0.025 below FP8 Q/K at s=65536 but 0.009 above it at s=16384. With six spread gaps --kernel=all is monotone in precision at s=65536, s=8192 and s=4096. --compare-to-ref gets a wide 25% bound for tiebreak, which diverges from BF16 by construction; it is a backstop against gross breakage, not a quality bar.
bf16fp8 could emit non-finite output on long sequences: a 75600-token head produced NaN on 18478 of its rows, which surfaced end to end as a black video. The frozen-max rollback rebases by shifting the scores exactly again, and the P pack divisor is now held in an SGPR for the life of the kernel. Validated over 64 captured 75600-token attention heads: all rows finite, worst cosine 0.999409 against the bf16 kernel. Dense and sparse recipe checks pass, including the bit-exact sparse cases. bf16 picks up the K-wait drain and softmax-tail scheduling from the same sweep.
Four V LDS slots with V loaded two tiles ahead, closing the cross-wave V publication race: the consuming band read a direct-to-LDS payload one barrier generation before the producing band's own vmcnt(0) drained it. Three slots are not enough -- the refill then targets the buffer the PV phase read earlier in the same matrix phase. Race under controlled HBM pressure, interleaved with the previous objects and order rotated: f8f6 50-62 -> 1 distinct output, mxfp6 50-59 -> 1, f6f4 54-64 -> 1, mxfp4 49-61 -> 1. Dense 24/24 ragged sweep and sparse 70/70 incl. bitexact for every recipe. b1/hq5/sq65536/d128 on an idle board: f8f6 -0.50%, mxfp6 -0.24%, f6f4 +0.04%, mxfp4 +0.84%. Occupancy is unchanged (VGPR-bound); only LDS grows, by one V slot.
Three V LDS slots and a retuned pad. Race probe states=1 at sk 384/512/576 under pressure; dense 264/264, sparse 770/770; quality cosines unchanged. 2379.0 vs 2378.7 TFLOPS at hq=40/sq=65536.
The Q-scale E8M0 exponent fold now guards against underflow: a Q block below 2^-118 previously wrapped past zero and clamped to the MAXIMUM exponent instead of the minimum. Alignment pads retuned for the resulting prologue size. mxfp6 sparse +1.2%; the other three objects hold their previous throughput. Dense 24/24, sparse 70/70 on both recipes.
Score conversion leaves the VALU: the i32 QK accumulator is seeded with the bit pattern of 1.5*2^23, so 64 v_cvt_f32_i32 per tile become two bf16 MFMAs (+5.07%). Both L matmuls issue at the PV front interleaved with the V LDS reads (+0.68%), and the alignment pad is retuned (+0.74%). 2376 -> 2534 TFLOPS. Dense 24/24, sparse 70/70, bitwise-stable over 256 launches.
Fixes two s_waitcnt lgkmcnt races where a K operand could be consumed before its LDS read had landed; output was bit-identical either way, so the window was closed by instruction timing rather than by construction. Enables the dynamic P pack gain, hoists the gain compare into the freeze gate's branch hazard shadow, and reclaims a dead VGPR pair (248 -> 246). Dense 24/24 and sparse 70/70 vs the torch reference; +0.21% at b1/hq40/sq65536/d128.
GPU7 paired +2.09% (20/20); dense 24/24, sparse 70/70.
GPU7 paired dense +1.10%, sparse +0.27%; dense 24/24, sparse 70/70.
gemm_QK skipped the lgkm wait for a PV-prefetched K subtile 0 on the assumption that "resident" implied landed. Nothing drains those reads between gemm_PV issuing them and the first QK MFMA consuming them, so the subtile and its descale were read racing. Emit the wait; the running budget is already exactly the queued reads. The VGPR file is reordered into QK sources -> PV sources -> matrix destinations so that no bank is both a matrix source and a matrix destination: the softmax denominator MFMA now owns its own accumulator instead of borrowing K bank 0, which lets the next tile's subtile 0 prefetch straight into bank 0 and retires both the separate prefetch bank and the phantom lgkm credit it needed. The fp6 block scales also move out of the shared scratch range, where the QK bias split clobbered them immediately after the pack. 238 VGPRs, still 2 waves/SIMD. Dense and sparse validated 24/24 and 70/70, race-free, output digests bit-identical to the previous objects. Dense is at parity on b1/hq5/sq65536/d128; sparse is ~2% behind, accepted for the race fix.
…n the fp6-P recipes Rebuilt across all nine hd128 recipes (dense and sparse). 1) The epilogue guard that substitutes 0 for a zero denominator compared L as an integer, so it matched only the bit pattern +0.0 and let -0.0 reach v_rcp_f32, yielding -Inf instead of 0. Now a float compare, which matches both zeros. Affects f4f4, f6f4, f6f8, f8f6, fp8, i8fp8, mxfp4, mxfp6 and mxfp8; bf16 and bf16fp8 already did this correctly. A tiny-but-nonzero L still reaches the reciprocal, which is intended. 2) f6f4, mxfp4 and mxfp6 pack P as fp6 with a per-K-block E8M0 scale, so the pack re-centres on each block's own peak and the old fp8 e4m3 ceiling no longer binds their softmax overflow-gate threshold. Raised to P = 2^32, matching f8f6. The recipes that pack P as fp8 are unchanged. 3) f4f4 also drops a dead per-row bias left over from before the exp-fold; with the running max frozen the QK prefill already biases the accumulator. This is the only size change here, -32 bytes. Every object: emitted assembly differs only in the intended instructions, 24/24 dense and 70/70 sparse correctness cases.
Both recipes now stage P through the pknorm conversion before reducing it, so the softmax block max, the overflow gate and the rare rollback all run on packed 16-bit patterns instead of f32 coordinates. The conversion is monotone, so the reduction yields the same block max as before; mxfp4 additionally uses the 3-input f16 max and parks P = 1 an octave band lower to keep that form well inside its valid range. f6f4 also moves its V priming reads and the V-slot m0 advance onto the QK tail MFMA shadows, holds the V-scale lane term in a resident register, and retunes its code-placement padding for both arms. Measured on gfx950 at b1 hq5 sq65536: mxfp4 and f6f4 both gain several percent over the previous objects. Register and LDS footprints are unchanged (mxfp4 240/90, f6f4 252 dense and 256 sparse). op_tests/test_mha_v4.py and test_mha_v4_sparse.py pass for both recipes.
Dense +3.40% and sparse +3.87% on b1 hq5 sq65536 d128, GPU 7, within-run paired over 15 and 12 rounds. The gain comes from moving the V primes and the m0 slot advance onto the QK tail shadows, hoisting the V-scale lane term into a resident VGPR, correcting the PV lgkm credits for the next-tile K0/scale prefetch, and retuning the code-alignment pad for both arms. Resources unchanged: dense 252 VGPR / 90 SGPR / 75,776 B LDS; sparse 256 / 98 / 108,552 B. Validated: test_mha_v4.py and test_mha_v4_sparse.py 361 passed / 1 skipped; the dense KV sweep 24/24 including ragged lengths; sparse 70/70 covering empty rows, discontinuous LUTs and GQA; a static lgkm race check clean on both arms; 400 and 60 frozen relaunches at sq 8192 and 65536 with zero nondeterminism; and the sparse empty-row case over 10 fresh processes. Output is bit-identical to the previous objects on synthetic shapes and on 288 of 292 real captured heads. The 4 that differ are massive-activation heads where the softmax rollback fires; they move by ~1e-7 in cosine (0.9993 against a 0.97 floor), from the rollback now rounding its octave shift in the integer domain.
Reschedules the V-slot m0 rotation onto the QK MFMA shadows and retunes the hot-loop alignment pad. Output is bit-identical on 6 shapes incl. 2 ragged. b1 hq5 sq65536 d128 on MI350, paired across 6 GPUs: dense +1.38%, sparse +0.40%. Validation: 361 passed / 1 skipped (test_mha_v4 + test_mha_v4_sparse), 24/24 dense sweep incl. ragged sk, 70/70 sparse, 0/400 and 0/60 mismatching frozen relaunches at sq 8192 and 65536, 10/10 sparse empty-row rounds in fresh processes, and the static LDS race check clean on both arms.
+0.38% dense / +0.61% sparse at b=1 hq=5 sq=65536 d=128, measured across the full code-placement sweep rather than a single layout. 24/24 dense and 70/70 sparse validation, race-clean, 361 mha_v4 tests pass. Quality over 288 captured heads: 274 metric triples unchanged, the rest within rel MAE +0.30% and cosine 2.3e-5, no non-finite outputs.
The line passes on main and was introduced by this branch, so CI's black job would have failed on it. Signed-off-by: jcaraban <jcaraban@amd.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Sparse GQA with LSE can fail at runtime, and architecture-dependent tests and benchmark head mapping need correction.
4 open issues
What changed in this PR
This PR extends the dense LSE support from #5798 to sorted-sparse MHA v4 on gfx950, enabling partial attention results to be merged for ring and context parallelism.
Changes:
- Threads optional LSE output through Python, custom-op, and C++ sparse launch interfaces.
- Adds sparse LSE and rejection tests, plus peaked and tied-key benchmark inputs.
- Defers a MoE import and clarifies LSE semantics and scale-buffer documentation.
| File | Description |
|---|---|
op_tests/test_mha_v4.py |
Updates the unsupported-LSE rejection case. |
op_tests/test_mha_v4_sparse.py |
Adds sparse LSE checks and host-guard tests. |
op_tests/op_benchmarks/triton/bench_sage.py |
Adds peaked and tied-key input distributions. |
csrc/pybind/mha_v4_fwd_pybind.cu |
Exposes optional sparse LSE output. |
csrc/py_itfs_cu/asm_mha_v4_fwd.cu |
Shares LSE validation and populates sparse launch arguments. |
csrc/include/torch/mha_v4_fwd.h |
Adds the optional sparse LSE parameter. |
aiter/ops/triton/quant/sage_attention_quant_wrappers.py |
Defers the MoE quantizer import. |
aiter/ops/mha_v4.py |
Enables sparse LSE and updates custom-op mutation declarations. |
aiter/ops/mha_v4.md |
Documents quantized-score LSE semantics. |
aiter/ops/mha_v4_quant.py |
Clarifies scale-buffer bounds in a docstring. |
Review effort: Balanced
Contributor
Author
@azaidy by different backends you mean mixing Triton into the PR? the issue is that these ASM Sage kernels are benched with the same Triton bench as the Triton Sage kernels. Are you simply asking that the changes to |
…hitecture Review findings on the LSE work, plus two of the same kind found alongside them. _restore_k_mean_in_lse contracted a per-KV-head mean against Q, which is a shape error under GQA. Every quantized recipe derives a k_mean, so return_lse with hq != hk raised `einsum(): subscript h has size ...` on the dense path too -- this predates the sorted-sparse work and was simply never covered. Expand the mean over the query-head group first, the way the kernel maps query head j to KV head j // ratio. New dense and sparse cases over six and three recipes at ratios 2 and 4 check the result against a reference built the same way, so a wrong expansion fails numerically rather than merely not crashing. bench_sage's tiebreak distribution mapped query directions with `arange(hq) % hk` and so disagreed with GQA's own grouping for any ratio above 1: two of four heads pointed at keys they never read. Matches test_mha_common's expansion now. Three guard tests asserted on a message that only the gfx950 ordering produces: - The LSE rejection reaches the format check only on gfx950; elsewhere the architecture gate fires first. Pick the expected message from the architecture. - test_mha_v4_rejects_unusable_seqlens_k wants the dtype and shape errors, which sit behind the same gate. Restrict it to gfx950; its sorted-sparse assertion is already covered arch-independently in test_mha_v4_sparse.py. - Two sparse guard tests built a mask from mha_v4_kv_tile() outside the raises block, so the manifest read, not the guard, would fail on an architecture with no sparse row. One takes a literal shape, which is enough because that guard runs before the LUT is built; the other reaches _block_mask_to_lut, which reads the manifest itself, so it takes the arch gate its siblings already have. Signed-off-by: jcaraban <jcaraban@amd.com>
One kernel backend per PR: these two files carry the only Triton-tagged changes here and neither half needs the other, so they go to a parallel PR off main rather than a stacked one. The wrapper change only moves an import inside a function, so the API this branch consumes is unchanged. Signed-off-by: jcaraban <jcaraban@amd.com>
github-actions
Bot
changed the title
Contributor
Author
mha_v4_kv_tile() is @functools.cache'd and derives from get_gfx(), and nothing clears it. A test that monkeypatches get_gfx and touches the tile leaves the cache holding the other architecture's value: monkeypatch restores the attribute, not the cache, so every later test in the process reads the wrong KV tile -- including test_mha_v4_sparse.py, which pytest runs next in the same process. That is silently wrong mask geometry rather than an error. Not live today: the only get_gfx patch stays on the dense path and never warms the tile. Signed-off-by: jcaraban <jcaraban@amd.com>
The ASM Q-scale gather used to be an unguarded global load covering a whole 256-row query tile, so a partial final tile addressed rows past the sequence and the producers reserved zeroed slack to keep those reads mapped. The gfx950 objects this branch ships promote q_descale to a buffer descriptor, so num_records clamps the overshoot in hardware and the slack buys nothing. Q scales now allocate at their logical extent like K's, which also makes the three MXFP*_q producers agree with their own register_fake instead of over-allocating relative to it, and removes the second copy of 256 that lived across the language boundary. Nothing numeric can police this. The overshot rows are past seqlen_q, so their values never reach the output: poisoning the slack with a large E8M0 exponent leaves every recipe bit-identical, on the unbounded objects as much as the bounded ones. The symptom of a regression is a fault at an unrelated later synchronization, and only when the over-read crosses an unmapped page. So the guard is now the object itself -- a disassembly check that no DISPATCHABLE fmha_v4 object fetches a scale with global_load_dword/ubyte. It is driven off the manifest, not the directory, because fwd_hd128_f4f4*.co still ship unguarded and are commented out there; re-enabling a row without rebuilding it fails the test. Verified to fail against main's objects (12 offenders) and pass against these. K keeps its padding, its slack constants and its C++ check: its two-tile producer lead is unchanged, as are the FP4 K/V gathers. 366 passed, 1 skipped (was 379: -15 Q backing-storage cases, +2 arch arms), and all five MX-Q recipes stay finite and fault-free at seqlen_q 1/100/255/257/513. Signed-off-by: jcaraban <jcaraban@amd.com>
f4f4 carries the same F4 Q/K/V triple as mxfp4, and mxfp4 took over its v_pack=1 FP6-P V order, so the two resolve to one row and mxfp4 wins it. Its manifest entries have been commented out for a while; this drops the dead entries and the two objects they named, which were the only ones still shipping an unguarded scale gather. Every manifest row now resolves to a shipped object and no object is unreferenced. Two tests said f4f4 but dispatched MXFP4/MXFP4, which is mxfp4 for as long as f4f4 has no manifest row; relabelled so the name matches the object under test. The gather invariant keeps reading the manifest rather than the directory -- that now guards a commented row being re-enabled without a rebuild, instead of the f4f4 objects specifically. The Triton-side f4f4 helpers (sage_quant_v_f4f4, _f4f4_v_kperm, the mha4_f4f4 bench spec) are left alone here; they belong to the bench_sage change set. 366 passed, 1 skipped. Signed-off-by: jcaraban <jcaraban@amd.com>
The FP6-P V order is consumed by mxfp4 and f6f4 now; f4f4 is deprecated and its row no longer exists. Last place in the tree that still implied otherwise. Signed-off-by: jcaraban <jcaraban@amd.com>
Contributor
This branch has not been deployed
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.
Motivation
#5798 added the LSE epilogue to the dense gfx950 recipes; sorted-sparse still raised. This completes it, so block-sparse MHA v4 can also run under ring / context parallelism, which merges per-rank partials via LSE.
Re-validating all 22 gfx950 hd128 objects for the LSE epilogue surfaced a set of correctness bugs that are fixed here, and the same sweep carries a throughput refresh.
Technical Details
Kernels
-inf.Las an integer, so-0.0might have turn into-Infbf16fp8might have produced non-finite output on long sequences, leading to black images/videos.2^-118and could overflowP = 2^32, from the previous 2^8.8Host
lseonfmha_v4_fwd_sparse, through the cu, header, pybind, custom op and both fakes.populate_lse_kernargfor dense and sparse, so_LSE_CAPABLE_QVis the only gate.return_lseraised under GQA — on the dense path too, not just the newly enabled sparse one..md._plain_block_scaledocstring: K's gather does overshoot the logical rows, andnum_recordsis what clamps it.Minor
bench_sage.pyaddspeakedandtiebreakdistributions; tiebreak's head mapping follows the GQA grouping.moe_wgradautotune was repeatedly invoked.Test Plan
torch.logsumexpover the selected blocks.black --checkandruffon PR-touched Python files.Test Result
test_mha_v4.py+test_mha_v4_sparse.pyon MI350/gfx950.sk, sparse 70/70 including empty rows, discontinuous LUTs and GQA; the bounded V-scale deploy ran 120/120 and 350/350, the i8fp8 refresh 264/264 and 770/770.torch.logsumexp: bf16 0.0009 nats, bf16fp8 0.03, fp8 0.05.black --checkandruffon the six PR-touched Python files: passed.Previous PRs: #5798, #5335, #5005, #4967, #4627