Repository navigation

[Triton/Gluon] [ASM] [HIP] Mha v4: adds bf16 sparse, LSE support, KV varlen, fixes, etc - #5798

Merged

jcaraban merged 66 commits into

mainfrom

mha_v4_fixes_sparse_etc

Sep 28, 2026

Merged

[Triton/Gluon] [ASM] [HIP] Mha v4: adds bf16 sparse, LSE support, KV varlen, fixes, etc#5798
jcaraban merged 66 commits into
mainfrom
mha_v4_fixes_sparse_etc

Commits on Sep 21, 2026
- Pin MX V partial-tile results against attention
  jcaraban
  committed
  d067061
- Cover block-sparse accuracy as the KV-tile count grows
  jcaraban
  committed
  4656866
- Refresh gfx950 dense MHA v4 objects for ragged seqlen_k
  jcaraban
  committed
  d2747d6
- Refresh the gfx950 FP8, MXFP8 and I8FP8 MHA v4 objects
  jcaraban
  committed
  f3b75fd
- Refresh the gfx950 F6F8 MHA v4 objects
  jcaraban
  committed
  7d973b9
- Accept per-batch key lengths in dense MHA v4
  jcaraban
  committed
  4ac7e7e
- Align sorted-sparse F8F6 MHA v4 with the dense FP6-P V packing
  jcaraban
  committed
  55869d9
- Update gfx950 F8F6 MHA v4 objects to the block-schedule build
  jcaraban
  committed
  e75d261
- Rebuild gfx950 F8F6 MHA v4 objects with a compacted register allocation
  jcaraban
  committed
  e8dff04
- Rebuild gfx950 F6F8 and F8F6 MHA v4 objects with a V-LDS race fix
  jcaraban
  committed
  ea62033
- Add a sorted-sparse gfx950 MXFP6 MHA v4 object and rebuild the dense one
  jcaraban
  committed
  bab0b5b
- Rebuild the gfx950 F6F4 and MXFP6 MHA v4 objects and move F6F4 sparse to FP6-P V
  jcaraban
  committed
  fa58266
- Fix all-MXFP4 sparse V packing and rebuild the gfx950 F4F4 objects
  jcaraban
  committed
  1fe9ae9
- Route all-MXFP4 attention to the canonical-V mxfp4 objects
  jcaraban
  committed
  cfb66a7
- Refresh gfx950 FMHA v4 f6f4 and mxfp8 objects
  jcaraban
  committed
  f6e1682
- Refresh gfx950 FMHA v4 mxfp6 objects
  jcaraban
  committed
  25a3dba
- Route gfx950 MXFP4 FMHA v4 through the FP6-P V packing
  jcaraban
  committed
  8ebcb8a
- Scope the MXFP4 FP6-P pack rule to the V format, and format
  jcaraban
  committed
  4c06e5f
- Retire MXFP4 Q/K with FP8 V
  jcaraban
  committed
  3497012
- Drop V dispatch branches unreachable after the FP8-V retirement
  jcaraban
  committed
  27103cf
- Refresh the MHA v4 doc for dense/sparse recipe parity
  jcaraban
  committed
  70f2529
- Remove the deprecated mxfp8 alias and the canonical MXFP4 V producer
  jcaraban
  committed
  b105a94
- Add coherent-attention input distributions to the sage benchmark
  jcaraban
  committed
  39e5a71
- Subtract K's token mean before quantizing in MHA v4
  jcaraban
  committed
  21511cd
- Add a K common-mode regression test for the quantized recipes
  jcaraban
  committed
  0b50b11
- Fuse the K mean subtraction into the MX quantizer kernels
  jcaraban
  committed
  417ecae
- Split the block-sparse cases out of test_mha_v4.py
  jcaraban
  committed
  adc319e
Commits on Sep 22, 2026
- Gate K smoothing on the measured shared component
  jcaraban
  committed
  5442b59
- Clamp the per-channel V amax so empty heads do not quantize to NaN
  jcaraban
  committed
  0a15f4d
- Raise the K-smoothing gate to the magnitude where it consistently pays
  jcaraban
  committed
  f5cf9bc
- Calibrate the benchmark's default distribution to diffusion attention
  jcaraban
  committed
  fbaf64b
- Emit the per-tensor amax from the hd128 rotate kernel
  jcaraban
  committed
  42b9b49
- Rebuild the MHA v4 fp8 hd128 dense kernel
  jcaraban
  committed
  e4d088c
- Add BF16 sparse MHA v4 support
  jcaraban
  committed
  6543f65
- Standardize gfx950 sparse LUT capacity
  jcaraban
  committed
  f2e399e
- Rebuild the MHA v4 fp8 hd128 sparse kernel
  jcaraban
  committed
  7e2c32a
- Rebuild the MHA v4 fp8 hd128 kernels after a scheduling simplification
  jcaraban
  committed
  61cf2fc
- Rebuild the MHA v4 mxfp8 hd128 kernels
  jcaraban
  committed
  b2ddc20
- Rebuild the MHA v4 i8fp8 hd128 kernels
  jcaraban
  committed
  b4ebc52
- Rebuild the MHA v4 i8fp8 hd128 kernels
  jcaraban
  committed
  4a8f7b5
- Rebuild the MHA v4 mxfp8 hd128 kernels
  jcaraban
  committed
  f0fc5e9
Commits on Sep 23, 2026
- Enable BF16 sparse benchmark modes
  jcaraban
  committed
  30abd62
- Rebuild the MHA v4 f8f6, f6f8, mxfp6, f6f4 and mxfp4 hd128 kernels
  jcaraban
  committed
  13a1d9d
- Rebuild the MHA v4 mxfp4 hd128 kernels
  jcaraban
  committed
  21a51a3
- bench_sage: drop the unsupported fp6_p argument on the f4f4 path
  jcaraban
  committed
  9149ef3
- mha v4: dense log-sum-exp output
  jcaraban
  committed
  cc06c3c
Commits on Sep 24, 2026
- Reject per-batch key lengths on rows that ignore them
  jcaraban
  committed
  4a908c7
- Add the K-smoothing shift back into the exported LSE
  jcaraban
  committed
  c600fae
- Require seqlens_k to live on Q's device
  jcaraban
  committed
  0345bb0
- Fix the BF16 Q/K, FP8 V exported LSE under a frozen-max rollback
  jcaraban
  committed
  1fb6b55
- Bound a per-batch key length inside the MHA v4 BF16 objects
  jcaraban
  committed
  cb9a297
- Require the K mean to share the input's GPU
  jcaraban
  committed
  e6d53d5
- Format test_mha_v4.py with black
  jcaraban
  committed
  31aa821
- Require the LSE buffer to live on Q's device, and document it
  jcaraban
  committed
  009c350
- Satisfy the pinned style checks
  jcaraban
  committed
  4d4a0ed
- Gate LSE to the architecture where it has been measured
  jcaraban
  committed
  15b591f
- Allocate the LSE output through one helper
  jcaraban
  committed
  9a6127e
- Trim the oversized comments in the sage benchmark
  jcaraban
  committed
  da09e19
- Fold the duplicated head and sparse-LSE validation into helpers
  jcaraban
  committed
  ece80ed
- Give the two custom-op recipes a shared exit
  jcaraban
  committed
  686e053
- Name the unpadded block-scale allocation
  jcaraban
  committed
  a8e6123
- Merge branch 'main' into mha_v4_fixes_sparse_etc
  jcaraban
  authored
  e22ba20
Commits on Sep 25, 2026
- Pin that the exported LSE bias does not move between key chunks
  jcaraban
  committed
  b9a01f0
- Bring both MHA v4 op tests up to the op-test standard
  jcaraban
  committed
  642c744
Commits on Sep 27, 2026
- Minor LUT scheduling tweaks across the gfx950 MHA v4 kernels
  jcaraban
  committed
  247a9d3
- Refresh the gfx950 MHA v4 bf16 and bf16fp8 objects
  jcaraban
  committed
  dadc3a7

Commits

Commits on Sep 21, 2026

Commits on Sep 22, 2026

Commits on Sep 23, 2026

Commits on Sep 24, 2026

Commits on Sep 25, 2026

Commits on Sep 27, 2026