Repository navigation
Conversation
… ahead, speculative softmax
The gfx950 fwd_hd128_mxfp6 kernel's pre-shifted Region-B K-scale gather reads each token's E8M0 dword at a +1-byte offset, which runs 1 byte past the scale tensor on the final (token,head,batch). Re-home `scale` into a buffer with trailing slack so that read stays mapped (the pad byte is never consumed by the MFMA op_sel).
Rebuilt from asm/fmha_v3_fwd/mi350/fwd_hd128_bf16_woven.py (diffusion repo bb95bcf0). Verified: builds, deploys, MAE 0 vs known-good, ~1417 TFLOPS.
quantize_fp6_k_lds_order_triton now always emits the 17408B-per-tile K layout (16384B chunk-major fp6 data + a 1024B lane-major K-scale tail: Region A unshifted + Region B pre-shifted), returning a stride-136 k_view. The fwd_hd128_mxfp6 kernel reads the scale straight from the K-buffer tail (coalesced buffer_load lds:1), so the separate K-scale global-load stream -- which was stalling -- is gone. Drops the AITER_MXFP6_PACK_KSCALE_TAIL env knob. Consumers use the returned k_view directly and auto-adapt to the new stride; there are no return_raw callers. Also updates the deployed fwd_hd128_mxfp6.co to the matching tail-consuming kernel (tail scale + bf16-MMA QK bias + address-gen hoist) so the packer and kernel stay coherent.
Replace the 32-iteration tl.static_range linear-search E2M3 encode in _pack_qk_fp6_kernel with a branchless round-half-even arithmetic encode. The E2M3 magnitude grid is a minifloat (2 exp bits, 3 mantissa bits, bias 1): normals (mag>=1) are 2^(exp2-1)*(1+m/8), subnormals (mag<1) are m/8. So the search collapses to (a) fp32 RNE-round-to-3-mantissa-bits for the normal range (add 0x7FFFF + kept-LSB, carry into exp) and (b) round-half-even of mag*8 for the subnormal range. Bit-identical to the search and the numpy packer (verified across scales). The pack kernel was compute-bound on the search (~360 GB/s, far below HBM); num_warps/BLOCK_N tuning was a wash. Result: fp6 pack 0.255ms -> 0.089ms; full mxfp6 e2e quant 0.496 -> 0.315ms, flipping mxfp6 e2e above fp8 (hq32 sq8192 1183->1462, hq5 sq65536 2537->2727). Kernel output bit-identical (cos 0.99745/0.99779). Drop the now-dead GRID constexpr param + call-site arg.
The arithmetic-encode _pack_qk_fp6_kernel is bandwidth/latency-bound on its permuted per-block gather, so many small single-warp programs hide the load latency better than fewer wide ones. Measured ~17-19% faster than 128/default-4 across shapes (bit-exact): e.g. N=382500 (H5 S76500) 0.127->0.104ms, N=262144 (H32 S8192) 0.090->0.073ms.
…slot Adds the f6f4 attention variant (fp6 Q/K, per-channel fp4 E2M1 V read via ds_read_b64_tr_b4) as a first-class kernel that coexists with the mainline f6f8 (fp6-QK / fp8-V) build instead of sharing/overwriting one .co slot. - sage_attention_quant_wrappers: add production sage_quant_mxfp6(..., f6f4=, v_fp4_packer=, q_packer=, k_packer=). Computes only the selected V operand (fp8 via sage_quant_v_kernel for f6f8, per-channel fp4 via the caller-supplied packer for f6f4). Q/K default to the in-tree Triton fp6 packers, but q_packer/k_packer can override (e.g. the bench's AITER_MXFP6_PACK swap / numpy fallback). - bench_sage: use the production sage_quant_mxfp6; keep the host-module (hp) Q/K packers (_build_fp6_qk_packer / _build_fp6_k_coalesced_packer via _load_host_fp6_pack, honoring AITER_MXFP6_PACK / AITER_MXFP6_QK_TRITON) and pass them as q_packer/k_packer; f6f4 packs V via the host Triton packer. - dispatch (asm_mha_fwd.cu): pick "f6f4bf16" vs "mxfp6bf16" by V dtype (fp4 uint8 vs fp8); treat f6f4bf16 as mxfp-packed with bf16 output. - mha_fwd.cu: accept f6f4bf16 in the fp6/descale path + v3 support check, and key the kernel cache by (symbol | .co path) so same-symbol/different-.co variants don't collide. - codegen.py: include co_name in the generated config-map key for the same reason. - fmha_fwd.csv: add the f6f4bf16 row -> fwd_hd128_f6f4.co (the .co is deployed separately, not tracked in git). Co-authored-by: Cursor <cursoragent@cursor.com>
valarLip
reviewed
Aug 11, 2026
Cap persistent MXFP8 quant launches at 32768 programs to avoid millions of tiny row programs on long-context head batches. Write F8F6 V payloads and scales directly into their final ABI buffers, cache permutation tables, and share the branchless E2M3 RNE encoder across V and Q/K paths. Match the custom-op fake strides to the real padded backing storage and add byte-exact full/tail coverage.
Synchronize the checked-in F4F4, F6F4, F8F6, FP8, and MXFP4 code objects with their validated kernel-repository implementations. Each replacement preserves the symbol consumed by fmha_v4_fwd.csv; functional and performance validation was performed on the deployed objects before this refresh.
Register the dedicated hd128 MXFP8 code object and allow FP8 Q/K/V data with per-1x32 E8M0 Q/K scales plus a per-tensor FP8 V descale. Validate E8M0 scale tensors by scale mode rather than data format, add production rotate-and-quant preprocessing to bench_sage, and cover the packed scale recipe. The kernel folds Q/K microscaling into scaled MFMA while retaining the existing FP8 V path.
Add deterministic all-zero inputs for degenerate quantization smoke tests and expose block, tensor, and head V-scale modes for F8F6 experiments. The non-block modes normalize V before fixed-E8M0 FP6 packing and store the float descale in the scale image consumed by the kernel; the production block-scale path remains the default.
Remove obsolete MXFP6 host/direct packing paths and unused quantization helpers left from bring-up. Centralize repeated BSHD validation, layout mapping, and format checks while preserving the fused HIP preprocessing and packed-buffer ABI. Refresh the MHA v4 documentation to describe the current MXFP6 path and keep only comments that explain layout or launch constraints. Ruff, Black, quantization, MHA v4, SageAttention, and distributed integration checks pass.
Document the Python format and scale-mode ABI, public quantization entrypoints, packed launcher contract, and code-object kernarg layout. Group enum aliases separately and apply Ruff's canonical import ordering in the Sage benchmark.
valarLip
approved these changes
Aug 17, 2026
This was referenced Aug 17, 2026
This was referenced Aug 17, 2026
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
Introduce an extensible, inference-only MHA v4 entrypoint for mixed low-precision ASM attention kernels, without overwhelming the existing
aiter.ops.mhamodule. In the future we can consider merging into the common entrypoint, but these kernels are best kept independent as they evolve.Q: why v4? A: IMO this is the closest thing to FAv4 on cdna3/4, also doing 1) conditional-softmax-rescaling, 2) exponential-approx/emu, 3) optimized pipelines for valu/mfma/loads utilization. But I'm open for naming suggestions.
Replaces #3495 and #3186
coauthors: @juuso-oskari , @Chi-Chu319
Technical Details
aiter.ops.mha_v4as an extensible interface for future architectures, formats, output types, sparse attention, GQA, causal attention, head dimensions, etctorch.compileexecution and distributed overlap by separating Q/K/V preprocessing and safely encapsulating specialized packed layouts.Testing
Current mi355x kernel throughput for H=5 S=65536 D=128, excluding Q/K/V preprocessing:
Submission Checklist