Repository navigation
[ASM] [HIP] Mha v4: LSE on sorted-sparse, kernel race fixes, gfx950 refresh - #6841
Merged
Merged
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
The rotate/quantize kernels process 8 rows per block and guard their stores and amax with row < m, but subtracted the K mean before that guard. The input load is bounded by m_oob through opus::make_gmem; the mean was a raw pointer with no bound, so any K row count (batch x keys x local heads) that is not a multiple of 8 read a mean row past the end of the buffer. It faults when the mean sits flush against the end of its allocation and silently reads a neighbour otherwise. All four kernels shared it: hd128, mxfp8, mxfp6 and mxfp4. Two field crashes trace to this: rotate_activation_hd128 under sequence-parallel FP8 attention, and the coalesced MXFP6 K quantizer on wan2.2. Both reproduce on the pre-fix kernel and are clean after. Results are unchanged. Out-of-range lanes are discarded by the row < m guards, and the rotation's inter-thread butterfly has max group_size = dim/vec_size, exactly the threads per row, so it never spans rows. Verified by 32 bitwise identical digests (4 kernels x 4 ragged lengths x 2 head counts) against the pre-fix build. The regression test covers all four kernels and runs its probe out-of-process, because the failure is a GPU memory fault that aborts the interpreter rather than failing an assertion. All 12 parametrizations were confirmed to fail against the reverted kernel.
jcaraban
force-pushed
the
mha_v4_lse_sparse_etc
branch
from
October 26, 2026 09:00
42811a9 to
983fc4b
Compare
CI runs op_tests/test_*.py as `python3 <file>`, not under pytest (see .github/scripts/aiter_test.sh), so only main() executes. Two gaps followed from that: the benchmark sweep called checkAllclose with default arguments, which logs a mismatch and returns a ratio rather than raising, so a completely wrong kernel still exited 0; and the invariants that no numeric sweep can reach were reachable only from pytest. Follow test_mxfp8fp4gemm.py: each recipe gets a `result` column from a shared _verdict(err), a catastrophic AssertionError is caught into that column rather than aborting, and main() raises once after printing the table so every recipe is reported instead of only the first bad one. catastrophic_check fires only on NaN/Inf or delta > ref_max/2, which no working recipe approaches, so the loose rtol/atol that lets the quantized recipes report their accuracy is unchanged. Lift the seqlens_k object gate and the K-mean bound probe into plain check_* functions called from main(), in the style of test_mha_flydsl.py's check_routing_declines. The pytest cases stay as thin wrappers, so nothing is dropped. Both failure paths verified: deploying main's fp8 object, which lacks the 0x1C0 load, makes main() exit 1 and name the object; forcing one recipe's verdict to "failed" makes it exit 1 after the table with that recipe named.
jcaraban
force-pushed
the
mha_v4_lse_sparse_etc
branch
from
October 26, 2026 12:24
9753a03 to
08207b4
Compare
CI runs op_tests/test_*.py as `python3 <file>` (.github/scripts/aiter_test.sh), so the 376 pytest cases in these two files never executed there. Invoke pytest from __main__ and propagate its status, the way test_gated_delta_rule and test_iq2r_glm53 already do. It sits in the __main__ guard rather than main() because main() returns early on an unsupported gfx, and the cases carry their own skip marks. Script mode now reports 215 passed and 160 passed/1 skipped, matching what pytest collects, with the benchmark tables unchanged. A seeded failure exits 1 naming the case.
… point Running pytest from __main__ made these two files print a verbose session into every CI log, so remove the dependency instead. op_tests/mha_v4_check_harness.py provides parametrize, skipif, param, expect_raises and patched_attr, and CheckRegistry.run walks the module's test_* functions. The checks keep their decorators, so the tests themselves barely move. The benchmark sweep already reports correctness per recipe, the way test_mxfp8fp4gemm does. The 96 checks are what that cannot express: rejection guards, schema assertions, work-table ordering, compile-without-graph-breaks, and out-of-process fault probes. Values bind by argument name, because stacked parametrize decorators do not stack in signature order. Cases a mark disables are counted as skipped rather than dropped. Case count was pinned against pytest's collection before pytest was removed: 376 cases, and the per-function counts still match exactly. 215 ran/0 skipped and 160 ran/1 skipped, both exit 0. A seeded failure exits 1 naming the case, and reverting the K-mean bound fix exits 1 with 12 failed rotate-kernel cases.
Running these checks from the script entry point executed them in CI for the first time, and three things only that run could surface. The harness reported str(exc), which is empty for a bare assert, so a CI failure named the check and nothing else. Log the traceback instead: it only prints on failure, and it is the whole report. Four checks had no architecture guard because nothing had ever run them off gfx950. The two seqlens_k ones and the LSE chunk-bias one hit _check_varlen_capable/_check_lse_capable, and test_mha_v4_rejects_unsupported_contracts matched the arch message rather than the one it meant to assert. All four are gfx950-only. The mxfp8 'extreme' case filled with bf16.max, which the butterfly grows 128x before its final divide; bf16 shares f32's exponent range, so it rotated to inf/NaN and the two paths then agreed only on an encoding neither promises. Use the largest input that stays finite. The Hadamard check demanded bitwise equality against an fp32 matmul. That sum is only exact when the 128 inputs share an exponent range, so it was pinning the kernel to whichever reduction order the installed torch picked. Allow one output ULP; a 2-ULP perturbation is still rejected. 215 ran/0 skipped and 160 ran/1 skipped locally, both exit 0.
… LUT probe cwd Read the MI300X and the second MI35X logs rather than inferring from a gfx950 box, which is what the first pass did; it guessed four checks and three of them were not the ones failing. On gfx942 the dense file aborted outright: the MX quantizers AITER_CHECK on gfx950, and that check aborts instead of raising, so the process died mid-run and never printed a summary. test_mha_v4_lse_survives_k_common_mode allowed gfx942 even though LSE itself is gfx950-only, and the two off-device launch rejections and the LSE gate check had no guard at all. All are gfx950-only now. Both runners failed the same two sparse cases, which had nothing to do with architecture: the malformed-LUT probe runs with cwd=AITER_ROOT_DIR and imports op_tests, and under a wheel install AITER_ROOT_DIR is the installed package, which has no op_tests. It passed locally only because an editable install points that at the checkout. Use the directory holding this file. gfx950 is unchanged at 215 ran/0 skipped and 160 ran/1 skipped.
The twin of lse_survives_k_common_mode, and the likeliest cause of the gfx942 core dump: its guard admits gfx942, but the mxfp8 and mxfp4 arms reach quantizers that AITER_CHECK on gfx950, and that check aborts rather than raising, so the file died mid-run with no summary. The FP8 arm is genuine gfx942 coverage, so skip the two MX cases instead of the whole check. Simulating a gfx942 import now leaves 26 device-touching cases, none of which reach an MX quantizer. gfx950 is unchanged at 215 ran/0 skipped.
valarLip
previously approved these changes
Oct 27, 2026
… scale values The k-mean-bound probe runs each rotate kernel in a subprocess and asserts it exits 0. On gfx942 the three MX recipes exit non-zero on "MX quantization requires gfx950", so all nine MX cases failed; hd128 passed and still runs. This is also why the earlier core dump is gone: the remaining MX launches are in a child process, so the abort no longer takes the test process with it. The two rotated-fp8 failures are left failing on purpose. Only the scale compares unequal there -- the rotated values and the quantized codes are bitwise identical -- on two of three shapes, gfx942 only. I could not reproduce it from gfx950 and have no defensible mechanism for it, so the assert now carries both values rather than guessing at a tolerance.
Merged
2 of 3 tasks
valarLip
approved these changes
Oct 28, 2026
johannes-graner
added a commit
that referenced
this pull request
Oct 28, 2026
Upstream mha_v4_quant no longer pads Q block scales; the ASM gather clamps past the sequence end via the buffer descriptor.
johannes-graner
added a commit
that referenced
this pull request
Oct 28, 2026
The MHA V4 V-LDS publication race is fixed upstream in #6841 (merged via main); strict ws4 validation passes.
azaidy
pushed a commit
that referenced
this pull request
Oct 28, 2026
…and a deferred MoE import (#6869) * Defer the MoE import in the sage quant wrappers 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. (cherry picked from commit ed3e932) * bench_sage: add `peaked` and `tiebreak` input distributions 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. (cherry picked from commit f5f9483) * bench_sage: satisfy black on the tiebreak v_axis call 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> (cherry picked from commit 1e701ae) * bench_sage: map tiebreak heads with the GQA ratio The tiebreak distribution pointed each query head at a pair of near-tied keys using `arange(hq) % hk`, which only agrees with GQA's own grouping when the ratio is 1. test_mha_common expands KV heads as `(h g)`, so query head j reads KV head j // ratio: at hq=4, hk=2 the mapping wanted [0, 0, 1, 1] and got [0, 1, 0, 1], leaving two heads pointed at keys they never read and so measuring ordinary diffuse attention rather than a tie. Found by the Copilot review on #6841. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: drop the f4f4 quant path f4f4 is deprecated: it carries the same F4 Q/K/V triple as mxfp4, which took over its v_pack=1 FP6-P V order, so the two are one row. Nothing here was distinguishing them -- _production_quantize_f4f4 was byte-identical to _production_quantize_mxfp4, and the mha4_f4f4 spec was identical to mha4_mxfp4's -- so `--kernel all` has been benchmarking that row twice, and `--kernel mha4_f4f4` only still resolved because it dispatched the mxfp4 object. sage_quant_v_f4f4 and sage_quant_v_fp4_colmajor_kernel go with it. The wrapper had no callers at all and the kernel had no caller but the wrapper; both packed per-channel FP4 V, which no live row consumes now that the F4 rows use the block-scaled MXFP4 packer. _f4f4_v_kperm is NOT dead and is renamed rather than removed: pack_v_mxfp4_colmajor_raw uses the same kv-column permutation, so it is a property of the col-major FP4 V layout, not of f4f4. 342 passed, 1 skipped; mha4_mxfp4 cosine 0.9915 against the reference, which is the path that exercises the renamed helper. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: cover the peaked and tiebreak generators The GQA head mapping had no regression test. Asserts the tie where it is observable -- the top-2 keys of the real score matrix -- rather than re-deriving the index arithmetic, so a wrong grouping fails numerically. Reverting to `arange(hq) % hk` fails ratios 2 and 4 and passes ratio 1, which is exactly the blind spot: at ratio 1 the two mappings are identical. Also covers the `sk >= 2*len(_TIEBREAK_GAPS)` guard and that `peaked` averages fewer V rows than `diffusion`, compared against diffusion rather than an absolute bound so the check tracks the generator and not a calibration constant. Not added to test_benchmark_smoke.py, which the review suggested: that harness is one parametrized case per script keyed by filename, and its only assertion is returncode == 0, which cannot see a wrong head mapping -- the bug this guards produced finite numbers and exit 0. Its shape also leaves --hk defaulted to --hq, so it never reaches a GQA ratio above 1. Its docstring defers correctness to the sibling file. Runs on CPU in ~3s; bench_sage.py itself is unchanged. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: move the distribution tests into op_tests/triton_tests They were under op_benchmarks, which select_triton_tests.py treats as BENCH and split_tests.sh never globs, so CI collected nothing: the triton shard lists only op_tests/triton_tests. The file is now index 130 of the 149 the shard picks up. Only the bench loader's relative path changes with the move. Reverting the GQA head mapping to the pre-fix `% hk` still fails gqa_ratio 2 and 4 and passes 1, where the two formulas agree, so the coverage survived the move. * bench_sage: cover the tiebreak reference bound The only correctness line this PR adds to the bench is tiebreak's entry in the per-distribution reference bound, and nothing collected exercised it. Assert it stays looser than the bounds the ordinary distributions use, which is the property it exists for: tiebreak leaves BF16 by construction, so a neighbouring bound would report breakage on the generator's own purpose. Deleting the entry fails this check and nothing else. * bench_sage: report quality instead of gating on it The bench already took its checker from the test tree. Move the last piece of correctness policy there too: fp8_max_diff_percentage now lives beside check_attention_outputs in test_fav3_sage and takes the distribution name, so no tolerance or comparison logic is defined in the bench. Split the two jobs --compare-to-ref had fused. The bench measures a perf/quality tradeoff, so the reference comparison now always reports MAE/MaxE/Cosine and never fails the run. --strict restores the old pass/fail behaviour for local bisects. The flag is --report-quality, with --compare-to-ref kept as an alias. Verified both branches directly: without --strict a non-finite and a grossly wrong output both return; with --strict both raise. * bench_sage: fold the tolerance check into test_fav3_sage and drop the generator tests test_bench_sage_distributions.py was the only file in triton_tests that imported a bench, through an importlib path hack and private module state. Its remaining subject, fp8_max_diff_percentage, now lives in test_fav3_sage, so the check moves there and the file goes. The generator tests go with it rather than following: bench_sage imports test_fav3_sage at module scope, so testing the generators from there would close a loop. The input generators are bench code, and this stops the bench from being a tested API surface. Removing the tiebreak entry still fails the moved check and nothing else. --------- Signed-off-by: jcaraban <jcaraban@amd.com>
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
…efresh (#6841) * Pin the sorted-sparse host guards and fix the K-scale docstring 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. * Defer the MoE import in the sage quant wrappers 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. * mha_v4: export LSE on the sorted-sparse path 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. * mha_v4: refresh the gfx950 quantized fmha code objects 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. * mha_v4: qualify the exported LSE as the quantized-score logsumexp 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. * bench_sage: add `peaked` and `tiebreak` input distributions 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. * mha_v4: refresh the gfx950 bf16 and bf16fp8 fmha code objects 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. * mha_v4: deploy bounded gfx950 V-scale objects; dense 120/120, sparse 350/350 * mha_v4: refresh the gfx950 f8f6/mxfp6/f6f4/mxfp4 code objects 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. * mha_v4: refresh the gfx950 i8fp8 code objects 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. * fmha v4 gfx950: rebuild hd128 mxfp4 and mxfp6 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. * fmha v4 gfx950: rebuild hd128 i8fp8 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. * fmha_v4: refresh gfx950 mxfp8 objects (dense + sparse) 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. * fmha gfx950: update f6f4 kernels GPU7 paired +2.09% (20/20); dense 24/24, sparse 70/70. * fmha gfx950: update i8fp8 kernels GPU7 paired dense +1.10%, sparse +0.27%; dense 24/24, sparse 70/70. * fmha_v4 hd128 mxfp4: fix K0-prefetch lgkm race, rebuild the VGPR map 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. * fmha_v4 gfx950 hd128: epilogue L==0 guard, and raise the freeze cap on 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. * fmha_v4 gfx950 hd128: refresh the mxfp4 and f6f4 code objects 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. * fmha_v4 gfx950 hd128: refresh the mxfp6 code objects 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. * fmha_v4 gfx950 hd128: refresh the mxfp4 code objects 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. * fmha_v4 gfx950 hd128: refresh the mxfp8 code objects +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. * bench_sage: satisfy black on the tiebreak v_axis call 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> * mha_v4: expand the K mean across GQA heads, and pin the guards by architecture 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> * Move the bench_sage and sage-wrapper changes to their own PR 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> * test_mha_v4: clear the kv-tile cache between tests 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> * mha_v4: drop the Q-scale gather padding, now that the gather is bounded 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> * mha_v4: retire the f4f4 row 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> * mha_v4: stop naming f4f4 in the MXFP4 V packer 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> * MHA v4: per-batch seqlens_k on the 8 quantized gfx950 hd128 rows The dense objects now read kernarg 0x1C0 and clamp the KV loop to each batch's key length; a null pointer keeps one object serving dense and varlen. The override is emitted after all base-address setup: the MX scale slices are strided by kv_seq_len, so shrinking it earlier relocates each batch's scale base and corrupts every batch after the first. _launch_mxfp4_coalesced and _launch_mxfp6 now forward seqlens_k; without it the MX recipes attended over the full padded key length. Replaces the now-unreachable rejection test with a manifest-driven invariant that disassembles every row _VARLEN_CAPABLE_QV claims and fails if it never loads 0x1C0. The gate is keyed on (q_format, v_format) alone, so a new object for an already-claimed pair would otherwise inherit the claim unrebuilt. Validated per recipe against a physically truncated K/V control, GQA 4x, and LSE (<=0.23 nats). Dense output is bitwise unchanged and the sparse objects are byte-identical. Dense ALIGN_PAD re-swept: fp8 +0.19%, f8f6 +0.32%, f6f8 +0.32%, mxfp6 +0.60%, f6f4 +0.63%, mxfp4 +0.44%; i8fp8 and mxfp8 unchanged. * mha_v4: refresh the doc and make the MXFP4 V rejection accurate seqlens_k now covers every dense gfx950 row, not just BF16 Q/K; the doc still claimed the old scope. Also record that _LSE_CAPABLE_QV and _VARLEN_CAPABLE_QV are claims about the deployed objects keyed on (q_format, v_format) alone, which is why the suite disassembles every claimed row instead of trusting the table. The MXFP4-Q/FP8-V row was dropped from the manifest, so "raw preprocessing is not implemented yet" read as pending work. State the requirement instead, and keep it accurate for an MXFP6 V, which never had a row either. * mha_v4_quant: bound the K-mean read in the rotate kernels The rotate/quantize kernels process 8 rows per block and guard their stores and amax with row < m, but subtracted the K mean before that guard. The input load is bounded by m_oob through opus::make_gmem; the mean was a raw pointer with no bound, so any K row count (batch x keys x local heads) that is not a multiple of 8 read a mean row past the end of the buffer. It faults when the mean sits flush against the end of its allocation and silently reads a neighbour otherwise. All four kernels shared it: hd128, mxfp8, mxfp6 and mxfp4. Two field crashes trace to this: rotate_activation_hd128 under sequence-parallel FP8 attention, and the coalesced MXFP6 K quantizer on wan2.2. Both reproduce on the pre-fix kernel and are clean after. Results are unchanged. Out-of-range lanes are discarded by the row < m guards, and the rotation's inter-thread butterfly has max group_size = dim/vec_size, exactly the threads per row, so it never spans rows. Verified by 32 bitwise identical digests (4 kernels x 4 ragged lengths x 2 head counts) against the pre-fix build. The regression test covers all four kernels and runs its probe out-of-process, because the failure is a GPU memory fault that aborts the interpreter rather than failing an assertion. All 12 parametrizations were confirmed to fail against the reverted kernel. * mha_v4 op_tests: make the CI entry point able to fail CI runs op_tests/test_*.py as `python3 <file>`, not under pytest (see .github/scripts/aiter_test.sh), so only main() executes. Two gaps followed from that: the benchmark sweep called checkAllclose with default arguments, which logs a mismatch and returns a ratio rather than raising, so a completely wrong kernel still exited 0; and the invariants that no numeric sweep can reach were reachable only from pytest. Follow test_mxfp8fp4gemm.py: each recipe gets a `result` column from a shared _verdict(err), a catastrophic AssertionError is caught into that column rather than aborting, and main() raises once after printing the table so every recipe is reported instead of only the first bad one. catastrophic_check fires only on NaN/Inf or delta > ref_max/2, which no working recipe approaches, so the loose rtol/atol that lets the quantized recipes report their accuracy is unchanged. Lift the seqlens_k object gate and the K-mean bound probe into plain check_* functions called from main(), in the style of test_mha_flydsl.py's check_routing_declines. The pytest cases stay as thin wrappers, so nothing is dropped. Both failure paths verified: deploying main's fp8 object, which lacks the 0x1C0 load, makes main() exit 1 and name the object; forcing one recipe's verdict to "failed" makes it exit 1 after the table with that recipe named. * mha_v4 op_tests: run the pytest cases from the script entry point CI runs op_tests/test_*.py as `python3 <file>` (.github/scripts/aiter_test.sh), so the 376 pytest cases in these two files never executed there. Invoke pytest from __main__ and propagate its status, the way test_gated_delta_rule and test_iq2r_glm53 already do. It sits in the __main__ guard rather than main() because main() returns early on an unsupported gfx, and the cases carry their own skip marks. Script mode now reports 215 passed and 160 passed/1 skipped, matching what pytest collects, with the benchmark tables unchanged. A seeded failure exits 1 naming the case. * mha_v4 op_tests: drop pytest and run the checks from the script entry point Running pytest from __main__ made these two files print a verbose session into every CI log, so remove the dependency instead. op_tests/mha_v4_check_harness.py provides parametrize, skipif, param, expect_raises and patched_attr, and CheckRegistry.run walks the module's test_* functions. The checks keep their decorators, so the tests themselves barely move. The benchmark sweep already reports correctness per recipe, the way test_mxfp8fp4gemm does. The 96 checks are what that cannot express: rejection guards, schema assertions, work-table ordering, compile-without-graph-breaks, and out-of-process fault probes. Values bind by argument name, because stacked parametrize decorators do not stack in signature order. Cases a mark disables are counted as skipped rather than dropped. Case count was pinned against pytest's collection before pytest was removed: 376 cases, and the per-function counts still match exactly. 215 ran/0 skipped and 160 ran/1 skipped, both exit 0. A seeded failure exits 1 naming the case, and reverting the K-mean bound fix exits 1 with 12 failed rotate-kernel cases. * mha_v4 op_tests: fix what the newly-enabled checks found in CI Running these checks from the script entry point executed them in CI for the first time, and three things only that run could surface. The harness reported str(exc), which is empty for a bare assert, so a CI failure named the check and nothing else. Log the traceback instead: it only prints on failure, and it is the whole report. Four checks had no architecture guard because nothing had ever run them off gfx950. The two seqlens_k ones and the LSE chunk-bias one hit _check_varlen_capable/_check_lse_capable, and test_mha_v4_rejects_unsupported_contracts matched the arch message rather than the one it meant to assert. All four are gfx950-only. The mxfp8 'extreme' case filled with bf16.max, which the butterfly grows 128x before its final divide; bf16 shares f32's exponent range, so it rotated to inf/NaN and the two paths then agreed only on an encoding neither promises. Use the largest input that stays finite. The Hadamard check demanded bitwise equality against an fp32 matmul. That sum is only exact when the 128 inputs share an exponent range, so it was pinning the kernel to whichever reduction order the installed torch picked. Allow one output ULP; a 2-ULP perturbation is still rejected. 215 ran/0 skipped and 160 ran/1 skipped locally, both exit 0. * mha_v4 op_tests: gate the rest of the gfx950-only checks, and fix the LUT probe cwd Read the MI300X and the second MI35X logs rather than inferring from a gfx950 box, which is what the first pass did; it guessed four checks and three of them were not the ones failing. On gfx942 the dense file aborted outright: the MX quantizers AITER_CHECK on gfx950, and that check aborts instead of raising, so the process died mid-run and never printed a summary. test_mha_v4_lse_survives_k_common_mode allowed gfx942 even though LSE itself is gfx950-only, and the two off-device launch rejections and the LSE gate check had no guard at all. All are gfx950-only now. Both runners failed the same two sparse cases, which had nothing to do with architecture: the malformed-LUT probe runs with cwd=AITER_ROOT_DIR and imports op_tests, and under a wheel install AITER_ROOT_DIR is the installed package, which has no op_tests. It passed locally only because an editable install points that at the checkout. Use the directory holding this file. gfx950 is unchanged at 215 ran/0 skipped and 160 ran/1 skipped. * test_mha_v4: keep the MX recipes of the k-common-mode check off gfx942 The twin of lse_survives_k_common_mode, and the likeliest cause of the gfx942 core dump: its guard admits gfx942, but the mxfp8 and mxfp4 arms reach quantizers that AITER_CHECK on gfx950, and that check aborts rather than raising, so the file died mid-run with no summary. The FP8 arm is genuine gfx942 coverage, so skip the two MX cases instead of the whole check. Simulating a gfx942 import now leaves 26 device-touching cases, none of which reach an MX quantizer. gfx950 is unchanged at 215 ran/0 skipped. * test_mha_v4: gate the MX k-mean probes off gfx942, and report the fp8 scale values The k-mean-bound probe runs each rotate kernel in a subprocess and asserts it exits 0. On gfx942 the three MX recipes exit non-zero on "MX quantization requires gfx950", so all nine MX cases failed; hd128 passed and still runs. This is also why the earlier core dump is gone: the remaining MX launches are in a child process, so the abort no longer takes the test process with it. The two rotated-fp8 failures are left failing on purpose. Only the scale compares unequal there -- the rotated values and the quantized codes are bitwise identical -- on two of three shapes, gfx942 only. I could not reproduce it from gfx950 and have no defensible mechanism for it, so the assert now carries both values rather than guessing at a tolerance. --------- Signed-off-by: jcaraban <jcaraban@amd.com>
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
…and a deferred MoE import (#6869) * Defer the MoE import in the sage quant wrappers 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. (cherry picked from commit ed3e932) * bench_sage: add `peaked` and `tiebreak` input distributions 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. (cherry picked from commit f5f9483) * bench_sage: satisfy black on the tiebreak v_axis call 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> (cherry picked from commit 1e701ae) * bench_sage: map tiebreak heads with the GQA ratio The tiebreak distribution pointed each query head at a pair of near-tied keys using `arange(hq) % hk`, which only agrees with GQA's own grouping when the ratio is 1. test_mha_common expands KV heads as `(h g)`, so query head j reads KV head j // ratio: at hq=4, hk=2 the mapping wanted [0, 0, 1, 1] and got [0, 1, 0, 1], leaving two heads pointed at keys they never read and so measuring ordinary diffuse attention rather than a tie. Found by the Copilot review on #6841. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: drop the f4f4 quant path f4f4 is deprecated: it carries the same F4 Q/K/V triple as mxfp4, which took over its v_pack=1 FP6-P V order, so the two are one row. Nothing here was distinguishing them -- _production_quantize_f4f4 was byte-identical to _production_quantize_mxfp4, and the mha4_f4f4 spec was identical to mha4_mxfp4's -- so `--kernel all` has been benchmarking that row twice, and `--kernel mha4_f4f4` only still resolved because it dispatched the mxfp4 object. sage_quant_v_f4f4 and sage_quant_v_fp4_colmajor_kernel go with it. The wrapper had no callers at all and the kernel had no caller but the wrapper; both packed per-channel FP4 V, which no live row consumes now that the F4 rows use the block-scaled MXFP4 packer. _f4f4_v_kperm is NOT dead and is renamed rather than removed: pack_v_mxfp4_colmajor_raw uses the same kv-column permutation, so it is a property of the col-major FP4 V layout, not of f4f4. 342 passed, 1 skipped; mha4_mxfp4 cosine 0.9915 against the reference, which is the path that exercises the renamed helper. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: cover the peaked and tiebreak generators The GQA head mapping had no regression test. Asserts the tie where it is observable -- the top-2 keys of the real score matrix -- rather than re-deriving the index arithmetic, so a wrong grouping fails numerically. Reverting to `arange(hq) % hk` fails ratios 2 and 4 and passes ratio 1, which is exactly the blind spot: at ratio 1 the two mappings are identical. Also covers the `sk >= 2*len(_TIEBREAK_GAPS)` guard and that `peaked` averages fewer V rows than `diffusion`, compared against diffusion rather than an absolute bound so the check tracks the generator and not a calibration constant. Not added to test_benchmark_smoke.py, which the review suggested: that harness is one parametrized case per script keyed by filename, and its only assertion is returncode == 0, which cannot see a wrong head mapping -- the bug this guards produced finite numbers and exit 0. Its shape also leaves --hk defaulted to --hq, so it never reaches a GQA ratio above 1. Its docstring defers correctness to the sibling file. Runs on CPU in ~3s; bench_sage.py itself is unchanged. Signed-off-by: jcaraban <jcaraban@amd.com> * bench_sage: move the distribution tests into op_tests/triton_tests They were under op_benchmarks, which select_triton_tests.py treats as BENCH and split_tests.sh never globs, so CI collected nothing: the triton shard lists only op_tests/triton_tests. The file is now index 130 of the 149 the shard picks up. Only the bench loader's relative path changes with the move. Reverting the GQA head mapping to the pre-fix `% hk` still fails gqa_ratio 2 and 4 and passes 1, where the two formulas agree, so the coverage survived the move. * bench_sage: cover the tiebreak reference bound The only correctness line this PR adds to the bench is tiebreak's entry in the per-distribution reference bound, and nothing collected exercised it. Assert it stays looser than the bounds the ordinary distributions use, which is the property it exists for: tiebreak leaves BF16 by construction, so a neighbouring bound would report breakage on the generator's own purpose. Deleting the entry fails this check and nothing else. * bench_sage: report quality instead of gating on it The bench already took its checker from the test tree. Move the last piece of correctness policy there too: fp8_max_diff_percentage now lives beside check_attention_outputs in test_fav3_sage and takes the distribution name, so no tolerance or comparison logic is defined in the bench. Split the two jobs --compare-to-ref had fused. The bench measures a perf/quality tradeoff, so the reference comparison now always reports MAE/MaxE/Cosine and never fails the run. --strict restores the old pass/fail behaviour for local bisects. The flag is --report-quality, with --compare-to-ref kept as an alias. Verified both branches directly: without --strict a non-finite and a grossly wrong output both return; with --strict both raise. * bench_sage: fold the tolerance check into test_fav3_sage and drop the generator tests test_bench_sage_distributions.py was the only file in triton_tests that imported a bench, through an importlib path hack and private module state. Its remaining subject, fp8_max_diff_percentage, now lives in test_fav3_sage, so the check moves there and the file goes. The generator tests go with it rather than following: bench_sage imports test_fav3_sage at module scope, so testing the generators from there would close a loop. The input generators are bench code, and this stops the bench from being a tested API surface. Removing the tiebreak entry still fails the moved check and nothing else. --------- Signed-off-by: jcaraban <jcaraban@amd.com>
ksikiric
added a commit
to ksikiric/xDiT
that referenced
this pull request
Nov 2, 2026
AITER's MHA-v4 kernel race fix (ROCm/aiter#6841) removed the nondeterminism that FP6-P V consumers showed when a rank's sequence is an odd multiple of 32. FP6-P profiles now pad each rank's sequence to 32 tokens, like the other profiles.
johannes-graner
added a commit
that referenced
this pull request
Nov 3, 2026
…nce-parallel attention (#5815) * Extract shared P2P table and wave-uniform helpers from FlyDSL dispatch/combine Add FP6 RNE and native conversion helpers to quant_utils. * Add FlyDSL fused intranode all-to-all for Ulysses attention Seven Q/K/V transport codecs, stream-ordered launches, MoRI symmetric buffers. * Add attention A2A multi-GPU test and CI skip entry * Move A2A payload, ordering and addressing onto FlyDSL typed views and layouts Typed buffer views, ordered-memory completion flags, layout algebra for work and address mappings, named helpers for FP8/native FP6 conversion and wave-local LDS waits. Non-wrapping peer idx2crd outer modes use UNBOUNDED_OUTER_EXTENT. * Match reference E4M3 Q/K scale rounding and e4m3_pc zero-channel clamp Packed E4M3 Q/K descale uses fp32 reciprocal multiply like torch amax/dtype_max; e4m3_pc clamps channel amax to 1e-12 like quantize_v_fp8. Adds tests. * Return per-role packed payload and scale from submit_role submit_role returns PackedRoleResult(payload, scale) for Q, K and V in packed mode so callers get each role's buffers right after submission. * Fix attention A2A test after main removed quantize_v_mxfp4 and e4m3_pc buffer reuse Use the private HIP DEFAULT-layout MXFP4 V producer for packed references, and run the e4m3_pc zero/tiny stress submits after the attention comparison so they do not overwrite the parity buffers it consumes. * Address Copilot review on A2A test and docs Collect any exception into the cross-rank verdict so a rank-local error cannot strand peers in the Gloo all_reduce; add a packed mxfp8 Q/K recipe; correct the submit_role return docstring. * Add valid_kv_len to AttentionA2AIntraNodeOp for compact K/V output Emit canonical compact K/V layouts while retaining the padded Q length. Match valid rather than padded MHA tile counts and remove the need for xDiT to copy packed K/V buffers. Reject transports with sequence-wide scales and padding outside the final rank. Co-authored-by: Kristian Sikiric <kristian.sikiric@amd.com> * Add --valid-kv-len and --heads to attention A2A test Pass an explicit compact K/V length through the existing consumer recipes while retaining padded Q. Bound the reference score matrix so Wan-length cases can run with a selected head count. * Restore A2A scheduler barrier after per-tensor Q/K preload Without sched_barrier(0), LLVM serializes the Wan-shape amax loads: 20 loads in flight drop to 1 (vmcnt(0) per load). Each Q/K amax launch goes from 40.4 to 56.8 us, and the E4M3 trio is 10% slower. * a2a: kernel guards, exact INT8 E8M0 exponent, trimmed descriptor extents * a2a: host init/config checks and docstring clarifications * a2a: async reuse, trimming and software-FP6 test coverage * a2a: optional group arg for cross-rank config agreement check * a2a: add maintained RCCL+quantizer baseline benchmark * a2a: sync only the init stream before the MoRI barrier * Update aiter/ops/flydsl/kernels/attention_a2a_intranode_kernel.py Co-authored-by: Mohsen <mohsen.saffari@amd.com> * Harden fused A2A: mandatory group, PE/offset checks, stream rule, debug spin trap Require group, check MoRI PE/npes, verify symmetric offsets across ranks, reject null p2p pointers, reject cross-stream trios, add AITER_A2A_SPIN_TIMEOUT_ITERS debug trap, negative tests, and restore aiter_test.sh to origin/main plus the one skip line. * Treat v4-consumer output mismatch as known upstream MHA V4 race The gfx950 MHA V4 asm consumers publish the V LDS slot before their buffer_load ... lds copies complete, so mha_v4_packed output varies under concurrent same-GPU work. Transport bytes are now checked strictly before the consumer; consumer-output equality and the accuracy checks that read it are recorded as known-upstream unless AITER_A2A_STRICT_V4_CONSUMER=1. * Drop Q-scale tile padding in fused A2A after #6841 Upstream mha_v4_quant no longer pads Q block scales; the ASM gather clamps past the sequence end via the buffer descriptor. * Make v4-consumer mismatch a hard failure again The MHA V4 V-LDS publication race is fixed upstream in #6841 (merged via main); strict ws4 validation passes. * Restore strict transport-byte checks in fused A2A test Keeps transport faults distinguishable from consumer faults; the KNOWN-UPSTREAM soft-fail stays removed. * Run fused A2A multi-GPU test in CI MoRI is now installed in the multi-GPU job, so drop test_attention_a2a.py from skip_tests and run it at world sizes 2/4/8 with heads 8 and 40. * [CI] Skip bench_attention_a2a.py in multi-GPU suite The benchmark is a perf harness with no correctness checks and requires -w, so the default no-argument invocation failed the job. * a2a: clamp FP6-P V partial-amax token to valid_kv_len-1 The partial-max path zero-masked tokens past valid_kv_len while the payload path clamps to the final token first, so a split group published amax 1 instead of the canonical value. Mirror the payload clamp for mxfp6_p only; add a ws=2 valid_kv_len=33 regression recipe. * bench a2a: gate on gfx950 before spawning workers Workers initialize distributed and MoRI before the op rejects gfx942 packed MX recipes; skip in main using get_gfx_runtime(). --------- Co-authored-by: Kristian Sikiric <kristian.sikiric@amd.com> Co-authored-by: Mohsen <mohsen.saffari@amd.com> Co-authored-by: TennyWang1223 <Tenny.Wang@amd.com>
avjves
pushed a commit
to xdit-project/xDiT
that referenced
this pull request
Nov 4, 2026
* feat(attention): add quantized A2A contracts * feat(distributed): add AITER attention A2A runtime * feat(wan): integrate packed attention A2A * test(attention): cover A2A routing and document usage * perf(attention): skip redundant A2A setup * perf(attention): expose packed A2A launch to Inductor * perf(attention): trim A2A reuse and compaction overhead * perf(attention): tune Attention A2A launch geometry * test(attention): pin A2A codec contracts * fix(attention): use MORI default symmetric heap * perf(attention): consume compact A2A K/V output * fix(ci): format A2A changes and isolate optional deps * fix(attention): adapt A2A to mandatory participant group * fix(attention): keep A2A group last for Inductor * perf(attention): collapse waits for MX packed A2A * test(attention): cover A2A waits behaviorally * fix(attention): wait on every packed A2A output * perf(attention): use 256 A2A workgroups for MXFP6 Q/K * fix(attention): drop the FP6-P 64-token padding mitigation AITER's MHA-v4 kernel race fix (ROCm/aiter#6841) removed the nondeterminism that FP6-P V consumers showed when a rank's sequence is an odd multiple of 32. FP6-P profiles now pad each rank's sequence to 32 tokens, like the other profiles. * refactor(attention): isolate Attention A2A paths
2 of 10 tasks
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 LSE to the dense gfx950 kernels, but the sparse variants still missed it. This PR 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 potential bugs that are fixed here, and the same sweep carries a throughput refresh and per-batch
seqlens_kon the quantized rows.Technical Details
Kernels
-inf.seqlens_kfrom BF16 Q/K to all 8 quantized gfx950 hd128 rows.Las integer, so-0.0might have turned into-Infbf16fp8might have produced non-finite output on long sequences, leading to black images/videos.2^-118and might have overflownP = 2^32, from the previous 2^8.8Host
lseonfmha_v4_fwd_sparse, through the cu, header, pybind, custom op and both fakes; one sharedpopulate_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.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