[Triton/Gluon] [ASM] [HIP] MHA v4: support GQA, add gfx950 bf16, add gfx942 i8/fp8

aiter/ops/mha_v4.md

@@ -5,23 +5,25 @@
 
 ## Current Status
 
-Dense BF16-output MHA v4 is implemented and validated on gfx950. A gfx942 signed INT8/FP8 row is
-also preserved under v4.
+Dense BF16-output MHA v4 is implemented and validated on gfx950. Gfx942 native FP8/FP8 and signed
+INT8/FP8 rows are also available under v4.
 
-The public raw and packed APIs support six dense combinations:
+The public raw and packed APIs support eight dense combinations:
 
 | Q/K | V | Output |
 |---|---|---|
+| BF16 | BF16 | BF16 |
 | INT8 | FP8 | BF16 |
 | FP8 | FP8 | BF16 |
+| MXFP8 | FP8 | BF16 |
 | MXFP6 E2M3 | FP8 | BF16 |
 | MXFP4 E2M1 | FP8 | BF16 |
 | MXFP6 E2M3 | MXFP4 E2M1 | BF16 |
 | MXFP4 E2M1 | MXFP4 E2M1 | BF16 |
 
-Current scope is batched, dense, non-causal MHA with matching Q/KV head counts, BF16 raw inputs,
-head dimension 128, and BF16 output. It is inference-only: no backward, dropout, RNG state, LSE,
-GQA, varlen, or sparse metadata. Unsupported requests fail explicitly and never fall back to
+Current scope is batched, dense, non-causal MHA with supported grouped-query head ratios, BF16 raw
+inputs, head dimension 128, and BF16 output. It is inference-only: no backward, dropout, RNG state,
+LSE, varlen, or sparse metadata. Unsupported requests fail explicitly and never fall back to
 `aiter.ops.mha`.
 
 ## Stable Decisions And Ownership
@@ -47,7 +49,7 @@ GQA, varlen, or sparse metadata. Unsupported requests fail explicitly and never
 The current implementation is intentionally one module, `aiter/ops/mha_v4.py`; a speculative
 subpackage split is not part of the design. It exports:
 
-- `mha_v4` and `mha_v4_packed`;
+- `mha_v4`, `mha_v4_mxfp8`, and `mha_v4_packed`;
 - `AttentionFormat`, `AttentionScaleMode`, `native_fp8_format`, and `scale_modes_for_formats`;
 - canonical per-tensor, MX Q/K, and V quantizers;
 - `mxfp4_k_view`, `mxfp6_k_view`, and `mxfp4_v_view` for raw-buffer reconstruction;
@@ -67,7 +69,7 @@ migration, and distributed integration are complete. Callers can delegate quanti
 scaling, scale recipes, and packed views to MHA v4 while retaining separate Q/K/V custom ops for
 communication overlap.
 
-Validation includes eager accuracy for all six combinations, fullgraph eager/compiled parity,
+Validation includes eager accuracy for all eight combinations, fullgraph eager/compiled parity,
 finite outputs, allocator churn with downstream consumers, explicit code-object dispatch,
 unaligned and unequal sequence lengths, retained model captures, and balanced multi-GPU target-shape
 benchmarks. Focused coverage lives in `op_tests/test_mha_v4.py`.
@@ -76,8 +78,8 @@ Still deferred:
 
 - sparse ragged-LUT execution, VSA/Sparge compatibility, and ring/LSE support;
 - low-precision output with an explicit data/scale ABI;
-- approximate BF16 input under a distinct identity from v3 BF16;
-- GQA, causal, varlen, other head dimensions, and more Q/K/V/O combinations;
+- additional BF16 kernel variants with distinct manifest identities;
+- causal, varlen, other head dimensions, and more Q/K/V/O combinations;
 - broader gfx942, CDNA5, and RDNA manifest/code-object coverage.
 
 ## Current Dense Performance
@@ -123,6 +125,22 @@ Inputs are contiguous BF16 BSHD tensors. The requested formats select canonical
 preprocessing and an explicit ASM row; unsupported combinations fail. Q/K must currently match.
 Output is BF16, and a supplied `out` must match Q's shape/device. Q, K, and V preprocessing remain
 separate custom ops so distributed schedulers can overlap each with its input communication.
+The canonical FP8 Q/K recipe applies normalized hd128 Walsh-Hadamard rotation before per-tensor
+quantization on both gfx942 and gfx950; V uses unrotated per-tensor FP8 quantization.
+
+#### Grouped-Query Attention
+
+Both raw entrypoints (`mha_v4` and `mha_v4_mxfp8`) and `mha_v4_packed` accept GQA directly. Q uses
+shape `[batch, query_length, query_heads, 128]`; K and V use
+`[batch, key_value_length, kv_heads, 128]`. K and V must have the same head count, `query_heads`
+must be divisible by `kv_heads`, and the ratio `query_heads / kv_heads` must be one of
+`1, 2, 4, 8, 16`. Ratio 1 is ordinary multi-head attention. The kernel maps each contiguous group
+of query heads to one K/V head; callers must not expand K or V to `query_heads`. Output retains Q's
+batch, sequence, and head dimensions.
+
+For example, Q with 32 heads and K/V with 8 heads selects GQA ratio 4. Q and K still use the same
+number format and canonical quantization recipe; "Q/K formats must match" refers to their encoding,
+not their head counts. Ratios outside the supported power-of-two set fail explicitly.
 
 ### Packed Expert API
 
@@ -289,8 +307,9 @@ code_object
 
 Kernel cache identity is `(kernel_symbol, code_object)`, never the symbol alone.
 
-The approximate BF16 kernel uses a distinct symbol, code-object slot, and manifest row, for example
-`fwd_hd128_bf16_approx.co`. It must not overwrite or reuse generic `fwd_hd128_bf16.co` dispatch.
+BF16 dispatch uses the same explicit format and scale-mode key as other rows. Each architecture
+owns its manifest row and code object under `hsa/<arch>/fmha_v4_fwd/`; adding gfx942 BF16 support
+does not require a Python-side architecture branch.
 
 ## Sparse Contract
 
@@ -362,7 +381,7 @@ Fix offsets with the first implementing kernel; existing v1 binaries retain thei
 1. Add sparse manifest rows, ragged-LUT validation, and exact 256x128/128x128 execution paths.
 2. Add VSA/Sparge adapters over the shared sparse descriptor and packed executor.
 3. Add LSE under a stable output schema for ring attention.
-4. Add approximate BF16 under a distinct symbol and code object from generic v3 BF16.
+4. Add further BF16 variants only under distinct manifest identities.
 5. Add a versioned low-precision-output ABI once data/scale ownership is concrete.
 6. Expand architectures, head dimensions, sequence modes, and format combinations only through
     explicit manifest rows.

aiter/ops/mha_v4.py

@@ -17,7 +17,6 @@
 from aiter import dtypes
 from aiter.jit.core import compile_ops
 from aiter.jit.utils.chip_info import get_gfx
-from aiter.ops.quant import rotate_activation
 from aiter.ops.triton._triton_kernels.quant.sage_attention_quant import (
     mha_v4_per_tensor_amax_kernel,
     mha_v4_per_tensor_quant_kernel,
@@ -46,6 +45,11 @@ def mha_v4_q_multiplier(softmax_scale: float) -> float:
     return softmax_scale * MHA_V4_LOG2E
 
 
+@compile_ops("module_fmha_v4_fwd")
+def rotate_activation_hd128(out: Tensor, input: Tensor) -> None:
+    """Apply normalized Walsh-Hadamard rotation to contiguous hd128 rows."""
+
+
 @compile_ops("module_fmha_v4_fwd")
 def rotate_activation_mxfp8_quant(
     out: Tensor,
@@ -138,6 +142,7 @@ class AttentionScaleMode(IntEnum):
 _FP8_FORMATS = (AttentionFormat.FP8_E4M3, AttentionFormat.FP8_E4M3_FNUZ)
 _MX_FORMATS = (AttentionFormat.FP6_E2M3, AttentionFormat.FP4_E2M1)
 _PACKED_QK_WIDTH = {
+    AttentionFormat.BF16: 128,
     AttentionFormat.INT8: 128,
     AttentionFormat.FP8_E4M3: 128,
     AttentionFormat.FP8_E4M3_FNUZ: 128,
@@ -176,6 +181,10 @@ def _validate_format_contract(
         )
     if q_format != k_format:
         raise ValueError("MHA v4 currently requires matching Q and K formats")
+    if q_format == AttentionFormat.BF16:
+        if v_format != AttentionFormat.BF16:
+            raise ValueError("BF16 Q/K currently requires BF16 V")
+        return
     if q_format not in _PACKED_QK_WIDTH:
         raise ValueError(f"unsupported Q/K format: {q_format!r}")
     if v_format not in (
@@ -200,6 +209,12 @@ def scale_modes_for_formats(
 ) -> tuple[AttentionScaleMode, AttentionScaleMode, AttentionScaleMode]:
     """Return the canonical Q, K, and V scale modes for a format recipe."""
     _validate_format_contract(q_format, k_format, v_format)
+    if q_format == AttentionFormat.BF16:
+        return (
+            AttentionScaleMode.NONE,
+            AttentionScaleMode.NONE,
+            AttentionScaleMode.NONE,
+        )
     if q_format == AttentionFormat.INT8 or q_format in _FP8_FORMATS:
         v_scale_mode = (
             AttentionScaleMode.F32_PER_TENSOR
@@ -380,10 +395,14 @@ def mha_v4_packed(
         raise ValueError("Q, K, and V must have the same batch size")
     if k.shape[1] != v.shape[1] or k.shape[2] != v.shape[2]:
         raise ValueError("K and V must have matching sequence and head dimensions")
-    if query_heads != k.shape[2]:
-        raise ValueError(
-            "MHA v4 initially supports MHA only; Q and KV heads must match"
-        )
+    kv_heads = k.shape[2]
+    if kv_heads == 0:
+        raise ValueError("MHA v4 requires non-empty KV heads")
+    if query_heads % kv_heads != 0:
+        raise ValueError("MHA v4 requires query heads to be divisible by KV heads")
+    gqa_ratio = query_heads // kv_heads
+    if gqa_ratio > 16 or gqa_ratio & (gqa_ratio - 1):
+        raise ValueError("MHA v4 supports power-of-two GQA ratios up to 16")
     if not q.is_cuda or not k.is_cuda or not v.is_cuda:
         raise ValueError("MHA v4 expects GPU tensors")
     if q.device != k.device or q.device != v.device:
@@ -508,7 +527,7 @@ def quantize_fp8_rotated(input: Tensor) -> tuple[Tensor, Tensor]:
     if input.shape[-1] != 128 or not input.is_contiguous():
         raise ValueError("rotated FP8 quantization requires contiguous hd128 input")
     rotated = torch.empty_like(input)
-    rotate_activation(rotated, input)
+    rotate_activation_hd128(rotated, input)
     return quantize_fp8(rotated)
 
 
@@ -902,6 +921,73 @@ def _launch_mxfp6_fake(
     del out
 
 
+def _validate_mha_v4_raw_inputs(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    out: Optional[Tensor],  # noqa: UP045
+    operation: str,
+) -> Tensor:
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(f"{operation} expects BSHD Q, K, and V tensors")
+    if (
+        q.dtype != torch.bfloat16
+        or k.dtype != torch.bfloat16
+        or v.dtype != torch.bfloat16
+    ):
+        raise ValueError(f"{operation} currently expects BF16 Q, K, and V inputs")
+    if q.shape[-1] != 128 or k.shape[-1] != 128 or v.shape[-1] != 128:
+        raise ValueError(f"{operation} currently supports head dimension 128 only")
+    if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
+        raise ValueError(f"{operation} currently requires contiguous BSHD inputs")
+    if out is None:
+        return torch.empty_like(q, dtype=torch.bfloat16)
+    if out.shape != q.shape or out.dtype != torch.bfloat16 or out.device != q.device:
+        raise ValueError("out must match Q's shape/device and have BF16 dtype")
+    return out
+
+
+def mha_v4_mxfp8(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    softmax_scale: Optional[float] = None,  # noqa: UP045
+    out: Optional[Tensor] = None,  # noqa: UP045
+    return_lse: bool = False,
+) -> Tensor:
+    """Quantize BF16 BSHD Q/K to MXFP8 and V to per-tensor FP8.
+
+    K and V may have fewer heads than Q for GQA. The Q-to-KV head ratio must
+    be a power of two no greater than 16; output retains Q's head count.
+    """
+    if return_lse:
+        raise NotImplementedError("MHA v4 kernels do not produce LSE yet")
+    out = _validate_mha_v4_raw_inputs(q, k, v, out, "mha_v4_mxfp8")
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** -0.5
+
+    fp8_format = native_fp8_format()
+    q_quantized, q_descale = quantize_mxfp8_q(q, mha_v4_q_multiplier(softmax_scale))
+    k_quantized, k_descale = quantize_mxfp8_k(k)
+    v_quantized, v_descale = quantize_fp8(v)
+    return mha_v4_packed(
+        q_quantized,
+        k_quantized,
+        v_quantized,
+        q_descale,
+        k_descale,
+        v_descale,
+        fp8_format,
+        fp8_format,
+        fp8_format,
+        AttentionScaleMode.E8M0_PER_1X32,
+        AttentionScaleMode.E8M0_PER_1X32,
+        AttentionScaleMode.F32_PER_TENSOR,
+        softmax_scale=softmax_scale,
+        out=out,
+    )
+
+
 def mha_v4(
     q: Tensor,
     k: Tensor,
@@ -917,29 +1003,33 @@ def mha_v4(
 
     Q and K formats must match. The selected Q/K/V recipe determines canonical
     quantizers, scale modes, and the packed ASM row; output is BF16 BSHD.
+    K and V may have fewer heads than Q for GQA. The Q-to-KV head ratio must
+    be a power of two no greater than 16; output retains Q's head count.
     """
     if return_lse:
         raise NotImplementedError("MHA v4 kernels do not produce LSE yet")
-    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
-        raise ValueError("mha_v4 expects BSHD Q, K, and V tensors")
-    if (
-        q.dtype != torch.bfloat16
-        or k.dtype != torch.bfloat16
-        or v.dtype != torch.bfloat16
-    ):
-        raise ValueError("mha_v4 currently expects BF16 Q, K, and V inputs")
-    if q.shape[-1] != 128 or k.shape[-1] != 128 or v.shape[-1] != 128:
-        raise ValueError("mha_v4 currently supports head dimension 128 only")
+    out = _validate_mha_v4_raw_inputs(q, k, v, out, "mha_v4")
     q_scale_mode, k_scale_mode, v_scale_mode = scale_modes_for_formats(
         q_format, k_format, v_format
     )
-    if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
-        raise ValueError("mha_v4 currently requires contiguous BSHD inputs")
-    if out is None:
-        out = torch.empty_like(q, dtype=torch.bfloat16)
-    elif out.shape != q.shape or out.dtype != torch.bfloat16 or out.device != q.device:
-        raise ValueError("out must match Q's shape/device and have BF16 dtype")
 
+    if q_format == AttentionFormat.BF16:
+        return mha_v4_packed(
+            q,
+            k,
+            v,
+            q,
+            k,
+            v,
+            q_format,
+            k_format,
+            v_format,
+            q_scale_mode,
+            k_scale_mode,
+            v_scale_mode,
+            softmax_scale=softmax_scale,
+            out=out,
+        )
     if q_format == AttentionFormat.INT8 and _is_fp8_format(v_format):
         q_quantized, q_descale = quantize_int8(q)
         k_quantized, k_descale = quantize_int8(k)
@@ -948,9 +1038,8 @@ def mha_v4(
         q_format,
         AttentionFormat.MXFP6,
     ):
-        quantize_qk = quantize_fp8 if _is_fp8_format(v_format) else quantize_fp8_rotated
-        q_quantized, q_descale = quantize_qk(q)
-        k_quantized, k_descale = quantize_qk(k)
+        q_quantized, q_descale = quantize_fp8_rotated(q)
+        k_quantized, k_descale = quantize_fp8_rotated(k)
         if _is_fp8_format(v_format):
             v_quantized, v_descale = quantize_fp8(v)
         else:

aiter/ops/triton/quant/mxfp6_fmha_pack.py

@@ -162,9 +162,11 @@ def quantize_fp6_v_clean_triton(
         v_fp8.stride(1),
         v_fp8.stride(2),
         v_fp8.stride(3),
+        sk,
         h_kv,
         nT,
         n_blocks,
+        CLAMP_TAIL=False,
         FIXED_E8M0=fixed_e8m0,
         SEPARATE_OUTPUT=False,
         BLOCK_N=BLOCK_N,
@@ -178,8 +180,8 @@ def quantize_fp6_v_data_scale_triton(
     """Pack F8F6 V directly into its separate data and scale ABI buffers."""
     assert _HAVE_TRITON, "triton/torch unavailable"
     b, sk, h_kv, d = v_fp8.shape
-    assert d == 128 and tile == 128 and sk % tile == 0, (d, sk, tile)
-    nT = sk // tile
+    assert d == 128 and tile == 128, (d, sk, tile)
+    nT = (sk + tile - 1) // tile
     n_blocks = b * h_kv * nT * 128 * 4
     data = torch.empty(
         b * h_kv * nT * 12288 + 256, dtype=torch.uint8, device=v_fp8.device
@@ -197,9 +199,11 @@ def quantize_fp6_v_data_scale_triton(
         v_fp8.stride(1),
         v_fp8.stride(2),
         v_fp8.stride(3),
+        sk,
         h_kv,
         nT,
         n_blocks,
+        CLAMP_TAIL=sk % tile != 0,
         FIXED_E8M0=fixed_e8m0,
         SEPARATE_OUTPUT=True,
         BLOCK_N=BLOCK_N,
@@ -284,9 +288,11 @@ def _pack_v_fp6_kernel(
         stride_vs,
         stride_vh,
         stride_vd,
+        sk,
         h_kv,
         nT,
         n_blocks,  # total 32-kv MX blocks
+        CLAMP_TAIL: tl.constexpr,
         FIXED_E8M0: tl.constexpr,
         SEPARATE_OUTPUT: tl.constexpr,
         BLOCK_N: tl.constexpr,
@@ -310,6 +316,8 @@ def _pack_v_fp6_kernel(
         f = tl.arange(0, 32)
         kt = tl.load(kvtab_ptr + L[:, None] * 32 + f[None, :])  # [BN,32]
         kv = (t * 128 + k * 64)[:, None] + kt  # [BN,32] kv-in-tile
+        if CLAMP_TAIL:
+            kv = tl.minimum(kv, sk - 1)
         voff = (
             bb[:, None] * stride_vb
             + kv * stride_vs
@@ -867,10 +875,6 @@ def pack_fp6_v_data_scale_views(
     assert _HAVE_TRITON, "triton/torch unavailable"
     b, sk, h_kv, d = v.shape
     n_tiles = (sk + tile - 1) // tile
-    sk_pad = n_tiles * tile
-    if sk_pad != sk:
-        tail = v[:, sk - 1 : sk].expand(b, sk_pad - sk, h_kv, d)
-        v = torch.cat([v, tail], dim=1)
 
     data_flat, scale_flat = quantize_fp6_v_data_scale_triton(
         v, tile=tile, fixed_e8m0=fixed_e8m0