[Triton/Gluon] [ASM] [HIP] MHA v4: fixes, refactor, new kernel, perf tweaks - #5335
Conversation
Dense MXFP4 kernels issue an overlapping final K-scale dword read, while the FP6-P F4F4 path prefetches one V-scale tile beyond the logical sequence. Exact-sized scale allocations can therefore fault when their ends coincide with an unmapped boundary. Allocate zeroed backing slack for both scale tensors while preserving their logical shapes and strides. Add layout tests that verify the padding size and contents.
Deploy the validated gfx950 MXFP4 attention binary and dispatch dense MXFP4 Q/K/V through the canonical V packing layout. Keep the existing sparse MXFP4 Q/K with FP8 V recipe intact, and align benchmark accounting, tests, and documentation with both paths.
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
# Conflicts: # aiter/ops/mha_v4.py # csrc/pybind/mha_v4_fwd_pybind.cu
There was a problem hiding this comment.
🟡 Changes recommended
Critical packed-buffer validation gaps permit speculative out-of-bounds reads, with additional compatibility, benchmark, and test issues unresolved.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Expands gfx950 MHA v4 with mixed-precision recipes, optimized V packing, and centralized quantization and dispatch.
Changes:
- Adds BF16/FP8 and FP6-P packed kernels.
- Extracts quantization and recipe planning.
- Updates dispatch manifests, tests, benchmarks, and documentation.
File summaries
| File | Description |
|---|---|
op_tests/test_mha_v4.py |
Expands recipe, packing, sparse, and compile tests. |
op_tests/op_benchmarks/triton/bench_sage.py |
Refactors recipe benchmarks and payload accounting. |
hsa/gfx950/fmha_v4_fwd/fmha_v4_fwd.csv |
Adds packing-aware dispatch entries. |
csrc/pybind/mha_v4_quant_pybind.cu |
Exposes new V quantizers. |
csrc/pybind/mha_v4_fwd_pybind.cu |
Adds V-packing launcher arguments. |
csrc/py_itfs_cu/asm_mha_v4_fwd.cu |
Implements recipe validation and packing-aware dispatch. |
csrc/kernels/mha_v4_quant.cu |
Implements FP6-P V-packing kernels. |
csrc/include/torch/mha_v4_fwd.h |
Extends forward interfaces for packing. |
csrc/include/mha_v4_quant.h |
Declares new quantizers. |
aiter/ops/mha_v4.py |
Centralizes recipe planning and launches. |
aiter/ops/mha_v4.md |
Updates MHA v4 usage and contracts. |
aiter/ops/mha_v4_quant.py |
Extracts quantization and layout producers. |
Review details
Suppressed comments (3)
csrc/py_itfs_cu/asm_mha_v4_fwd.cu:688
- This validates only the logical V-scale shape, while
quantize_v_mxfp4_fp6_pnow allocates and zeroes an additional 512-byte scale tile specifically for speculative lookahead reads. A packed caller can supply an exact-size tensor, pass validation, and leave that read unmapped or nonzero. Require the extra zeroed backing storage when the selected dense V pack needs it, or copy into guarded storage.
TORCH_CHECK(recipe.v_scale_mode == scale_mode_id(AttentionScaleMode::E8M0Per1x32) &&
v_descale.scalar_type() == at::ScalarType::Byte,
"MX V descale must use uint8 E8M0 per-1x32 scales");
TORCH_CHECK(v_descale.sizes() ==
torch::IntArrayRef({shapes.batch, shapes.nhead_k, tiles * 512}),
"MX V descale must have shape [batch, key_heads, tiles * 512]");
op_tests/op_benchmarks/triton/bench_sage.py:1378
- In
--e2emode,mha4_f4f4takes this raw path, but_resolve_raw_recipe(MXFP4, MXFP4, MXFP4)deliberately selectsAttentionPack.DEFAULT; the manifest therefore dispatchesfwd_hd128_mxfp4.co, not the requested pack-1fwd_hd128_f4f4.co. The reported F4F4 end-to-end timing is consequently for a different kernel. Keep F4F4 on the packed helper so quantization is timed while preservingV_FOR_FP6_P.
if args.e2e:
if args.hadamard_rotate:
return lambda: launch_mha_v4(
q_bshd,
k_bshd,
v_bshd,
AttentionFormat.MXFP4,
AttentionFormat.MXFP4,
v_format,
softmax_scale=softmax_scale,
)
return lambda: _kernel_mxfp4(*_quantize_mxfp4())
op_tests/test_mha_v4.py:2352
- This adds a benchmark CLI and
@benchmarkworkload to the unit-test module. Repository convention requires benchmark entry points underop_tests/op_benchmarks/triton/asbench_<op>.py; keeping it here also makes test collection import benchmark-only dependencies and mixes two execution modes. Move this benchmark into the benchmark tree and leavetest_mha_v4.pyfocused on pytest coverage (CodingGuidelineID 1000002).
@benchmark()
def benchmark_mha_v4(batch, sequence_q, sequence_k, heads, dtype):
"""Benchmark the public dense BF16 MHA v4 path against a Torch reference."""
- Files reviewed: 12/27 changed files
- Comments generated: 5
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
The MHA v4 ASM kernels gather Q and K block scales with unguarded global loads that address every row of a tile, plus F4F4's two-tile producer lookahead for K. A partial final tile therefore reads past the logical sequence, and the resulting out-of-bounds access faults the GPU whenever the memory after the tensor is unmapped. The fault surfaces asynchronously, so it appeared as an unrelated abort inside a later Triton launch and made the suite crash depending on test selection and allocator state. Size the scale backing storage from the tile the kernel actually gathers and zero the slack, keeping each tensor's logical shape and strides.
The overallocated kv_block_indices tail stays uninitialized and an empty row's lut_start points into it, so kernels that read LUT[lut_start] before testing lut_count have to clamp the offsets they derive from it.
Both paths launch the same deterministic code object, so the 2% tolerance only hid a dispatch, scale or lifetime regression that every other recipe catches exactly. The recipe is bitwise stable across repeated eager and compiled launches, so assert equality like the rest.
The producer zeroes a 256-byte trailing buffer for the ASM's speculative reads, but the layout test compared only data_size bytes, so a regression there would have left those reads nondeterministic while the test still passed.
The refactor folded this recipe into mha_v4 but dropped the previously public entry point outright, breaking existing importers even though the recipe is unchanged. Forward to mha_v4 with a DeprecationWarning for the compatibility window.
The packed API only checked descale shapes, so an external caller could pass an exact-size Q or MXFP4 K scale, satisfy every check, and still fault the GPU: the ASM gathers scales across the whole tile it is running, plus MXFP4 K's two-tile producer lead, and reads past the logical rows. Reject unbacked scales with a clear error instead. The aiter.ops.mha_v4_quant producers already reserve the zeroed slack, so internal callers are unaffected.
|
copilot review addressed @Boss2002n , do you have further comments? edit: I solved other potential issues in the same line of what copilot pointed out |
The MXFP4 Q/K rows gather V scales two 512-byte tiles ahead of the tile they are running, and that lead does not shrink at the end of the sequence, so the last tiles address scale bytes past the final one at every shape. The canonical producer reserved no slack and the FP6-P producer reserved one tile, so both walked off the tensor. Under the caching allocator that silently reads a neighbouring tensor; at a segment boundary it faults the GPU at an unrelated later synchronization. Measured on gfx950 with the caching allocator disabled and the scale placed at the end of its own mapping: 1023 trailing mapped bytes still fault, 1024 do not, constant in both sequence length and head count. The canonical packing lived in a Triton packer that allocated its own scale, so the requirement could only be expressed by editing a shared quantization module whose other callers do not want it. The HIP producer already emits this layout and differs only by a token pairing, so make that pairing a template parameter and route both packings through it. Buffer ownership now sits in one helper in the mha_v4 domain, which also retires the nested-custom-op TODO, and the launcher refuses an externally supplied scale that is not backed. Verified the HIP canonical output is bit-identical to the retired packer across sequence lengths, and kept that packer as a second test oracle.
The fake allocated on a hard-coded "cuda", so under torch.compile it reported device 0 whatever device the LUT actually lived on, and the traced metadata disagreed with the real op on any other device.
Every stride and extent in the kernarg block is a uint32, but the launcher assigned int64 byte strides into them unchecked. A tensor whose batch stride exceeds 4 GiB truncated silently and the kernel then addressed the wrong rows, so the failure was a wrong answer rather than a fault. BF16 hd128 reaches this once sequence * heads passes 16.7M, which needs batch > 1 to matter but is not out of reach for long context. Route the assignments through a checked conversion so the launcher names the offending field and refuses the shape instead. No test: triggering it needs a multi-gigabyte allocation, so it was verified by hand rather than in the suite.
if we are adding a new kernel - id say pls open a new PR, makes it so much more easier to review it - right now its just a very dense PR and very hard to follow atleast for the triton changes |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
you are so rude that you don't even look the code |
Sorry @Boss2002n, I agree the PR is dense. There are many kernels under this MHA v4 group and typically one PR needs to address them all together. I also need to pack bugfixes along, as we don't want segfaults in production. I'm actively migrating from the Triton logic to self-contained HIP. The history behind is that originally we had Triton Sage v1 (i8fp8) and Save v3? (mxfp4) and the ASM kernels have been replacing them. The ASM variants still used the quantization / packing / rotation logic from Triton, but I'm cutting that dependency for better isolated HIP kernels. Right now this PR only touches aiter/ops/triton/attention/utils.py with a TODO (for the next PR) and op_tests/op_benchmarks/triton/bench_sage.py which tests the perf of both Triton and ASM sage kernels so unfortunately always shows large diffs |
Oh no, That makes perfect sense :) |
Motivation
Expand MHA v4 on gfx950 with additional low-precision dense attention recipes, improved quantization and packing paths, and unified recipe dispatch.
Technical Details
Kernels
block-sparsesorted-dispatchcompatibleHost
mha_v4_quant.py.Minor
Test Plan
Test Result
Previous PRs: #4627, #4967, #5005