diff --git a/aiter/ops/mha_v4.md b/aiter/ops/mha_v4.md
index 286bf43995b..c9f702ef7ff 100644
--- a/aiter/ops/mha_v4.md
+++ b/aiter/ops/mha_v4.md
@@ -8,23 +8,31 @@ dispatch metadata. Unsupported recipes fail instead of falling back to another a
 - Contiguous BF16 BSHD inputs with head dimension 128.
 - BF16 BSHD output.
 - Dense and sorted block-sparse inference.
+- GFX950 LDS-staged sparse recipes support up to 8,192 KV tiles per sequence.
 - Grouped-query ratios `1, 2, 4, 8, 16`.
-- No backward, dropout, RNG state, LSE, causal, or varlen support yet.
-
-Supported dense recipes:
-
-| Q/K | V |
-|---|---|
-| BF16 | BF16 |
-| BF16 | FP8 |
-| INT8 | FP8 |
-| MXFP8 | FP8 |
-| FP8 | FP8 |
-| FP8 | MXFP6 |
-| MXFP6 | FP8 |
-| MXFP6 | MXFP6 (dense only) |
-| MXFP6 | MXFP4 |
-| MXFP4 | MXFP4 |
+- Per-batch key lengths via `seqlens_k`, on the dense GFX950 `BF16 Q/K` rows only. Every other
+  recipe, architecture, and the sorted-sparse path reject it.
+- Log-sum-exp via `return_lse`, on the dense GFX950 rows. Sorted sparse rejects it, and so does
+  GFX942 until its exported value is measured.
+- No backward, dropout, RNG state, causal, or Q-side varlen support yet.
+
+Supported recipes. Every recipe is available in both dense and sorted-sparse mode with the same
+V packing and scale modes.
+
+| Q/K | V | Modes |
+|---|---|---|
+| BF16 | BF16 | dense, sparse |
+| BF16 | FP8 | dense, sparse |
+| INT8 | FP8 | dense, sparse |
+| MXFP8 | FP8 | dense, sparse |
+| FP8 | FP8 | dense, sparse |
+| FP8 | MXFP6 | dense, sparse |
+| MXFP6 | FP8 | dense, sparse |
+| MXFP6 | MXFP6 | dense, sparse |
+| MXFP6 | MXFP4 | dense, sparse |
+| MXFP4 | MXFP4 | dense, sparse |
+
+MXFP4 Q/K requires MXFP4 V. The FP8-V variant is retired.
 
 ## Ownership
 
@@ -60,11 +68,13 @@ Backend choice is private to `mha_v4_quant`; recipe selection does not branch on
 | Per-tensor INT8/FP8 | Triton |
 | Rotated FP8 and FP8 V | Triton |
 | Canonical MXFP6 V | Triton |
-| Canonical MXFP4 V | Triton |
 | MXFP8/MXFP6/MXFP4 Q and K | HIP `module_mha_v4_quant` |
 | FP6-P MXFP6 V | HIP `module_mha_v4_quant` |
 | FP6-P MXFP4 V | HIP `module_mha_v4_quant` |
 
+No manifest row consumes canonical MX V: every MXFP6-V and MXFP4-V row selects the FP6-P pack.
+`quantize_v_mxfp6` is retained only as the reference the FP6-P layout test permutes against.
+
 ## APIs
 
 Use `mha_v4` for BF16 inputs and canonical preprocessing:
@@ -115,10 +125,10 @@ calling `mha_v4_packed`.
 MXFP4 V uses E2M1 values with one E8M0 scale per `(channel, 32-token)` block. Each 128-token tile
 contributes 8,192 data bytes and 512 scale bytes. The data buffer includes 64 bytes of launch slack.
 
-`AttentionPack.DEFAULT` is the canonical V token order used by sparse kernels and FP8-P rows.
-`AttentionPack.V_FOR_FP6_P` selects the shared dense V token order for FP6-P and FP4-P consumers.
-Numeric format and consumer pairing are separate dispatch contracts even when the physical V
-layout is identical.
+`AttentionPack.DEFAULT` is the canonical V token order. `AttentionPack.V_FOR_FP6_P` selects the V
+token order that FP6-P and FP4-P consumers require, and a row's packing no longer depends on the
+mode: dense and sparse rows for the same recipe select the same pack. Numeric format and consumer
+pairing remain separate dispatch contracts even when the physical V layout is identical.
 
 Changing a custom op's output shape or packed layout requires a versioned custom-op name.
 
diff --git a/aiter/ops/mha_v4.py b/aiter/ops/mha_v4.py
index 4748fe83d65..ba808dd6391 100644
--- a/aiter/ops/mha_v4.py
+++ b/aiter/ops/mha_v4.py
@@ -12,9 +12,8 @@
 import csv
 import functools
 import os
-import warnings
 from enum import IntEnum
-from typing import NamedTuple, Optional
+from typing import NamedTuple, Optional, Union
 
 import torch
 from torch import Tensor
@@ -38,7 +37,6 @@
     quantize_mxfp8_k,
     quantize_mxfp8_q,
     quantize_v_fp8,
-    quantize_v_mxfp4,
     quantize_v_mxfp4_fp6_p,
     quantize_v_mxfp6,
     quantize_v_mxfp6_fp6_p,
@@ -54,7 +52,6 @@
     "AttentionScaleMode",
     "mha_v4",
     "mha_v4_kv_tile",
-    "mha_v4_mxfp8",
     "mha_v4_packed",
     "mha_v4_q_multiplier",
     "mha_v4_sparse_work_table",
@@ -72,7 +69,6 @@
     "quantize_mxfp8_k",
     "quantize_mxfp8_q",
     "quantize_v_fp8",
-    "quantize_v_mxfp4",
     "quantize_v_mxfp4_fp6_p",
     "quantize_v_mxfp6",
     "quantize_v_mxfp6_fp6_p",
@@ -193,6 +189,14 @@ class _RawRecipePlan(NamedTuple):
 # find_config(..., mode=1).
 _MHA_V4_SPARSE_MODE = 1
 
+# Shared-K component, relative to a typical K row, above which removing it improves quantization.
+# Real video models run to 0.67 (HunyuanVideo 1.5) and 0.76 (Wan) at the extreme, and on the one
+# captured layer above 0.7 the recipes disagree: FP8 lost 7% while MXFP4 and F8F6 gained. The win
+# is only consistent past ~0.85, where it is already 1.16x-1.31x, so the gate sits there and
+# leaves every layer either model actually produces untouched.
+_K_SMOOTH_MIN_COMMON = 0.85
+_K_SMOOTH_SAMPLE_ROWS = 2048
+
 
 def native_fp8_format() -> AttentionFormat:
     """Return the FP8 E4M3 encoding native to the active GPU architecture."""
@@ -405,10 +409,6 @@ def _resolve_raw_recipe(
     )
 
     if q_format == AttentionFormat.BF16:
-        if sparse:
-            raise NotImplementedError(
-                "sorted-sparse MHA v4 does not have a BF16 manifest row yet"
-            )
         kind = (
             _RawRecipeKind.BF16
             if v_format == AttentionFormat.BF16
@@ -421,23 +421,13 @@ def _resolve_raw_recipe(
     elif q_format in _FP8_FORMATS:
         kind = _RawRecipeKind.FP8
     elif q_format == AttentionFormat.MXFP4:
-        if v_format == AttentionFormat.MXFP6:
+        if v_format != AttentionFormat.MXFP4:
             raise NotImplementedError(
                 "raw preprocessing is not implemented yet for "
                 f"Q={q_format.name}, K={k_format.name}, V={v_format.name}"
             )
-        # Sparse still uses the legacy FP8-V row; update this mode split when its MXFP4-V row lands.
-        if not sparse and _is_fp8_format(v_format):
-            raise NotImplementedError(
-                "dense MXFP4 Q/K with FP8 V does not have a kernel row yet"
-            )
         kind = _RawRecipeKind.MXFP4
     elif q_format == AttentionFormat.MXFP6:
-        # Dense already has MXFP6 Q/K/V; remove this guard when the matching sparse row lands.
-        if sparse and v_format == AttentionFormat.MXFP6:
-            raise NotImplementedError(
-                "sorted-sparse MXFP6 Q/K/V does not have a kernel row yet"
-            )
         kind = _RawRecipeKind.MXFP6
     else:
         raise NotImplementedError(
@@ -445,15 +435,17 @@ def _resolve_raw_recipe(
             f"Q={q_format.name}, K={k_format.name}, V={v_format.name}"
         )
 
-    # Sparse FP6-P recipes still use canonical V packing; align them when those rows are updated.
-    uses_dense_p_pack = not sparse and (
+    # FP6-P rows need V repacked to match the FP6 P operand's K layout, or the kernel reads V rows
+    # in the wrong order. Dense and sparse agree on this per recipe; the mode never changes it.
+    uses_fp6_p_pack = (
         (kind == _RawRecipeKind.FP8 and v_format == AttentionFormat.MXFP6)
         or (
             kind == _RawRecipeKind.MXFP6
             and v_format in (AttentionFormat.MXFP6, AttentionFormat.MXFP4)
         )
+        or (kind == _RawRecipeKind.MXFP4 and v_format == AttentionFormat.MXFP4)
     )
-    v_pack = AttentionPack.V_FOR_FP6_P if uses_dense_p_pack else AttentionPack.DEFAULT
+    v_pack = AttentionPack.V_FOR_FP6_P if uses_fp6_p_pack else AttentionPack.DEFAULT
     return _RawRecipePlan(kind, scale_modes, v_pack)
 
 
@@ -542,11 +534,13 @@ def _fmha_v4_fwd_fake(
     k_scale_mode: int,
     v_scale_mode: int,
     softmax_scale: float,
+    seqlens_k: Optional[Tensor],  # noqa: UP045
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     del q, k, v, q_descale, k_descale, v_descale
     del q_format, k_format, v_format, v_pack
     del q_scale_mode, k_scale_mode, v_scale_mode, softmax_scale
-    del out
+    del out, seqlens_k, lse
 
 
 @compile_ops(
@@ -570,10 +564,12 @@ def _fmha_v4_fwd(
     k_scale_mode: int,
     v_scale_mode: int,
     softmax_scale: float,
+    seqlens_k: Optional[Tensor],  # noqa: UP045
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None: ...
 
 
-@torch.library.custom_op("aiter::mha_v4_fwd_launch", mutates_args=("out",))
+@torch.library.custom_op("aiter::mha_v4_fwd_launch", mutates_args=("out", "lse"))
 def _mha_v4_fwd_launch(
     q: Tensor,
     k: Tensor,
@@ -590,6 +586,8 @@ def _mha_v4_fwd_launch(
     k_scale_mode: int,
     v_scale_mode: int,
     softmax_scale: float,
+    seqlens_k: Optional[Tensor],  # noqa: UP045
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     _fmha_v4_fwd(
         q,
@@ -607,6 +605,8 @@ def _mha_v4_fwd_launch(
         k_scale_mode,
         v_scale_mode,
         softmax_scale,
+        seqlens_k,
+        lse,
     )
 
 
@@ -627,10 +627,13 @@ def _mha_v4_fwd_launch_fake(
     k_scale_mode: int,
     v_scale_mode: int,
     softmax_scale: float,
+    seqlens_k: Optional[Tensor],  # noqa: UP045
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     del q, k, v, q_descale, k_descale, v_descale, out
     del q_format, k_format, v_format, v_pack
     del q_scale_mode, k_scale_mode, v_scale_mode, softmax_scale
+    del seqlens_k, lse
 
 
 def _fmha_v4_fwd_sparse_fake(
@@ -757,6 +760,98 @@ def _mha_v4_fwd_sparse_launch_fake(
     del kv_block_indices, lut_start, lut_count
 
 
+# (q_format, v_format) rows whose code object implements the LSE epilogue. The kernels branch on
+# the s_lse kernarg, so an object without the epilogue would silently leave the buffer untouched
+# instead of failing. Keyed on the pair because one Q format spans several rows (FP8 Q reaches both
+# the fp8 and the f8f6 object).
+_LSE_CAPABLE_QV = frozenset(
+    {
+        (AttentionFormat.BF16, AttentionFormat.BF16),
+        (AttentionFormat.BF16, AttentionFormat.FP8_E4M3),
+        (AttentionFormat.BF16, AttentionFormat.FP8_E4M3_FNUZ),
+        (AttentionFormat.FP8_E4M3, AttentionFormat.FP8_E4M3),
+        (AttentionFormat.FP8_E4M3_FNUZ, AttentionFormat.FP8_E4M3_FNUZ),
+        (AttentionFormat.FP8_E4M3, AttentionFormat.FP6_E2M3),
+        (AttentionFormat.FP8_E4M3_FNUZ, AttentionFormat.FP6_E2M3),
+        (AttentionFormat.INT8, AttentionFormat.FP8_E4M3),
+        (AttentionFormat.INT8, AttentionFormat.FP8_E4M3_FNUZ),
+        (AttentionFormat.MXFP4, AttentionFormat.MXFP4),
+        (AttentionFormat.MXFP6, AttentionFormat.MXFP6),
+        (AttentionFormat.MXFP6, AttentionFormat.MXFP4),
+        (AttentionFormat.MXFP6, AttentionFormat.FP8_E4M3),
+        (AttentionFormat.MXFP6, AttentionFormat.FP8_E4M3_FNUZ),
+    }
+)
+
+
+def _check_lse_capable(
+    q_format: AttentionFormat, v_format: AttentionFormat, sparse: bool
+) -> None:
+    """Reject LSE where the exported value has not been measured."""
+    if sparse:
+        raise NotImplementedError(
+            "MHA v4 does not produce LSE on the sorted-sparse path yet"
+        )
+    arch = get_gfx()
+    if arch != "gfx950":
+        # The gfx942 objects do carry the epilogue, but they predate the frozen-max correction
+        # and their LSE has never been compared against torch.logsumexp. O never reads LSE, so
+        # passing every output test says nothing about it. Lift this once MI300 is measured.
+        raise NotImplementedError(f"MHA v4 LSE is not validated on {arch} yet")
+    if (q_format, v_format) not in _LSE_CAPABLE_QV:
+        raise NotImplementedError(
+            f"MHA v4 LSE is not implemented for Q={q_format.name} "
+            f"V={v_format.name} yet"
+        )
+
+
+# (q_format, v_format) rows whose code object consumes the seqlens_k kernarg. Every dense source
+# carries the load, but only these objects were rebuilt with it; the rest would silently attend
+# over the full padded key length, so they are rejected rather than left to return a wrong answer.
+_VARLEN_CAPABLE_QV = frozenset(
+    {
+        (AttentionFormat.BF16, AttentionFormat.BF16),
+        (AttentionFormat.BF16, AttentionFormat.FP8_E4M3),
+        (AttentionFormat.BF16, AttentionFormat.FP8_E4M3_FNUZ),
+    }
+)
+
+
+def _check_varlen_capable(q_format: AttentionFormat, v_format: AttentionFormat) -> None:
+    """Reject per-batch key lengths on rows whose code object ignores them."""
+    arch = get_gfx()
+    if arch != "gfx950" or (q_format, v_format) not in _VARLEN_CAPABLE_QV:
+        raise NotImplementedError(
+            f"MHA v4 per-batch key lengths are not implemented for "
+            f"Q={q_format.name} V={v_format.name} on {arch} yet"
+        )
+
+
+def _empty_lse(q: Tensor, return_lse: bool) -> Optional[Tensor]:  # noqa: UP045
+    """An FP32 ``[batch, heads, Sq]`` LSE buffer for a BSHD ``q``, or None.
+
+    The kernel derives the LSE batch stride as heads * head stride, so this must stay contiguous.
+    """
+    if not return_lse:
+        return None
+    return torch.empty(
+        (q.shape[0], q.shape[2], q.shape[1]), dtype=torch.float32, device=q.device
+    )
+
+
+def _validate_gqa_heads(query_heads: int, kv_heads: int, operation: str) -> None:
+    """K and V may carry fewer heads than Q; the kernel addresses them through the ratio."""
+    if kv_heads == 0:
+        raise ValueError(f"{operation} requires non-empty KV heads")
+    if query_heads % kv_heads != 0:
+        raise ValueError(
+            f"{operation} 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(f"{operation} supports power-of-two GQA ratios up to 16")
+
+
 def mha_v4_packed(
     q: Tensor,
     k: Tensor,
@@ -774,21 +869,25 @@ def mha_v4_packed(
     v_pack: AttentionPack = AttentionPack.DEFAULT,
     softmax_scale: Optional[float] = None,  # noqa: UP045
     out: Optional[Tensor] = None,  # noqa: UP045
+    lse: Optional[Tensor] = None,  # noqa: UP045
     return_lse: bool = False,
+    seqlens_k: Optional[Tensor] = None,  # noqa: UP045
     kv_block_indices: Optional[Tensor] = None,  # noqa: UP045
     lut_start: Optional[Tensor] = None,  # noqa: UP045
     lut_count: Optional[Tensor] = None,  # noqa: UP045
-) -> Tensor:
+) -> Union[Tensor, tuple[Tensor, Tensor]]:  # noqa: UP007
     """Launch non-causal MHA v4 over pre-quantized BSHD operands.
 
     Formats, packing, and scale modes select an explicit ASM row. Packed widths and
     nonstandard K layouts are validated before launch; output is BF16 BSHD.
     Pass the ragged LUT triple to select the sorted-sparse row; omit all three
     tensors for dense. The work table is built inside the sparse custom op.
+    With ``return_lse`` the call returns ``(out, lse)``, where ``lse`` is FP32
+    ``[batch, heads, Sq]`` holding ``ln(sum exp(s - max)) + max``.
     """
-    if return_lse:
-        raise NotImplementedError("MHA v4 kernels do not produce LSE yet")
     lut = _packed_lut_triple(kv_block_indices, lut_start, lut_count)
+    if return_lse:
+        _check_lse_capable(q_format, v_format, lut is not None)
     _validate_pack_contract(v_format, v_pack)
     scale_modes = (q_scale_mode, k_scale_mode, v_scale_mode)
     _validate_scale_recipe(q_format, k_format, v_format, scale_modes)
@@ -800,14 +899,7 @@ 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")
-    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")
+    _validate_gqa_heads(query_heads, k.shape[2], "MHA v4")
     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:
@@ -832,6 +924,20 @@ def mha_v4_packed(
 
     if softmax_scale is None:
         softmax_scale = logical_head_dim**-0.5
+    if seqlens_k is not None:
+        if kv_block_indices is not None:
+            raise NotImplementedError(
+                "sorted-sparse MHA v4 does not accept per-batch key lengths yet"
+            )
+        _check_varlen_capable(q_format, v_format)
+        if seqlens_k.dtype != torch.int32 or seqlens_k.device != q.device:
+            raise ValueError(
+                "seqlens_k must be an int32 tensor on the same device as Q"
+            )
+        if seqlens_k.numel() < batch:
+            raise ValueError("seqlens_k needs one entry per batch")
+        if not seqlens_k.is_contiguous():
+            raise ValueError("seqlens_k must be contiguous")
     if out is None:
         out = torch.empty(
             (batch, query_length, query_heads, logical_head_dim),
@@ -843,6 +949,9 @@ def mha_v4_packed(
     elif out.dtype != torch.bfloat16 or out.device != q.device:
         raise ValueError("out must be a BF16 tensor on the same device as Q")
 
+    if lse is None or not return_lse:
+        lse = _empty_lse(q, return_lse)
+
     launch_args = (
         q,
         k,
@@ -861,12 +970,8 @@ def mha_v4_packed(
         softmax_scale,
     )
     if lut is None:
-        _mha_v4_fwd_launch(*launch_args)
+        _mha_v4_fwd_launch(*launch_args, seqlens_k, lse)
     else:
-        if q_format == AttentionFormat.BF16:
-            raise NotImplementedError(
-                "sorted-sparse MHA v4 does not have a BF16 manifest row yet"
-            )
         kv_tile = mha_v4_kv_tile()
         if k.shape[1] % kv_tile != 0:
             raise ValueError(
@@ -874,11 +979,13 @@ def mha_v4_packed(
                 f"multiple of {kv_tile}"
             )
         _mha_v4_fwd_sparse_launch(*launch_args, *lut)
+    if return_lse:
+        return out, lse
     return out
 
 
 @torch.library.custom_op(
-    "aiter::mha_v4_launch_mxfp4_coalesced_v3", mutates_args=("out",)
+    "aiter::mha_v4_launch_mxfp4_coalesced_v3", mutates_args=("out", "lse")
 )
 def _launch_mxfp4_coalesced(
     q: Tensor,
@@ -891,6 +998,7 @@ def _launch_mxfp4_coalesced(
     v_format: int,
     v_pack: int,
     softmax_scale: float,
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     resolved_v_format = AttentionFormat(v_format)
     resolved_v_pack = AttentionPack(v_pack)
@@ -918,6 +1026,8 @@ def _launch_mxfp4_coalesced(
         *scale_modes,
         softmax_scale=softmax_scale,
         out=out,
+        lse=lse,
+        return_lse=lse is not None,
         v_pack=resolved_v_pack,
     )
 
@@ -934,6 +1044,7 @@ def _launch_mxfp4_coalesced_fake(
     v_format: int,
     v_pack: int,
     softmax_scale: float,
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     del (
         q,
@@ -945,11 +1056,12 @@ def _launch_mxfp4_coalesced_fake(
         v_format,
         v_pack,
         softmax_scale,
+        lse,
     )
     del out
 
 
-@torch.library.custom_op("aiter::mha_v4_launch_mxfp6_v3", mutates_args=("out",))
+@torch.library.custom_op("aiter::mha_v4_launch_mxfp6_v3", mutates_args=("out", "lse"))
 def _launch_mxfp6(
     q: Tensor,
     q_descale: Tensor,
@@ -963,6 +1075,7 @@ def _launch_mxfp6(
     v_format: int,
     v_pack: int,
     softmax_scale: float,
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     resolved_v_format = AttentionFormat(v_format)
     resolved_v_pack = AttentionPack(v_pack)
@@ -990,6 +1103,8 @@ def _launch_mxfp6(
         *scale_modes,
         softmax_scale=softmax_scale,
         out=out,
+        lse=lse,
+        return_lse=lse is not None,
         v_pack=resolved_v_pack,
     )
 
@@ -1008,12 +1123,62 @@ def _launch_mxfp6_fake(
     v_format: int,
     v_pack: int,
     softmax_scale: float,
+    lse: Optional[Tensor],  # noqa: UP045
 ) -> None:
     del q, q_descale, k_raw, k_descale_raw, v_data, v_descale
-    del sequence_k, heads, v_format, v_pack, softmax_scale
+    del sequence_k, heads, v_format, v_pack, softmax_scale, lse
     del out
 
 
+def _k_mean(k: Tensor, kind: _RawRecipeKind) -> Optional[Tensor]:  # noqa: UP045
+    """A per-(batch, head, channel) constant to remove from K, or None if it would not pay.
+
+    Softmax is shift-invariant in a component shared by every key, but quantization noise is not,
+    so once that component dominates, removing it is a large win: measured 3.5x on FP8 and 5.1x on
+    MXFP4 at the extreme. Recipes that never quantize K have nothing to gain.
+
+    Shift-invariance holds for *any* constant vector, not just the exact mean, so this estimates it
+    from a strided sample. That keeps the cost independent of sequence length -- a full reduction
+    over K is what made centering too expensive to keep before -- and a few thousand rows estimate
+    the shared component to well under the accuracy it is worth removing.
+
+    It is also not a win at every magnitude. The subtraction takes energy out of K's RMS but not
+    out of its outliers, so a per-tensor scale (set by amax) gets relatively coarser. Below a
+    common mode of ~0.85 that costs more than the shared component does, and real traces sit at
+    0.12-0.76, so the gate leaves them untouched.
+    """
+    if kind in (_RawRecipeKind.BF16, _RawRecipeKind.BF16_FP8):
+        return None
+    stride = max(1, k.shape[1] // _K_SMOOTH_SAMPLE_ROWS)
+    sample = k[:, ::stride].float()
+    mean = sample.mean(dim=1)
+    # Compared squared to keep the whole gate to a handful of tiny kernels; it is launch-bound.
+    row_sq = sample.pow(2).sum(dim=-1).mean(dim=1)
+    gate = mean.pow(2).sum(dim=-1) > _K_SMOOTH_MIN_COMMON**2 * row_sq
+    return (mean * gate.unsqueeze(-1)).contiguous()
+
+
+def _restore_k_mean_in_lse(
+    lse: Optional[Tensor],  # noqa: UP045
+    q: Tensor,
+    k_mean: Optional[Tensor],  # noqa: UP045
+    softmax_scale: Optional[float],  # noqa: UP045
+) -> Optional[Tensor]:  # noqa: UP045
+    """Undo K smoothing in the exported LSE.
+
+    Smoothing runs the kernel against ``k - k_mean``, which shifts every score in the call by the
+    per-query constant ``q @ k_mean``. Output is unaffected because a shift shared by all keys
+    cancels in the softmax, but the LSE inherits it. Ring attention weights each chunk by
+    ``exp(lse)`` and derives a separate ``k_mean`` per chunk, so leaving the shift in mis-weights
+    the chunks against one another.
+    """
+    if lse is None or k_mean is None:
+        return lse
+    scale = softmax_scale if softmax_scale is not None else q.shape[-1] ** -0.5
+    lse += torch.einsum("bshd,bhd->bhs", q, k_mean.to(q.dtype)).float() * scale
+    return lse
+
+
 def _validate_mha_v4_raw_inputs(
     q: Tensor,
     k: Tensor,
@@ -1039,16 +1204,7 @@ def _validate_mha_v4_raw_inputs(
         raise ValueError(
             f"{operation} requires K and V with matching sequence and head dimensions"
         )
-    kv_heads = k.shape[2]
-    if kv_heads == 0:
-        raise ValueError(f"{operation} requires non-empty KV heads")
-    if q.shape[2] % kv_heads != 0:
-        raise ValueError(
-            f"{operation} requires query heads to be divisible by KV heads"
-        )
-    gqa_ratio = q.shape[2] // kv_heads
-    if gqa_ratio > 16 or gqa_ratio & (gqa_ratio - 1):
-        raise ValueError(f"{operation} supports power-of-two GQA ratios up to 16")
+    _validate_gqa_heads(q.shape[2], k.shape[2], operation)
     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:
@@ -1070,7 +1226,8 @@ def mha_v4(
     q_scale_mode: Optional[AttentionScaleMode] = None,  # noqa: UP045
     k_scale_mode: Optional[AttentionScaleMode] = None,  # noqa: UP045
     v_scale_mode: Optional[AttentionScaleMode] = None,  # noqa: UP045
-) -> Tensor:
+    seqlens_k: Optional[Tensor] = None,  # noqa: UP045
+) -> Union[Tensor, tuple[Tensor, Tensor]]:  # noqa: UP007
     """Quantize BF16 BSHD operands and run non-causal MHA v4.
 
     Q and K formats must match. Formats select the canonical quantizers and
@@ -1083,9 +1240,19 @@ def mha_v4(
     and 256x64 on gfx942. Sparse LUT rows are one per query head; K/V addressing
     uses the GQA ratio. A row may select nothing: an all-False row is a no-op
     that writes a zero output tile.
+    With ``return_lse`` the call returns ``(out, lse)``, where ``lse`` is FP32
+    ``[batch, heads, Sq]`` holding ``ln(sum exp(s - max)) + max``.
     """
     if return_lse:
-        raise NotImplementedError("MHA v4 kernels do not produce LSE yet")
+        _check_lse_capable(q_format, v_format, block_mask is not None)
+    # Checked here as well as in mha_v4_packed: the MXFP4 and MXFP6 recipes return through their
+    # own launchers, which never forward seqlens_k.
+    if seqlens_k is not None:
+        if block_mask is not None:
+            raise NotImplementedError(
+                "sorted-sparse MHA v4 does not accept per-batch key lengths yet"
+            )
+        _check_varlen_capable(q_format, v_format)
     out = _validate_mha_v4_raw_inputs(q, k, v, out, "mha_v4")
     sparse = block_mask is not None
     recipe = _resolve_raw_recipe(
@@ -1099,6 +1266,15 @@ def mha_v4(
     )
     q_scale_mode, k_scale_mode, v_scale_mode = recipe.scale_modes
 
+    # Every quantized K path fuses the subtraction into its rotation kernel, except INT8, whose
+    # quantizer is still Triton and so needs a materialised K.
+    k_mean = _k_mean(k, recipe.kind)
+    # INT8 materialises the subtraction below and drops k_mean, so keep it for the LSE correction.
+    k_mean_lse = k_mean
+    if k_mean is not None and recipe.kind is _RawRecipeKind.INT8_FP8:
+        k = (k.float() - k_mean.unsqueeze(1)).to(k.dtype)
+        k_mean = None
+
     lut_indices: Optional[Tensor] = None  # noqa: UP045
     lut_start: Optional[Tensor] = None  # noqa: UP045
     lut_count: Optional[Tensor] = None  # noqa: UP045
@@ -1109,6 +1285,9 @@ def mha_v4(
         "lut_start": lut_start,
         "lut_count": lut_count,
     }
+    # Set by the two recipes that reach the kernel through their own custom op instead of calling
+    # mha_v4_packed directly; the shared tail below runs it.
+    launch = None
     if recipe.kind == _RawRecipeKind.BF16:
         q_quantized, q_descale = q, q
         k_quantized, k_descale = k, k
@@ -1121,7 +1300,7 @@ def mha_v4(
         if softmax_scale is None:
             softmax_scale = 128**-0.5
         q_quantized, q_descale = quantize_mxfp8_q(q, mha_v4_q_multiplier(softmax_scale))
-        k_quantized, k_descale = quantize_mxfp8_k(k)
+        k_quantized, k_descale = quantize_mxfp8_k(k, k_mean)
         v_quantized, v_descale = quantize_fp8(v)
     elif recipe.kind == _RawRecipeKind.INT8_FP8:
         q_quantized, q_descale = quantize_int8(q)
@@ -1129,7 +1308,7 @@ def mha_v4(
         v_quantized, v_descale = quantize_fp8(v)
     elif recipe.kind == _RawRecipeKind.FP8:
         q_quantized, q_descale = quantize_fp8_rotated(q)
-        k_quantized, k_descale = quantize_fp8_rotated(k)
+        k_quantized, k_descale = quantize_fp8_rotated(k, k_mean)
         if _is_fp8_format(v_format):
             v_quantized, v_descale = quantize_fp8(v)
         elif recipe.v_pack == AttentionPack.V_FOR_FP6_P:
@@ -1140,13 +1319,11 @@ def mha_v4(
         if softmax_scale is None:
             softmax_scale = 128**-0.5
         q_quantized, q_descale = quantize_mxfp4_q(q, mha_v4_q_multiplier(softmax_scale))
-        k_quantized, k_descale = quantize_mxfp4_k(k)
-        if _is_fp8_format(v_format):
-            v_quantized, v_descale = quantize_v_fp8(v)
-        else:
-            v_quantized, v_descale = quantize_v_mxfp4(v)
+        k_quantized, k_descale = quantize_mxfp4_k(k, k_mean)
+        v_quantized, v_descale = quantize_v_mxfp4_fp6_p(v)
         if lut_indices is None:
-            _launch_mxfp4_coalesced(
+            launch = functools.partial(
+                _launch_mxfp4_coalesced,
                 q_quantized,
                 q_descale,
                 k_quantized,
@@ -1158,30 +1335,25 @@ def mha_v4(
                 int(recipe.v_pack),
                 softmax_scale,
             )
-            return out
-        k_view = mxfp4_k_view(k_quantized, k_descale)
-        v_view = (
-            v_quantized
-            if _is_fp8_format(v_format)
-            else mxfp4_v_view(v_quantized, v_descale, k.shape[1])
-        )
-        k_quantized = k_view
-        v_quantized = v_view
+        else:
+            k_view = mxfp4_k_view(k_quantized, k_descale)
+            v_view = mxfp4_v_view(v_quantized, v_descale, k.shape[1])
+            k_quantized = k_view
+            v_quantized = v_view
     elif recipe.kind == _RawRecipeKind.MXFP6:
         if softmax_scale is None:
             softmax_scale = 128**-0.5
         q_quantized, q_descale = quantize_mxfp6_q(q, mha_v4_q_multiplier(softmax_scale))
-        k_quantized, k_descale = quantize_mxfp6_k(k)
+        k_quantized, k_descale = quantize_mxfp6_k(k, k_mean)
         if _is_fp8_format(v_format):
             v_quantized, v_descale = quantize_v_fp8(v)
         elif v_format == AttentionFormat.MXFP6:
             v_quantized, v_descale = quantize_v_mxfp6_fp6_p(v)
-        elif recipe.v_pack == AttentionPack.V_FOR_FP6_P:
-            v_quantized, v_descale = quantize_v_mxfp4_fp6_p(v)
         else:
-            v_quantized, v_descale = quantize_v_mxfp4(v)
+            v_quantized, v_descale = quantize_v_mxfp4_fp6_p(v)
         if lut_indices is None:
-            _launch_mxfp6(
+            launch = functools.partial(
+                _launch_mxfp6,
                 q_quantized,
                 q_descale,
                 k_quantized,
@@ -1195,22 +1367,29 @@ def mha_v4(
                 int(recipe.v_pack),
                 softmax_scale,
             )
-            return out
-        k_view, k_descale_view = mxfp6_k_view(
-            k_quantized, k_descale, q.shape[0], k.shape[1], k.shape[2]
-        )
-        v_view = (
-            v_quantized
-            if v_format != AttentionFormat.MXFP4
-            else mxfp4_v_view(v_quantized, v_descale, k.shape[1])
-        )
-        k_quantized = k_view
-        k_descale = k_descale_view
-        v_quantized = v_view
+        else:
+            k_view, k_descale_view = mxfp6_k_view(
+                k_quantized, k_descale, q.shape[0], k.shape[1], k.shape[2]
+            )
+            v_view = (
+                v_quantized
+                if v_format != AttentionFormat.MXFP4
+                else mxfp4_v_view(v_quantized, v_descale, k.shape[1])
+            )
+            k_quantized = k_view
+            k_descale = k_descale_view
+            v_quantized = v_view
     else:
         raise AssertionError(f"unhandled MHA v4 raw recipe: {recipe.kind!r}")
 
-    return mha_v4_packed(
+    if launch is not None:
+        lse_out = _empty_lse(q, return_lse)
+        launch(lse_out)
+        if return_lse:
+            return out, _restore_k_mean_in_lse(lse_out, q, k_mean_lse, softmax_scale)
+        return out
+
+    result = mha_v4_packed(
         q_quantized,
         k_quantized,
         v_quantized,
@@ -1227,43 +1406,12 @@ def mha_v4(
         out=out,
         return_lse=return_lse,
         v_pack=recipe.v_pack,
+        seqlens_k=seqlens_k,
         **packed_lut,
     )
-
-
-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,
-    block_mask: Optional[Tensor] = None,  # noqa: UP045
-) -> Tensor:
-    """Quantize BF16 BSHD Q/K to MXFP8 and V to per-tensor FP8.
-
-    Deprecated: this recipe is reachable through :func:`mha_v4` by passing FP8
-    formats with E8M0 per-1x32 Q/K scale modes.
-    """
-    warnings.warn(
-        "mha_v4_mxfp8 is deprecated; call mha_v4 with FP8 formats and "
-        "E8M0_PER_1X32 Q/K scale modes instead",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    fp8_format = native_fp8_format()
-    return mha_v4(
-        q,
-        k,
-        v,
-        fp8_format,
-        fp8_format,
-        fp8_format,
-        softmax_scale=softmax_scale,
-        out=out,
-        return_lse=return_lse,
-        block_mask=block_mask,
-        q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-        k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-        v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
-    )
+    if return_lse:
+        packed_out, packed_lse = result
+        return packed_out, _restore_k_mean_in_lse(
+            packed_lse, q, k_mean_lse, softmax_scale
+        )
+    return result
diff --git a/aiter/ops/mha_v4_quant.py b/aiter/ops/mha_v4_quant.py
index e12eae31595..3407efa33d3 100644
--- a/aiter/ops/mha_v4_quant.py
+++ b/aiter/ops/mha_v4_quant.py
@@ -3,6 +3,8 @@
 
 """Quantization and packed-layout producers for MHA v4."""
 
+from typing import Optional
+
 import torch
 import triton
 from torch import Tensor
@@ -60,17 +62,41 @@ def mha_v4_q_multiplier(softmax_scale: float) -> float:
     return softmax_scale * MHA_V4_LOG2E
 
 
-@compile_ops("module_mha_v4_quant", develop=True)
-def rotate_activation_hd128(out: Tensor, input: Tensor) -> None:
+@compile_ops("module_mha_v4_quant", fc_name="rotate_activation_hd128", develop=True)
+def _rotate_activation_hd128(
+    out: Tensor, input: Tensor, mean: Tensor, partial_amax: Tensor
+) -> None:
     """Apply normalized Walsh-Hadamard rotation to contiguous hd128 rows."""
 
 
+def _or_empty(input: Tensor, mean: Optional[Tensor]) -> Tensor:  # noqa: UP045
+    """The quantizer kernels read an empty tensor as "no K mean"."""
+    return input.new_empty((0,), dtype=torch.float32) if mean is None else mean
+
+
+def rotate_activation_hd128(
+    out: Tensor,
+    input: Tensor,
+    mean: Optional[Tensor] = None,  # noqa: UP045
+    partial_amax: Optional[Tensor] = None,  # noqa: UP045
+) -> None:
+    """Rotate hd128 rows, first subtracting a (batch, heads, 128) fp32 `mean` when given.
+
+    Pass `partial_amax` to have the kernel also emit a per-block amax of the rotated values,
+    which saves a per-tensor quantizer the full extra read its own amax pass would cost.
+    """
+    _rotate_activation_hd128(
+        out, input, _or_empty(input, mean), _or_empty(input, partial_amax)
+    )
+
+
 @compile_ops("module_mha_v4_quant", develop=True)
 def rotate_activation_mxfp8_quant(
     out: Tensor,
     scale: Tensor,
     input: Tensor,
     multiplier: float,
+    mean: Tensor,
 ) -> None:
     """Apply hd128 Walsh-Hadamard rotation and quantize directly to MXFP8."""
 
@@ -90,6 +116,7 @@ def rotate_activation_mxfp6_quant_k(
     out: Tensor,
     scale: Tensor,
     input: Tensor,
+    mean: Tensor,
 ) -> None:
     """Rotate and pack hd128 K directly into the MXFP6 LDS-order buffers."""
 
@@ -118,6 +145,7 @@ def rotate_activation_mxfp4_quant_k(
     out: Tensor,
     scale: Tensor,
     input: Tensor,
+    mean: Tensor,
 ) -> None:
     """Apply hd128 Walsh-Hadamard rotation and pack K in the MXFP4 ASM tile order."""
 
@@ -146,32 +174,52 @@ def _validate_bshd_hd128(input: Tensor, operation: str) -> tuple[int, int, int,
     return input.shape
 
 
+# vec_size(16) * warp(64) / dim(128); must track the hd128 rotate kernel's m_block.
+MHA_V4_ROTATE_ROWS_PER_BLOCK = 8
+
+
 def _quantize_per_tensor(
-    input: Tensor, output_dtype: torch.dtype, dtype_max: float, clip: float
+    input: Tensor,
+    output_dtype: torch.dtype,
+    dtype_max: float,
+    clip: float,
+    partial: Optional[Tensor] = None,  # noqa: UP045
 ) -> tuple[Tensor, Tensor]:
+    """Quantize to a single scale. `partial` supplies precomputed per-block amaxes."""
     if not input.is_contiguous():
         raise ValueError("MHA v4 per-tensor quantization requires contiguous input")
     numel = input.numel()
     blocks = triton.cdiv(numel, MHA_V4_PER_TENSOR_BLOCK_SIZE)
-    partial = input.new_empty((blocks,), dtype=torch.float32)
     scale = input.new_empty((1,), dtype=torch.float32)
     output = input.new_empty(input.shape, dtype=output_dtype)
-    mha_v4_per_tensor_amax_kernel[(blocks,)](
-        input,
-        partial,
-        numel,
-        BLOCK_SIZE=MHA_V4_PER_TENSOR_BLOCK_SIZE,
-        num_warps=8,
-    )
-    scale_block = triton.next_power_of_2(blocks)
-    mha_v4_per_tensor_scale_kernel[(1,)](
-        partial,
-        scale,
-        blocks,
-        dtype_max=dtype_max / clip,
-        BLOCK_SIZE=scale_block,
-        num_warps=8,
-    )
+    if partial is None:
+        partial = input.new_empty((blocks,), dtype=torch.float32)
+        mha_v4_per_tensor_amax_kernel[(blocks,)](
+            input,
+            partial,
+            numel,
+            BLOCK_SIZE=MHA_V4_PER_TENSOR_BLOCK_SIZE,
+            num_warps=8,
+        )
+        scale_block = triton.next_power_of_2(blocks)
+        mha_v4_per_tensor_scale_kernel[(1,)](
+            partial,
+            scale,
+            blocks,
+            dtype_max=dtype_max / clip,
+            BLOCK_SIZE=scale_block,
+            num_warps=8,
+        )
+    else:
+        # The rotate kernel emits one partial per 8-row block, far more than the per-tensor amax
+        # kernel would, so a single-program Triton reduction over them would need an unusable
+        # BLOCK_SIZE. torch reduces the small fp32 vector instead.
+        # An all-zero tensor must yield scale 1.0, not 0 (which would make the quantized
+        # data NaN) and not an epsilon; this matches the Triton scale kernel's contract.
+        amax = partial.amax()
+        scale.copy_(
+            torch.where(amax > 0, amax / (dtype_max / clip), amax.new_ones(())).view(1)
+        )
     mha_v4_per_tensor_quant_kernel[(blocks,)](
         input,
         output,
@@ -211,13 +259,23 @@ def _quantize_fp8_fake(input: Tensor) -> tuple[Tensor, Tensor]:
     )
 
 
-def quantize_fp8_rotated(input: Tensor) -> tuple[Tensor, Tensor]:
-    """Apply normalized hd128 Walsh-Hadamard rotation, then per-tensor FP8 quantize."""
+def quantize_fp8_rotated(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
+    """Rotate hd128 rows, optionally removing `mean` first, then per-tensor FP8 quantize.
+
+    The rotation kernel emits the amax as it goes, so this costs two passes over the tensor
+    rather than three.
+    """
     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_hd128(rotated, input)
-    return quantize_fp8(rotated)
+    blocks = triton.cdiv(input.numel() // 128, MHA_V4_ROTATE_ROWS_PER_BLOCK)
+    partial = input.new_empty((blocks,), dtype=torch.float32)
+    rotate_activation_hd128(rotated, input, mean, partial)
+    return _quantize_per_tensor(
+        rotated, dtypes.fp8, torch.finfo(dtypes.fp8).max, 1.0, partial
+    )
 
 
 def block_scale_storage(
@@ -254,43 +312,51 @@ def query_block_scale(
     )
 
 
+def _plain_block_scale(input: Tensor) -> Tensor:
+    """One E8M0 scale per 32-element block, without the gather padding.
+
+    Used where the consumer addresses scales within the logical sequence, and by the fakes, which
+    only have to agree on shape and dtype.
+    """
+    batch, sequence, heads, head_dim = input.shape
+    return input.new_empty((batch, sequence, heads, head_dim // 32), dtype=torch.uint8)
+
+
 @torch.library.custom_op("aiter::mha_v4_quantize_mxfp8_q", mutates_args=())
 def quantize_mxfp8_q(input: Tensor, multiplier: float) -> tuple[Tensor, Tensor]:
     """Rotate and quantize hd128 BSHD Q to MXFP8 data and E8M0 block scales."""
     batch, sequence, heads, head_dim = _validate_bshd_hd128(input, "MXFP8 quantization")
     quantized = input.new_empty(input.shape, dtype=dtypes.fp8)
     scale = query_block_scale(input, batch, sequence, heads, head_dim // 32)
-    rotate_activation_mxfp8_quant(quantized, scale, input, multiplier)
+    rotate_activation_mxfp8_quant(
+        quantized, scale, input, multiplier, _or_empty(input, None)
+    )
     return quantized, scale
 
 
 @quantize_mxfp8_q.register_fake
 def _quantize_mxfp8_q_fake(input: Tensor, multiplier: float) -> tuple[Tensor, Tensor]:
     del multiplier
-    batch, sequence, heads, head_dim = input.shape
-    return input.new_empty(input.shape, dtype=dtypes.fp8), input.new_empty(
-        (batch, sequence, heads, head_dim // 32), dtype=torch.uint8
-    )
+    return input.new_empty(input.shape, dtype=dtypes.fp8), _plain_block_scale(input)
 
 
 @torch.library.custom_op("aiter::mha_v4_quantize_mxfp8_k", mutates_args=())
-def quantize_mxfp8_k(input: Tensor) -> tuple[Tensor, Tensor]:
+def quantize_mxfp8_k(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
     """Rotate and quantize hd128 BSHD K to MXFP8 data and E8M0 block scales."""
-    batch, sequence, heads, head_dim = _validate_bshd_hd128(
-        input, "MXFP8 K quantization"
-    )
+    _validate_bshd_hd128(input, "MXFP8 K quantization")
     quantized = input.new_empty(input.shape, dtype=dtypes.fp8)
-    scale = input.new_empty((batch, sequence, heads, head_dim // 32), dtype=torch.uint8)
-    rotate_activation_mxfp8_quant(quantized, scale, input, 1.0)
+    scale = _plain_block_scale(input)
+    rotate_activation_mxfp8_quant(quantized, scale, input, 1.0, _or_empty(input, mean))
     return quantized, scale
 
 
 @quantize_mxfp8_k.register_fake
-def _quantize_mxfp8_k_fake(input: Tensor) -> tuple[Tensor, Tensor]:
-    batch, sequence, heads, head_dim = input.shape
-    return input.new_empty(input.shape, dtype=dtypes.fp8), input.new_empty(
-        (batch, sequence, heads, head_dim // 32), dtype=torch.uint8
-    )
+def _quantize_mxfp8_k_fake(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
+    return input.new_empty(input.shape, dtype=dtypes.fp8), _plain_block_scale(input)
 
 
 @torch.library.custom_op("aiter::mha_v4_quantize_mxfp4", mutates_args=())
@@ -311,7 +377,7 @@ def _quantize_mxfp4_q_fake(input: Tensor, multiplier: float) -> tuple[Tensor, Te
     batch, sequence, heads, head_dim = input.shape
     return input.new_empty(
         (batch, sequence, heads, head_dim // 2), dtype=torch.uint8
-    ), input.new_empty((batch, sequence, heads, head_dim // 32), dtype=torch.uint8)
+    ), _plain_block_scale(input)
 
 
 def mxfp4_k_raw_buffer_size(batch: int, sequence: int, heads: int) -> int:
@@ -357,7 +423,9 @@ def mxfp6_v_raw_buffer_size(batch: int, sequence: int, heads: int) -> int:
 
 
 @torch.library.custom_op("aiter::mha_v4_quantize_mxfp4_k_raw", mutates_args=())
-def quantize_mxfp4_k(input: Tensor) -> tuple[Tensor, Tensor]:
+def quantize_mxfp4_k(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
     """Rotate and pack hd128 BSHD K into the coalesced MXFP4 ASM layout."""
     batch, sequence, heads, head_dim = _validate_bshd_hd128(
         input, "MXFP4 K quantization"
@@ -375,12 +443,14 @@ def quantize_mxfp4_k(input: Tensor) -> tuple[Tensor, Tensor]:
         lookahead_rows=MHA_V4_KV_SCALE_LOOKAHEAD_ROWS,
         extra=MHA_V4_MXFP4_K_SCALE_SLACK_BYTES,
     )
-    rotate_activation_mxfp4_quant_k(raw, scale, input)
+    rotate_activation_mxfp4_quant_k(raw, scale, input, _or_empty(input, mean))
     return raw, scale
 
 
 @quantize_mxfp4_k.register_fake
-def _quantize_mxfp4_k_fake(input: Tensor) -> tuple[Tensor, Tensor]:
+def _quantize_mxfp4_k_fake(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
     batch, sequence, heads, head_dim = input.shape
     scale_elements = batch * sequence * heads * (head_dim // 32)
     scale_storage = input.new_empty(
@@ -434,22 +504,26 @@ def _quantize_mxfp6_q_fake(input: Tensor, multiplier: float) -> tuple[Tensor, Te
     batch, sequence, heads, head_dim = input.shape
     return input.new_empty(
         (batch, sequence, heads, head_dim // 32 * 24), dtype=torch.uint8
-    ), input.new_empty((batch, sequence, heads, head_dim // 32), dtype=torch.uint8)
+    ), _plain_block_scale(input)
 
 
 @torch.library.custom_op("aiter::mha_v4_quantize_mxfp6_k_raw", mutates_args=())
-def quantize_mxfp6_k(input: Tensor) -> tuple[Tensor, Tensor]:
+def quantize_mxfp6_k(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
     """Rotate and pack hd128 BSHD K into raw MXFP6 ASM data and scale buffers."""
     batch, sequence, heads, _ = _validate_bshd_hd128(input, "MXFP6 E2M3 K quantization")
     data_size, scale_size = fp6_k_raw_buffer_sizes(batch, sequence, heads)
     raw = input.new_empty((data_size,), dtype=torch.uint8)
     scale_raw = input.new_empty((scale_size,), dtype=torch.uint8)
-    rotate_activation_mxfp6_quant_k(raw, scale_raw, input)
+    rotate_activation_mxfp6_quant_k(raw, scale_raw, input, _or_empty(input, mean))
     return raw, scale_raw
 
 
 @quantize_mxfp6_k.register_fake
-def _quantize_mxfp6_k_raw_fake(input: Tensor) -> tuple[Tensor, Tensor]:
+def _quantize_mxfp6_k_raw_fake(
+    input: Tensor, mean: Optional[Tensor] = None  # noqa: UP045
+) -> tuple[Tensor, Tensor]:
     batch, sequence, heads, _ = input.shape
     data_size, scale_size = fp6_k_raw_buffer_sizes(batch, sequence, heads)
     return input.new_empty((data_size,), dtype=torch.uint8), input.new_empty(
@@ -545,21 +619,6 @@ def _mxfp4_v_buffers(
     return raw, scale
 
 
-@torch.library.custom_op("aiter::mha_v4_quantize_v_mxfp4_raw_v2", mutates_args=())
-def quantize_v_mxfp4(input: Tensor) -> tuple[Tensor, Tensor]:
-    """Pack hd128 BSHD V into raw column-major MXFP4 data and scale buffers."""
-    batch, sequence, heads, _ = _validate_bshd_hd128(input, "MXFP4 V quantization")
-    raw, scale = _mxfp4_v_buffers(input, batch, sequence, heads)
-    _quantize_v_mxfp4_hip(raw, scale, input)
-    return raw, scale
-
-
-@quantize_v_mxfp4.register_fake
-def _quantize_v_mxfp4_raw_fake(input: Tensor) -> tuple[Tensor, Tensor]:
-    batch, sequence, heads, _ = input.shape
-    return _mxfp4_v_buffers(input, batch, sequence, heads)
-
-
 @torch.library.custom_op("aiter::mha_v4_quantize_v_mxfp4_fp6_p_raw", mutates_args=())
 def quantize_v_mxfp4_fp6_p(input: Tensor) -> tuple[Tensor, Tensor]:
     """Pack MXFP4 V in the token order consumed by the FP6-P F4F4 kernel."""
diff --git a/aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py b/aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py
index a904e5f0858..a098f081f00 100644
--- a/aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py
+++ b/aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py
@@ -339,7 +339,9 @@ def sage_quant_v_amax_finalize_kernel(
         mask=(offs_n[:, None] < NUM_SEQ_BLKS) & (offs_d[None, :] < D),
         other=0.0,
     )
-    scale = tl.max(partial, axis=0) * (1.0 / FP8_MAX)
+    # An all-zero channel would otherwise give scale 0, and dividing by it turns the whole
+    # channel into NaN rather than zero. Sequence-parallel head padding produces exactly that.
+    scale = tl.maximum(tl.max(partial, axis=0), 1e-12) * (1.0 / FP8_MAX)
     tl.store(V_Scale + bh * D + offs_d, scale, mask=offs_d < D)
 
 
diff --git a/csrc/include/mha_v4_quant.h b/csrc/include/mha_v4_quant.h
index cd59f222ff8..28fae61da11 100644
--- a/csrc/include/mha_v4_quant.h
+++ b/csrc/include/mha_v4_quant.h
@@ -7,14 +7,21 @@
 namespace aiter {
 namespace torch_itfs {
 
-// Apply normalized Walsh-Hadamard rotation to contiguous hd128 rows.
-void rotate_activation_hd128(aiter_tensor_t& out, const aiter_tensor_t& input);
-
-// Rotate hd128 rows and emit token-major MX data plus one E8M0 scale per 32 values.
+// Apply normalized Walsh-Hadamard rotation to contiguous hd128 rows. When `mean` is a valid
+// (batch, heads, 128) fp32 tensor its per-(batch, head, channel) values are subtracted first,
+// fusing K-smoothing into the rotation pass; pass an empty tensor to skip it.
+void rotate_activation_hd128(aiter_tensor_t& out,
+                             const aiter_tensor_t& input,
+                             const aiter_tensor_t& mean,
+                             aiter_tensor_t& partial_amax);
+
+// Rotate hd128 rows and emit token-major MX data plus one E8M0 scale per 32 values. `mean` is
+// the optional K mean documented above; pass an empty tensor for Q.
 void rotate_activation_mxfp8_quant(aiter_tensor_t& out,
                                    aiter_tensor_t& scale,
                                    const aiter_tensor_t& input,
-                                   double multiplier);
+                                   double multiplier,
+                                   const aiter_tensor_t& mean);
 
 void rotate_activation_mxfp6_quant(aiter_tensor_t& out,
                                    aiter_tensor_t& scale,
@@ -24,7 +31,8 @@ void rotate_activation_mxfp6_quant(aiter_tensor_t& out,
 // The K variants write the coalesced tile layouts consumed directly by MHA v4 code objects.
 void rotate_activation_mxfp6_quant_k(aiter_tensor_t& out,
                                      aiter_tensor_t& scale,
-                                     const aiter_tensor_t& input);
+                                     const aiter_tensor_t& input,
+                                     const aiter_tensor_t& mean);
 
 void quantize_v_mxfp6_fp6_p(aiter_tensor_t& out,
                             aiter_tensor_t& scale,
@@ -37,7 +45,8 @@ void rotate_activation_mxfp4_quant(aiter_tensor_t& out,
 
 void rotate_activation_mxfp4_quant_k(aiter_tensor_t& out,
                                      aiter_tensor_t& scale,
-                                     const aiter_tensor_t& input);
+                                     const aiter_tensor_t& input,
+                                     const aiter_tensor_t& mean);
 
 void quantize_v_mxfp4_fp6_p(aiter_tensor_t& out,
                             aiter_tensor_t& scale,
diff --git a/csrc/include/torch/mha_v4_fwd.h b/csrc/include/torch/mha_v4_fwd.h
index f365b0c0a59..ce7ac3f3f38 100644
--- a/csrc/include/torch/mha_v4_fwd.h
+++ b/csrc/include/torch/mha_v4_fwd.h
@@ -22,7 +22,9 @@ void fmha_v4_fwd(const at::Tensor& q,
                  int64_t q_scale_mode,
                  int64_t k_scale_mode,
                  int64_t v_scale_mode,
-                 double softmax_scale);
+                 double softmax_scale,
+                 std::optional<at::Tensor> seqlens_k = std::nullopt,
+                 std::optional<at::Tensor> lse       = std::nullopt);
 
 // Sorted block-sparse sibling. Same packed operands as fmha_v4_fwd, plus a ragged LUT.
 // Builds the work table internally (identity raster if lut_count is uniform, else LPT).
diff --git a/csrc/kernels/mha_v4_quant.cu b/csrc/kernels/mha_v4_quant.cu
index 6e1c4989efb..bd4669b8c4f 100644
--- a/csrc/kernels/mha_v4_quant.cu
+++ b/csrc/kernels/mha_v4_quant.cu
@@ -52,12 +52,36 @@ __device__ float swap_thread_data(float data)
     return data;
 }
 
+// Remove K's per-(batch, head, channel) token mean. Softmax is shift-invariant in a component
+// shared by every key, but quantization noise is not. No-op when `mean` is null.
+template <int vec_size, int dim, typename floatxvec_t>
+__device__ inline void subtract_k_mean(floatxvec_t& af,
+                                       float const* __restrict__ mean,
+                                       const int32_t row,
+                                       const int32_t lane,
+                                       const int32_t heads,
+                                       const int32_t seq_heads)
+{
+    if(mean == nullptr)
+        return;
+    const int32_t mean_row = (row / seq_heads) * heads + (row % heads);
+    float const* __restrict__ mrow =
+        mean + static_cast<int64_t>(mean_row) * dim + lane * vec_size;
+#pragma unroll
+    for(int i = 0; i < vec_size; i++)
+        af[i] -= mrow[i];
+}
+
 template <typename DTYPE_I, int vec_size = 16>
 __global__ void hadamard_rotate_activation_hd128_kernel(DTYPE_I* __restrict__ out,
                                                          DTYPE_I const* __restrict__ input,
+                                                         float const* __restrict__ mean,
+                                                         const int32_t heads,
+                                                         const int32_t seq_heads,
                                                          const int32_t m,
                                                          const int32_t in_stride,
-                                                         const int32_t out_stride)
+                                                         const int32_t out_stride,
+                                                         float* __restrict__ partial_amax = nullptr)
 {
     constexpr int dim         = kHeadDim;
     constexpr int warp_size   = opus::get_warp_size();
@@ -80,6 +104,9 @@ __global__ void hadamard_rotate_activation_hd128_kernel(DTYPE_I* __restrict__ ou
     for(int i = 0; i < vec_size; i++)
         af[i] = static_cast<float>(a[i]);
 
+    // K-smoothing: remove the per-(batch, head, channel) token mean before rotating.
+    subtract_k_mean<vec_size, dim>(af, mean, row, lane, heads, seq_heads);
+
     constexpr int intra_thread_loop = __builtin_ctz(vec_size);
     opus::static_for<intra_thread_loop>([&](auto i) {
         constexpr int h = 1 << i.value;
@@ -114,6 +141,24 @@ __global__ void hadamard_rotate_activation_hd128_kernel(DTYPE_I* __restrict__ ou
         *reinterpret_cast<outxvec_t*>(out + static_cast<int64_t>(row) * out_stride +
                                       lane * vec_size) = rotated;
     }
+
+    // The rotated values are already in registers, so a per-tensor amax costs a reduction rather
+    // than the extra full read a separate pass would need.
+    if(partial_amax != nullptr)
+    {
+        float local = 0.0f;
+        if(row < m)
+        {
+#pragma unroll
+            for(int i = 0; i < vec_size; i++)
+                local = fmaxf(local, fabsf(static_cast<float>(
+                                  static_cast<DTYPE_I>(af[i] * dim_rsqrt))));
+        }
+        auto max_op  = [](float a, float b) { return fmaxf(a, b); };
+        local        = multithread_reduce(local, max_op, opus::get_warp_size());
+        if(threadIdx.x == 0)
+            partial_amax[blockIdx.x] = local;
+    }
 }
 
 template <typename DTYPE_I, int vec_size = 16>
@@ -123,7 +168,10 @@ __global__ void hadamard_rotate_activation_mxfp8_quant_kernel(
     DTYPE_I const* __restrict__ input,
     const int32_t m,
     const int32_t stride,
-    const float multiplier)
+    const float multiplier,
+    float const* __restrict__ mean = nullptr,
+    const int32_t heads           = 0,
+    const int32_t seq_heads       = 0)
 {
     constexpr int dim         = 128;
     constexpr int warp_size   = opus::get_warp_size();
@@ -147,6 +195,8 @@ __global__ void hadamard_rotate_activation_mxfp8_quant_kernel(
     for(int i = 0; i < vec_size; i++)
         af[i] = static_cast<float>(a[i]);
 
+    subtract_k_mean<vec_size, dim>(af, mean, row, lane, heads, seq_heads);
+
     constexpr int intra_thread_loop = __builtin_ctz(vec_size);
     opus::static_for<intra_thread_loop>([&](auto i) {
         constexpr int h = 1 << i.value;
@@ -215,7 +265,8 @@ __global__ void hadamard_rotate_activation_mxfp6_quant_kernel(
     const float multiplier,
     const int32_t sequence = 0,
     const int32_t heads    = 0,
-    const int32_t tiles    = 0)
+    const int32_t tiles    = 0,
+    float const* __restrict__ mean = nullptr)
 {
     constexpr int dim         = 128;
     constexpr int warp_size   = opus::get_warp_size();
@@ -239,6 +290,8 @@ __global__ void hadamard_rotate_activation_mxfp6_quant_kernel(
     for(int i = 0; i < vec_size; i++)
         af[i] = static_cast<float>(a[i]);
 
+    subtract_k_mean<vec_size, dim>(af, mean, row, lane, heads, sequence * heads);
+
     constexpr int intra_thread_loop = __builtin_ctz(vec_size);
     opus::static_for<intra_thread_loop>([&](auto i) {
         constexpr int h = 1 << i.value;
@@ -423,7 +476,8 @@ __global__ void hadamard_rotate_activation_mxfp4_quant_kernel(
     const float multiplier,
     const int32_t sequence = 0,
     const int32_t heads    = 0,
-    const int32_t tiles    = 0)
+    const int32_t tiles    = 0,
+    float const* __restrict__ mean = nullptr)
 {
     constexpr int dim         = 128;
     constexpr int warp_size   = opus::get_warp_size();
@@ -447,6 +501,8 @@ __global__ void hadamard_rotate_activation_mxfp4_quant_kernel(
     for(int i = 0; i < vec_size; i++)
         af[i] = static_cast<float>(a[i]);
 
+    subtract_k_mean<vec_size, dim>(af, mean, row, lane, heads, sequence * heads);
+
     constexpr int intra_thread_loop = __builtin_ctz(vec_size);
     opus::static_for<intra_thread_loop>([&](auto i) {
         constexpr int h = 1 << i.value;
@@ -771,9 +827,34 @@ void launch_quant(aiter_tensor_t& out,
     kernel(grid, dim3(block_size), stream, m, static_cast<float>(multiplier));
 }
 
+// Validates the optional (batch, heads, 128) fp32 K mean against a BSHD input. Returns nullptr
+// and leaves the outputs untouched when the tensor is empty.
+static const float* check_k_mean(const aiter_tensor_t& mean,
+                                 const aiter_tensor_t& input,
+                                 int32_t& heads,
+                                 int32_t& seq_heads)
+{
+    if(mean.numel() == 0)
+        return nullptr;
+    AITER_CHECK(mean.is_gpu() && mean.device_id == input.device_id,
+                "K mean must be on the same GPU as input");
+    AITER_CHECK(mean.dtype() == AITER_DTYPE_fp32, "K mean must be fp32");
+    AITER_CHECK(mean.is_contiguous(), "K mean must be contiguous");
+    AITER_CHECK(mean.dim() == 3 && mean.size(2) == kHeadDim,
+                "K mean must be (batch, heads, 128)");
+    AITER_CHECK(input.dim() == 4, "K mean requires BSHD input");
+    heads     = static_cast<int32_t>(mean.size(1));
+    seq_heads = static_cast<int32_t>(input.size(1)) * heads;
+    AITER_CHECK(static_cast<int32_t>(input.size(2)) == heads, "K mean heads must match input");
+    return reinterpret_cast<const float*>(mean.data_ptr());
+}
+
 } // namespace
 
-void rotate_activation_hd128(aiter_tensor_t& out, const aiter_tensor_t& input)
+void rotate_activation_hd128(aiter_tensor_t& out,
+                             const aiter_tensor_t& input,
+                             const aiter_tensor_t& mean,
+                             aiter_tensor_t& partial_amax)
 {
     constexpr int32_t dim        = kHeadDim;
     constexpr int32_t block_size = WARP_SIZE;
@@ -802,24 +883,44 @@ void rotate_activation_hd128(aiter_tensor_t& out, const aiter_tensor_t& input)
     const int32_t out_stride = dim;
     const dim3 grid((m + m_block - 1) / m_block);
     HipDeviceGuard device_guard(input.device_id);
+    const float* mean_ptr = nullptr;
+    int32_t heads         = 1;
+    int32_t seq_heads     = m;
+    mean_ptr              = check_k_mean(mean, input, heads, seq_heads);
+    float* amax_ptr = nullptr;
+    if(partial_amax.numel() > 0)
+    {
+        AITER_CHECK(partial_amax.dtype() == AITER_DTYPE_fp32, "partial amax must be fp32");
+        AITER_CHECK(partial_amax.numel() >= grid.x,
+                    "partial amax needs one slot per rotate block");
+        amax_ptr = reinterpret_cast<float*>(partial_amax.data_ptr());
+    }
     const hipStream_t stream = aiter::getCurrentHIPStream();
     AITER_DISPATCH_FLOATING16_TYPES_rmTorch(input.dtype(), "rotate_activation_hd128", [&] {
         using DTYPE_I = typename aiter::hip2opus<scalar_t>::type;
         hadamard_rotate_activation_hd128_kernel<DTYPE_I><<<grid, dim3(block_size), 0, stream>>>(
             reinterpret_cast<DTYPE_I*>(out.data_ptr()),
             reinterpret_cast<DTYPE_I const*>(input.data_ptr()),
+            mean_ptr,
+            heads,
+            seq_heads,
             m,
             in_stride,
-            out_stride);
+            out_stride,
+            amax_ptr);
     });
 }
 
 void rotate_activation_mxfp8_quant(aiter_tensor_t& out,
                                    aiter_tensor_t& scale,
                                    const aiter_tensor_t& input,
-                                   const double multiplier)
+                                   const double multiplier,
+                                   const aiter_tensor_t& mean)
 {
     check_inputs<128, AITER_DTYPE_fp8>(out, scale, input);
+    int32_t mean_heads     = 1;
+    int32_t mean_seq_heads = static_cast<int32_t>(input.numel() / kHeadDim);
+    const float* mean_ptr  = check_k_mean(mean, input, mean_heads, mean_seq_heads);
     AITER_DISPATCH_FLOATING16_TYPES_rmTorch(input.dtype(), "rotate_activation_mxfp8_quant", [&] {
         using DTYPE_I = typename aiter::hip2opus<scalar_t>::type;
         launch_quant(out, scale, input, multiplier, [&](dim3 grid,
@@ -833,7 +934,10 @@ void rotate_activation_mxfp8_quant(aiter_tensor_t& out,
                 reinterpret_cast<DTYPE_I const*>(input.data_ptr()),
                 m,
                 128,
-                factor);
+                factor,
+                mean_ptr,
+                mean_heads,
+                mean_seq_heads);
         });
     });
 }
@@ -864,7 +968,8 @@ void rotate_activation_mxfp6_quant(aiter_tensor_t& out,
 
 void rotate_activation_mxfp6_quant_k(aiter_tensor_t& out,
                                      aiter_tensor_t& scale,
-                                     const aiter_tensor_t& input)
+                                     const aiter_tensor_t& input,
+                                     const aiter_tensor_t& mean)
 {
     constexpr int64_t tile = kHeadDim;
     AITER_CHECK(input.dim() == 4, "input must be BSHD");
@@ -880,6 +985,9 @@ void rotate_activation_mxfp6_quant_k(aiter_tensor_t& out,
     // aiter_tensor_t has no as_strided; pass the logical (unpadded) out/scale numel
     // so check_inputs validates the m*96 / m*4 layout without a strided view.
     check_inputs<96>(out, scale, input, m * 96, m * 4);
+    int32_t mean_heads     = 1;
+    int32_t mean_seq_heads = static_cast<int32_t>(m);
+    const float* mean_ptr  = check_k_mean(mean, input, mean_heads, mean_seq_heads);
     AITER_DISPATCH_FLOATING16_TYPES_rmTorch(input.dtype(), "rotate_activation_mxfp6_quant_k", [&] {
         using DTYPE_I = typename aiter::hip2opus<scalar_t>::type;
         constexpr int32_t block_size = WARP_SIZE;
@@ -896,7 +1004,8 @@ void rotate_activation_mxfp6_quant_k(aiter_tensor_t& out,
                                                     1.0f,
                                                     sequence,
                                                     heads,
-                                                    tiles);
+                                                    tiles,
+                                                    mean_ptr);
     });
 }
 
@@ -973,7 +1082,8 @@ void rotate_activation_mxfp4_quant(aiter_tensor_t& out,
 
 void rotate_activation_mxfp4_quant_k(aiter_tensor_t& out,
                                      aiter_tensor_t& scale,
-                                     const aiter_tensor_t& input)
+                                     const aiter_tensor_t& input,
+                                     const aiter_tensor_t& mean)
 {
     constexpr int64_t tile = kHeadDim;
     AITER_CHECK(input.dim() == 4, "input must be BSHD");
@@ -985,6 +1095,9 @@ void rotate_activation_mxfp4_quant_k(aiter_tensor_t& out,
                 "out must have one padded 8192-byte tile per batch and head");
     // aiter_tensor_t has no as_strided; pass the logical (unpadded) out numel.
     check_inputs<64>(out, scale, input, static_cast<int64_t>(input.numel()) / 2);
+    int32_t mean_heads     = 1;
+    int32_t mean_seq_heads = static_cast<int32_t>(input.numel() / kHeadDim);
+    const float* mean_ptr  = check_k_mean(mean, input, mean_heads, mean_seq_heads);
     AITER_DISPATCH_FLOATING16_TYPES_rmTorch(input.dtype(), "rotate_activation_mxfp4_quant_k", [&] {
         using DTYPE_I = typename aiter::hip2opus<scalar_t>::type;
         constexpr int32_t block_size = WARP_SIZE;
@@ -1002,7 +1115,8 @@ void rotate_activation_mxfp4_quant_k(aiter_tensor_t& out,
                                                     1.0f,
                                                     sequence,
                                                     heads,
-                                                    tiles);
+                                                    tiles,
+                                                    mean_ptr);
     });
 }
 
diff --git a/csrc/py_itfs_cu/asm_mha_v4_fwd.cu b/csrc/py_itfs_cu/asm_mha_v4_fwd.cu
index 4c1b6d96da3..bfd765a27ca 100644
--- a/csrc/py_itfs_cu/asm_mha_v4_fwd.cu
+++ b/csrc/py_itfs_cu/asm_mha_v4_fwd.cu
@@ -73,6 +73,15 @@ struct MhaV4Recipe
 };
 
 constexpr int64_t kHeadDim = 128;
+constexpr int64_t kGfx950StagedSparseMaxKvTiles = 8192;
+
+bool uses_gfx950_staged_sparse_lut(const MhaV4Recipe& recipe)
+{
+    return recipe.q_format == format_id(AttentionFormat::Bf16) ||
+           recipe.q_format == format_id(AttentionFormat::Fp8E4M3) ||
+           recipe.q_format == format_id(AttentionFormat::Fp4E2M1) ||
+           recipe.q_format == format_id(AttentionFormat::Int8);
+}
 
 struct PointerSlot
 {
@@ -817,7 +826,9 @@ void fmha_v4_fwd(const at::Tensor& q,
                  int64_t q_scale_mode,
                  int64_t k_scale_mode,
                  int64_t v_scale_mode,
-                 double softmax_scale)
+                 double softmax_scale,
+                 std::optional<at::Tensor> seqlens_k,
+                 std::optional<at::Tensor> lse)
 {
     const MhaV4Recipe recipe{q_format,
                              k_format,
@@ -859,6 +870,52 @@ void fmha_v4_fwd(const at::Tensor& q,
                            shapes.gqa_ratio,
                            softmax_scale);
 
+    // Per-batch key lengths are optional: the kernels read this slot only when it is non-null, so a
+    // dense launch leaves it zero rather than selecting a different code object.
+    if(seqlens_k.has_value())
+    {
+        TORCH_CHECK(seqlens_k->is_cuda() && seqlens_k->device() == q.device(),
+                    "MHA v4 seqlens_k must be a GPU tensor on the same device as Q");
+        TORCH_CHECK(seqlens_k->scalar_type() == at::kInt,
+                    "MHA v4 seqlens_k must be int32, got ",
+                    seqlens_k->scalar_type());
+        TORCH_CHECK(seqlens_k->is_contiguous(), "MHA v4 seqlens_k must be contiguous");
+        TORCH_CHECK(seqlens_k->numel() >= shapes.batch,
+                    "MHA v4 seqlens_k needs one entry per batch: got ",
+                    seqlens_k->numel(),
+                    " for batch ",
+                    shapes.batch);
+        // Contents stay the caller's contract, as they do for cu_seqlens on the varlen path:
+        // they live on the device, so bounding them here would synchronize every launch.
+        args.ptr_kseq.value = seqlens_k->data_ptr();
+    }
+
+    // LSE is opt-in the same way: the kernels branch on s_lse and skip the store when it is zero,
+    // so a launch without it leaves the reserved slots at zero and keeps the same code object.
+    if(lse.has_value())
+    {
+        TORCH_CHECK(lse->is_cuda() && lse->device() == q.device(),
+                    "MHA v4 lse must be a GPU tensor on the same device as Q");
+        TORCH_CHECK(lse->scalar_type() == at::kFloat,
+                    "MHA v4 lse must be float32, got ",
+                    lse->scalar_type());
+        TORCH_CHECK(lse->is_contiguous(), "MHA v4 lse must be contiguous");
+        TORCH_CHECK(lse->dim() == 3 && lse->size(0) == shapes.batch &&
+                        lse->size(1) == shapes.nhead_q && lse->size(2) == shapes.seqlen_q,
+                    "MHA v4 lse must be [batch, nhead_q, seqlen_q] = [",
+                    shapes.batch,
+                    ", ",
+                    shapes.nhead_q,
+                    ", ",
+                    shapes.seqlen_q,
+                    "]");
+        // The kernel derives the batch stride as q_head_num * s_lse_Hs, which only holds for a
+        // contiguous [batch, head, seqlen_q] buffer.
+        args.ptr_lse.value  = lse->data_ptr();
+        args.s_lse.value    = 1;
+        args.s_lse_Hs.value = byte_stride(*lse, 1, "LSE head stride");
+    }
+
     static SynchronizedCache<std::string, AiterAsmKernel> kernels;
     const std::string cache_key = arch + "|" + cfg.knl_name + "|" + cfg.co_name;
     auto& kernel                = kernels.get_or_create(
@@ -920,6 +977,14 @@ void fmha_v4_fwd_sparse(const at::Tensor& q,
 
     const int64_t q_tiles  = (shapes.seqlen_q + cfg.ts_qo - 1) / cfg.ts_qo;
     const int64_t kv_tiles = (shapes.seqlen_k + cfg.ts_kv - 1) / cfg.ts_kv;
+    if(arch == "gfx950" && uses_gfx950_staged_sparse_lut(recipe))
+    {
+        TORCH_CHECK(kv_tiles <= kGfx950StagedSparseMaxKvTiles,
+                    "gfx950 LDS-staged sparse MHA v4 supports at most ",
+                    kGfx950StagedSparseMaxKvTiles,
+                    " KV tiles, got ",
+                    kv_tiles);
+    }
     const int64_t lut_rows = shapes.batch * shapes.nhead_q * q_tiles;
     TORCH_CHECK(kv_block_indices.is_cuda() && lut_start.is_cuda() && lut_count.is_cuda(),
                 "LUT tensors must be GPU tensors");
diff --git a/csrc/pybind/mha_v4_fwd_pybind.cu b/csrc/pybind/mha_v4_fwd_pybind.cu
index c22005d781e..a1fed0d6f98 100644
--- a/csrc/pybind/mha_v4_fwd_pybind.cu
+++ b/csrc/pybind/mha_v4_fwd_pybind.cu
@@ -28,7 +28,9 @@ PYBIND11_MODULE(AITER_EXTENSION_NAME, m)
           py::arg("q_scale_mode"),
           py::arg("k_scale_mode"),
           py::arg("v_scale_mode"),
-          py::arg("softmax_scale"));
+          py::arg("softmax_scale"),
+          py::arg("seqlens_k") = std::nullopt,
+          py::arg("lse")       = std::nullopt);
     m.def("fmha_v4_fwd_sparse",
           &aiter::torch_itfs::fmha_v4_fwd_sparse,
           py::arg("q"),
diff --git a/csrc/pybind/mha_v4_quant_pybind.cu b/csrc/pybind/mha_v4_quant_pybind.cu
index 24923434fa0..ed45931a0ce 100644
--- a/csrc/pybind/mha_v4_quant_pybind.cu
+++ b/csrc/pybind/mha_v4_quant_pybind.cu
@@ -19,13 +19,16 @@ PYBIND11_MODULE(AITER_EXTENSION_NAME, m)
     m.def("rotate_activation_hd128",
           &aiter::torch_itfs::rotate_activation_hd128,
           py::arg("out"),
-          py::arg("input"));
+          py::arg("input"),
+          py::arg("mean"),
+          py::arg("partial_amax"));
     m.def("rotate_activation_mxfp8_quant",
           &aiter::torch_itfs::rotate_activation_mxfp8_quant,
           py::arg("out"),
           py::arg("scale"),
           py::arg("input"),
-          py::arg("multiplier"));
+          py::arg("multiplier"),
+          py::arg("mean"));
     m.def("rotate_activation_mxfp6_quant",
           &aiter::torch_itfs::rotate_activation_mxfp6_quant,
           py::arg("out"),
@@ -36,7 +39,8 @@ PYBIND11_MODULE(AITER_EXTENSION_NAME, m)
           &aiter::torch_itfs::rotate_activation_mxfp6_quant_k,
           py::arg("out"),
           py::arg("scale"),
-          py::arg("input"));
+          py::arg("input"),
+          py::arg("mean"));
     m.def("_quantize_v_mxfp6_fp6_p_hip",
           &aiter::torch_itfs::quantize_v_mxfp6_fp6_p,
           py::arg("out"),
@@ -52,7 +56,8 @@ PYBIND11_MODULE(AITER_EXTENSION_NAME, m)
           &aiter::torch_itfs::rotate_activation_mxfp4_quant_k,
           py::arg("out"),
           py::arg("scale"),
-          py::arg("input"));
+          py::arg("input"),
+          py::arg("mean"));
     m.def("_quantize_v_mxfp4_fp6_p_hip",
           &aiter::torch_itfs::quantize_v_mxfp4_fp6_p,
           py::arg("out"),
diff --git a/hsa/gfx950/fmha_v4_fwd/fmha_v4_fwd.csv b/hsa/gfx950/fmha_v4_fwd/fmha_v4_fwd.csv
index 2ba1d3b2002..9516b843175 100644
--- a/hsa/gfx950/fmha_v4_fwd/fmha_v4_fwd.csv
+++ b/hsa/gfx950/fmha_v4_fwd/fmha_v4_fwd.csv
@@ -1,5 +1,6 @@
 # mode: 0 = dense 3-D grid, 1 = sorted block-sparse (flat grid + 752-byte kernarg, work_table @ 0x2D0)
 q_format,k_format,v_format,v_pack,o_format,q_scale_mode,k_scale_mode,v_scale_mode,o_scale_mode,hdim_q,hdim_v,mask,mode,ts_qo,ts_kv,knl_name,co_name
+# DENSE kernels
 2,2,2,0,2,0,0,0,0,128,128,0,0,256,64,_ZN5aiter19fmha_fwd_hd128_bf16E,fwd_hd128_bf16.co
 2,2,3,0,2,0,0,1,0,128,128,0,0,256,64,_ZN5aiter22fmha_fwd_hd128_bf16fp8E,fwd_hd128_bf16fp8.co
 10,10,3,0,2,1,1,1,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_i8fp8_gfx950E,fwd_hd128_i8fp8.co
@@ -9,13 +10,18 @@ q_format,k_format,v_format,v_pack,o_format,q_scale_mode,k_scale_mode,v_scale_mod
 7,7,3,0,2,5,5,4,0,128,128,0,0,256,128,_ZN5aiter26fmha_fwd_hd128_f6f8_gfx950E,fwd_hd128_f6f8.co
 7,7,7,1,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter27fmha_fwd_hd128_mxfp6_gfx950E,fwd_hd128_mxfp6.co
 7,7,9,1,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_f6f4_gfx950E,fwd_hd128_f6f4.co
-9,9,9,1,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_f4f4_gfx950E,fwd_hd128_f4f4.co
-9,9,9,0,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_mxfp4_gfx950E,fwd_hd128_mxfp4.co
+9,9,9,1,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_mxfp4_gfx950E,fwd_hd128_mxfp4.co
+# f4f4 disabled: mxfp4 now covers this row and claims the same v_pack=1 FP6-P V order.
+#9,9,9,1,2,5,5,5,0,128,128,0,0,256,128,_ZN5aiter28fmha_fwd_hd128_f4f4_gfx950E,fwd_hd128_f4f4.co
+# SPARSE kernels
+2,2,2,0,2,0,0,0,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_bf16_sparse_gfx950E,fwd_hd128_bf16_sparse.co
+2,2,3,0,2,0,0,1,0,128,128,0,1,256,128,_ZN5aiter36fmha_fwd_hd128_bf16fp8_sparse_gfx950E,fwd_hd128_bf16fp8_sparse.co
 10,10,3,0,2,1,1,1,0,128,128,0,1,256,128,_ZN5aiter34fmha_fwd_hd128_i8fp8_sparse_gfx950E,fwd_hd128_i8fp8_sparse.co
 3,3,3,0,2,5,5,1,0,128,128,0,1,256,128,_ZN5aiter34fmha_fwd_hd128_mxfp8_sparse_gfx950E,fwd_hd128_mxfp8_sparse.co
 3,3,3,0,2,1,1,1,0,128,128,0,1,256,128,_ZN5aiter32fmha_fwd_hd128_fp8_sparse_gfx950E,fwd_hd128_fp8_sparse.co
-3,3,7,0,2,1,1,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f8f6_sparse_gfx950E,fwd_hd128_f8f6_sparse.co
+3,3,7,1,2,1,1,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f8f6_sparse_gfx950E,fwd_hd128_f8f6_sparse.co
 7,7,3,0,2,5,5,4,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f6f8_sparse_gfx950E,fwd_hd128_f6f8_sparse.co
-7,7,9,0,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f6f4_sparse_gfx950E,fwd_hd128_f6f4_sparse.co
-9,9,3,0,2,5,5,4,0,128,128,0,1,256,128,_ZN5aiter34fmha_fwd_hd128_mxfp4_sparse_gfx950E,fwd_hd128_mxfp4_sparse.co
-9,9,9,0,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f4f4_sparse_gfx950E,fwd_hd128_f4f4_sparse.co
\ No newline at end of file
+7,7,7,1,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter34fmha_fwd_hd128_mxfp6_sparse_gfx950E,fwd_hd128_mxfp6_sparse.co
+7,7,9,1,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f6f4_sparse_gfx950E,fwd_hd128_f6f4_sparse.co
+9,9,9,1,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter34fmha_fwd_hd128_mxfp4_sparse_gfx950E,fwd_hd128_mxfp4_sparse.co
+#9,9,9,1,2,5,5,5,0,128,128,0,1,256,128,_ZN5aiter33fmha_fwd_hd128_f4f4_sparse_gfx950E,fwd_hd128_f4f4_sparse.co
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16.co
index fcdcb62cb5f..08391f92643 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16_sparse.co
new file mode 100755
index 00000000000..7947b5032f2
Binary files /dev/null and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8.co
index 03573106f1f..591a107fb2e 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8_sparse.co
new file mode 100755
index 00000000000..b21fe3c1906
Binary files /dev/null and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_bf16fp8_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4.co
index 3920ef61ce6..9495615f47c 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4_sparse.co
index 379428e2568..3b21fea476f 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f4f4_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4.co
index ba9e1673c02..a7c93ee00ea 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4_sparse.co
index 242045aedc4..8d6d5240abd 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f4_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8.co
old mode 100755
new mode 100644
index cfbcd4f6e87..7ff33e72547
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8_sparse.co
old mode 100755
new mode 100644
index c0adcd5a857..afb246327af
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f6f8_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6.co
old mode 100644
new mode 100755
index 8601f65acff..ce7420d8fd4
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6_sparse.co
index a261f71ee5e..c7fe23193b4 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_f8f6_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8.co
index d7464901ed2..816656b82c5 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8_sparse.co
index c53eca66f4c..60755f48e9d 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_fp8_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8.co
index ec31b513436..b04793c4362 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8_sparse.co
index 748215a7d88..5258ad3fbc4 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_i8fp8_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4.co
old mode 100755
new mode 100644
index 73af44903b5..c9f59d3a70d
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4_sparse.co
old mode 100755
new mode 100644
index be177684e12..f2e432ee706
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp4_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6.co
old mode 100755
new mode 100644
index 8f595055305..bad06df259e
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6_sparse.co
new file mode 100644
index 00000000000..221f1107229
Binary files /dev/null and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp6_sparse.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8.co
index 50b0844a8bc..805e5d5020d 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8.co differ
diff --git a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8_sparse.co b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8_sparse.co
index a869586e030..040bc7cc417 100755
Binary files a/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8_sparse.co and b/hsa/gfx950/fmha_v4_fwd/fwd_hd128_mxfp8_sparse.co differ
diff --git a/op_tests/op_benchmarks/triton/bench_sage.py b/op_tests/op_benchmarks/triton/bench_sage.py
index b9f0862616a..670f4165e2f 100644
--- a/op_tests/op_benchmarks/triton/bench_sage.py
+++ b/op_tests/op_benchmarks/triton/bench_sage.py
@@ -5,6 +5,7 @@
 import glob
 import json
 import logging
+import math
 import os
 import re
 import sys
@@ -45,7 +46,6 @@
     quantize_mxfp8_k,
     quantize_mxfp8_q,
     quantize_v_fp8,
-    quantize_v_mxfp4,
     quantize_v_mxfp4_fp6_p,
     quantize_v_mxfp6,
     quantize_v_mxfp6_fp6_p,
@@ -86,7 +86,7 @@ def _production_quantize_mxfp4(query, key, value, softmax_scale):
     q_fp4, q_scale = quantize_mxfp4_q(query, mha_v4_q_multiplier(softmax_scale))
     k_raw, k_scale = quantize_mxfp4_k(key)
     k_fp4 = mxfp4_k_view(k_raw, k_scale)
-    v_raw, v_scale = quantize_v_mxfp4(value)
+    v_raw, v_scale = quantize_v_mxfp4_fp6_p(value)
     v_fp4 = mxfp4_v_view(v_raw, v_scale, value.shape[1])
     return q_fp4, q_scale, k_fp4, k_scale, v_fp4, v_scale
 
@@ -98,12 +98,11 @@ def _production_quantize_mxfp8(query, key, value, softmax_scale):
     return q_fp8, k_fp8, v_fp8, q_scale, k_scale, v_scale
 
 
-def _production_quantize_f4f4(query, key, value, softmax_scale, fp6_p=True):
+def _production_quantize_f4f4(query, key, value, softmax_scale):
     q_fp4, q_scale = quantize_mxfp4_q(query, mha_v4_q_multiplier(softmax_scale))
     k_raw, k_scale = quantize_mxfp4_k(key)
     k_fp4 = mxfp4_k_view(k_raw, k_scale)
-    quantize_v = quantize_v_mxfp4_fp6_p if fp6_p else quantize_v_mxfp4
-    v_raw, v_scale = quantize_v(value)
+    v_raw, v_scale = quantize_v_mxfp4_fp6_p(value)
     v_fp4 = mxfp4_v_view(v_raw, v_scale, value.shape[1])
     return q_fp4, q_scale, k_fp4, k_scale, v_fp4, v_scale
 
@@ -124,8 +123,7 @@ def _production_quantize_mxfp6(
         return q_fp6, q_scale, k_fp6, k_scale, *quantize_v(value)
     if v_format != AttentionFormat.MXFP4:
         raise ValueError(f"unsupported MXFP6 Q/K V format: {v_format!r}")
-    quantize_v = quantize_v_mxfp4_fp6_p if fp6_p else quantize_v_mxfp4
-    v_raw, v_scale = quantize_v(value)
+    v_raw, v_scale = quantize_v_mxfp4_fp6_p(value)
     v_quantized = mxfp4_v_view(v_raw, v_scale, value.shape[1])
     return q_fp6, q_scale, k_fp6, k_scale, v_quantized, v_scale
 
@@ -180,8 +178,10 @@ def _mha_v4_spec(
     ),
     "fav3_fp8": KernelSpec((1.0, 1.0, 1.0), quantized=True, uses_hadamard=True),
     "aiter_bf16": KernelSpec((2.0, 2.0, 2.0), include_in_all=True),
-    "mha4_bf16": _mha_v4_spec((2.0, 2.0, 2.0), quantized=False),
-    "mha4_bf16fp8": _mha_v4_spec((2.0, 2.0, 1.0)),
+    "mha4_bf16": _mha_v4_spec(
+        (2.0, 2.0, 2.0), quantized=False, supports_block_sparse=True
+    ),
+    "mha4_bf16fp8": _mha_v4_spec((2.0, 2.0, 1.0), supports_block_sparse=True),
     "mha4_i8fp8": _mha_v4_spec((1.0, 1.0, 1.0), supports_block_sparse=True),
     "mha4_mxfp8": _mha_v4_spec(
         (1.0, 1.0, 1.0),
@@ -205,6 +205,7 @@ def _mha_v4_spec(
     ),
     "mha4_mxfp6": _mha_v4_spec(
         (0.75, 0.75, 0.75),
+        supports_block_sparse=True,
         uses_hadamard=True,
     ),
     "mha4_f6f4": _mha_v4_spec(
@@ -295,7 +296,9 @@ def _generate_transformer_qkv(
     d_head_v: int,
     device: str,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    # Realistic LLM activations: RMS-norm + per-channel log-normal scales + shared low-rank Q/K component + V outlier dims/tokens. Returns fp32 q/k/v.
+    # Realistic LLM activations: RMS-norm, per-channel log-normal scales, a per-token shared Q/K
+    # term, and V outlier dims/tokens. The shared term gives self-affinity, not a token-common
+    # mode: that stays near 0.02 here, against 0.13-0.42 on dumped video models. Returns fp32.
     q = torch.randn((batch, hq, sq, d_head), device=device, dtype=torch.float32)
     k = torch.randn((batch, hk, sk, d_head), device=device, dtype=torch.float32)
     v = torch.randn((batch, hk, sk, d_head_v), device=device, dtype=torch.float32)
@@ -339,6 +342,97 @@ def _generate_transformer_qkv(
     return q, k, v
 
 
+# Calibrated against dumped Wan, HunyuanVideo 1.5 and Flux2 attention, which agree closely despite
+# differing architectures: shared components 0.37-0.57 (Q), 0.37-0.47 (K), 0.31-0.36 (V),
+# cancellation 2.4-2.8, logit spread 2.0-2.6, amax/RMS near 7 on Q/K and 8.8-25.2 on V.
+_DIFFUSION_RHO_Q = 0.45
+_DIFFUSION_RHO_K = 0.40
+_DIFFUSION_RHO_V = 0.34
+_DIFFUSION_CHANNEL_SIGMA = 0.16
+_DIFFUSION_LOGIT_STD = 2.2
+_DIFFUSION_V_OUTLIER_DIM_GAIN = 2.0
+_DIFFUSION_V_OUTLIER_TOKEN_GAIN = 4.0
+
+
+def _coherent_rows(batch, heads, seq, dim, rho, sigma, device):
+    """Rows sharing a per-head direction, so the token-common mode is `rho` by construction.
+
+    The per-channel scale then spreads channel magnitudes, which is what lifts amax above the RMS
+    even though each channel on its own stays near-Gaussian, as the dumps do.
+    """
+    shared = torch.randn((1, heads, 1, dim), device=device, dtype=torch.float32)
+    shared *= math.sqrt(dim) / shared.norm(dim=-1, keepdim=True)
+    noise = torch.randn((batch, heads, seq, dim), device=device, dtype=torch.float32)
+    rows = rho * shared + math.sqrt(1.0 - rho * rho) * noise
+    channel_scale = (
+        torch.randn((1, heads, 1, dim), device=device, dtype=torch.float32)
+        .mul(sigma)
+        .exp()
+    )
+    return rows * channel_scale
+
+
+def _generate_diffusion_qkv(
+    batch: int,
+    hq: int,
+    hk: int,
+    sq: int,
+    sk: int,
+    d_head: int,
+    d_head_v: int,
+    device: str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Diffusion-transformer activations. Returns fp32 q/k/v.
+
+    The defining difference from LLM activations is that Q, K and V all carry a large component
+    shared across tokens. For V that is what holds cancellation near 2.6 rather than the 8-40 that
+    iid V produces, because averaging many keys cancels the noise and leaves the shared part.
+    """
+    q = _coherent_rows(
+        batch, hq, sq, d_head, _DIFFUSION_RHO_Q, _DIFFUSION_CHANNEL_SIGMA, device
+    )
+    k = _coherent_rows(
+        batch, hk, sk, d_head, _DIFFUSION_RHO_K, _DIFFUSION_CHANNEL_SIGMA, device
+    )
+    v = _coherent_rows(
+        batch, hk, sk, d_head_v, _DIFFUSION_RHO_V, _DIFFUSION_CHANNEL_SIGMA, device
+    )
+
+    outlier_dims = torch.randperm(d_head_v, device=device)[: max(1, d_head_v // 16)]
+    v[..., outlier_dims] *= _DIFFUSION_V_OUTLIER_DIM_GAIN
+    outlier_tokens = torch.randperm(sk, device=device)[: max(1, sk // 256)]
+    v[:, :, outlier_tokens, :] *= _DIFFUSION_V_OUTLIER_TOKEN_GAIN
+
+    # Scaling Q moves the logit spread, and so the diffuseness, without touching any of the
+    # common modes, which are ratios.
+    probe = torch.arange(0, sq, max(1, sq // 64), device=device)[:64]
+    measured = ((q[0, 0, probe] @ k[0, 0].T) * (d_head**-0.5)).std(-1).median()
+    q *= _DIFFUSION_LOGIT_STD / measured.clamp_min(1e-9)
+    return q, k, v
+
+
+_NAMED_DISTRIBUTIONS = (
+    "zero",
+    "normal",
+    "diffusion",
+    "padded",
+    "transformer",
+    "sink",
+    "underflow",
+    "latepeak",
+    "maxstair",
+    "kcommon",
+)
+
+
+def _input_distribution(value: str) -> str:
+    if value in _NAMED_DISTRIBUTIONS or value.startswith("dump:"):
+        return value
+    raise argparse.ArgumentTypeError(
+        f"invalid choice: {value!r} (choose from {', '.join(_NAMED_DISTRIBUTIONS)}, or dump:PATH)"
+    )
+
+
 def generate_test_tensors(
     batch: int,
     hq: int,
@@ -352,6 +446,38 @@ def generate_test_tensors(
     distribution: str,
     hadamard_rotate: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # "dump:<path>": replay Q/K/V captured from a real model. Every synthetic distribution here is
+    # ~10x too optimistic even when matched on common mode, diffuseness, cancellation, logit
+    # spread and amax/RMS, so replaying sidesteps having to know what actually drives real error.
+    if distribution.startswith("dump:"):
+        sample = torch.load(
+            distribution[len("dump:") :], map_location=device, weights_only=False
+        )
+        # Captures are BSHD; the bench works in BHSD.
+        q, k, v = (
+            sample[name].transpose(1, 2)
+            for name in ("query_bshd", "key_bshd", "value_bshd")
+        )
+        have = (
+            q.shape[0],
+            q.shape[1],
+            k.shape[1],
+            q.shape[2],
+            k.shape[2],
+            q.shape[3],
+            v.shape[3],
+        )
+        want = (batch, hq, hk, sq, sk, d_head, d_head_v)
+        if any(h < w for h, w in zip(have, want)):
+            raise ValueError(
+                f"dump is {have} (batch, hq, hk, sq, sk, d, dv) and cannot cover the "
+                f"requested {want}; ask for a shape it contains"
+            )
+        q = q[:batch, :hq, :sq, :d_head]
+        k = k[:batch, :hk, :sk, :d_head]
+        v = v[:batch, :hk, :sk, :d_head_v]
+        return tuple(t.contiguous().to(dtype) for t in (q, k, v))
+
     # "zero": all-zero Q/K/V for degenerate-input and quantization smoke tests.
     if distribution == "zero":
         q = torch.zeros((batch, hq, sq, d_head), device=device, dtype=dtype)
@@ -381,25 +507,9 @@ def generate_test_tensors(
         return q.to(dtype), k.to(dtype), v.to(dtype)
 
     if distribution == "underflow":
-        # Reproduces the fp8 underflow tile-skip regression on the microbench.
-        #
-        # A strong "hotspot" in the first KV tile (keys [0:128]) establishes a
-        # high frozen softmax max. Every later KV tile then sits far below the
-        # e4m3 round-to-zero floor (~2^-11 of the max ≈ 7.62 nats), so the
-        # kernel's underflow tile-skip path fires on those tiles. A per-query-row
-        # jitter on the hotspot strength makes the all-underflow condition
-        # row-dependent, so the two anti-phase co-resident wave groups (which own
-        # different query-row blocks) disagree on which tiles to skip. The
-        # shared-VALU lockstep barrier then eats the saving while the extra
-        # underflow compare is still paid on every no-mask tile -> net slowdown,
-        # matching the observed model-level result.
-        #
-        # Tunables (env):
-        #   AITER_UNDERFLOW_GAP    max hotspot logit in nats (default 16.0)
-        #   AITER_UNDERFLOW_JITTER per-row hotspot factor ~ U[jitter, 1]
-        #                          (default 0.4 -> asymmetric/realistic regression;
-        #                           set 1.0 for the symmetric best-case where every
-        #                           later tile underflows for both partner waves)
+        # Reproduces the fp8 underflow tile-skip regression: a hotspot in the first KV tile sets a
+        # high frozen max, so later tiles fall below the e4m3 round-to-zero floor and the skip path
+        # fires. Per-row jitter makes partner waves disagree on which tiles to skip, eating it.
         gap = float(os.environ.get("AITER_UNDERFLOW_GAP", "16.0"))
         jitter = float(os.environ.get("AITER_UNDERFLOW_JITTER", "0.4"))
         jitter = min(max(jitter, 0.0), 1.0)
@@ -439,15 +549,9 @@ def generate_test_tensors(
         return q.to(dtype), k.to(dtype), v.to(dtype)
 
     if distribution == "latepeak":
-        # ADVERSARIAL TRIPWIRE for the frozen-max rollback (added 2026-06-14 after the black-video
-        # regression). Mirrors `underflow` but places the high-norm "attention sink" hotspot in the
-        # LAST KV tile instead of the first. With a frozen-max rollback that seeds from tile 0, the
-        # seed is LOW and the late hotspot's logit blows far past it -> the Schraudolph u32-cvt
-        # saturates to 0xFFFFFFFF (NaN bits) -> corrupt P -> NaN/black. The exact (proper running
-        # max) path is immune (S - m_new <= 0 always). Random transformer/normal/underflow never
-        # produce a late-tile outlier, so this is the structured input cosine-on-random missed.
-        #   AITER_LATEPEAK_GAP : late-peak logit in nats (default 40.0 -> well past the cvt
-        #                        saturation at scale_log2e*(S-seed) > 128 for 1/sqrt(d) scaling)
+        # Tripwire for the frozen-max rollback: the sink sits in the LAST KV tile, so a seed taken
+        # from tile 0 is low and the late logit saturates the Schraudolph u32 cvt to NaN bits.
+        # Random distributions never produce a late-tile outlier, so cosine-on-random missed it.
         gap = float(
             os.environ.get(
                 "AITER_LATEPEAK_GAP", os.environ.get("AITER_LATESINK_GAP", "40.0")
@@ -475,19 +579,9 @@ def generate_test_tensors(
         return q.to(dtype), k.to(dtype), v.to(dtype)
 
     if distribution == "maxstair":
-        # Frozen-max rollback stress: make every 128-token KV tile establish a new row max, with
-        # alternating 128-query-row groups above and below the rollback threshold. This keeps
-        # freeze-max active for half the rows while stressing rollback in the other half.
-        # Build in the post-Hadamard domain so MXFP6 quantization preserves each tile step.
-        #
-        # Tunables (env):
-        #   AITER_MAXSTAIR_STEP       score increase in kernel log2 units per KV tile
-        #                              (default 12.0; rollback threshold is about 8.87).
-        #   AITER_MAXSTAIR_LOW_FACTOR alternate 128-query-row groups between factors 1 and this
-        #                              value (default 0.5: half the rows roll back). Set 1.0 for
-        #                              the less representative every-row/every-tile rollback mode.
-        #                              Values below ~0.74 with the default step keep low groups
-        #                              below the threshold and stress paired-wave disagreement.
+        # Frozen-max rollback stress: every 128-token KV tile sets a new row max, with alternating
+        # query-row groups above and below the rollback threshold (about 8.87 in kernel log2
+        # units). Built post-Hadamard so MXFP6 quantization preserves each tile step.
         step = float(os.environ.get("AITER_MAXSTAIR_STEP", "12.0"))
         low_factor = float(os.environ.get("AITER_MAXSTAIR_LOW_FACTOR", "0.5"))
         if step <= 0:
@@ -552,11 +646,44 @@ def generate_test_tensors(
         )
         return q.to(dtype), k.to(dtype), v.to(dtype)
 
-    if distribution != "transformer":
+    if distribution == "kcommon":
+        # Diffusion activations plus a large shared-K direction. Softmax is shift-invariant in it
+        # so output is unchanged, but per-tensor Q/K quantization noise scales with |q.k| rather
+        # than its spread. Real models reach 0.47; this goes further to stress the K-smoothing gate.
+        q, k, v = _generate_diffusion_qkv(
+            batch, hq, hk, sq, sk, d_head, d_head_v, device
+        )
+        direction = torch.randn((hk, d_head), device=device, dtype=torch.float32)
+        direction /= direction.norm(dim=-1, keepdim=True)
+        k += 16.0 * k.norm(dim=-1).mean() * direction[None, :, None, :]
+        return q.to(dtype), k.to(dtype), v.to(dtype)
+
+    if distribution == "padded":
+        # Ulysses head padding: a shard can hold (batch, head) slices entirely zero in Q, K and V.
+        # Z-Image rank 7 has 4 of 8 heads empty. Found a NaN in the per-channel FP8 V quantizer,
+        # which divides by a zero amax.
+        q, k, v = _generate_diffusion_qkv(
+            batch, hq, hk, sq, sk, d_head, d_head_v, device
+        )
+        q[:, hq - max(1, hq // 4) :] = 0
+        empty_kv = max(1, hk // 4)
+        k[:, hk - empty_kv :] = 0
+        v[:, hk - empty_kv :] = 0
+        return q.to(dtype), k.to(dtype), v.to(dtype)
+
+    if distribution == "transformer":
+        # LLM activations. Kept as its own distribution rather than folded into "diffusion":
+        # MHA v4 runs LLM workloads too, and the two look nothing alike, the token-common mode
+        # staying near 0.02 here against 0.31-0.57 measured on diffusion models.
+        q, k, v = _generate_transformer_qkv(
+            batch, hq, hk, sq, sk, d_head, d_head_v, device
+        )
+        return q.to(dtype), k.to(dtype), v.to(dtype)
+
+    if distribution != "diffusion":
         raise ValueError(f"Unsupported input distribution: {distribution}")
 
-    # "transformer": realistic LLM activation statistics (see _generate_transformer_qkv).
-    q, k, v = _generate_transformer_qkv(batch, hq, hk, sq, sk, d_head, d_head_v, device)
+    q, k, v = _generate_diffusion_qkv(batch, hq, hk, sq, sk, d_head, d_head_v, device)
     return q.to(dtype), k.to(dtype), v.to(dtype)
 
 
@@ -1253,12 +1380,9 @@ def launch_mha_v4(*tensors, **kwargs):
 
         q_quantized, q_descale = quantize_fp8_rotated(q_bshd)
         k_quantized, k_descale = quantize_fp8_rotated(k_bshd)
-        if block_lut is None:
-            v_quantized, v_descale = quantize_v_mxfp6_fp6_p(v_bshd)
-            v_pack = AttentionPack.V_FOR_FP6_P
-        else:
-            v_quantized, v_descale = quantize_v_mxfp6(v_bshd)
-            v_pack = AttentionPack.DEFAULT
+        # Both modes run the FP6 P pack, so both need V staged in the matching layout.
+        v_quantized, v_descale = quantize_v_mxfp6_fp6_p(v_bshd)
+        v_pack = AttentionPack.V_FOR_FP6_P
         return lambda: launch_mha_v4_packed(
             q_quantized,
             k_quantized,
@@ -1318,17 +1442,12 @@ def launch_mha_v4(*tensors, **kwargs):
             raise ValueError(f"{args.kernel} does not support --qsmooth")
 
         is_f4f4 = args.kernel == "mha4_f4f4"
-        sparse_mxfp4 = args.kernel == "mha4_mxfp4" and block_lut is not None
-        v_format = fp8_format if sparse_mxfp4 else AttentionFormat.MXFP4
+        v_format = AttentionFormat.MXFP4
         scale_modes = scale_modes_for_formats(
             AttentionFormat.MXFP4, AttentionFormat.MXFP4, v_format
         )
-        use_dense_p_pack = block_lut is None
-        v_pack = (
-            AttentionPack.V_FOR_FP6_P
-            if is_f4f4 and use_dense_p_pack
-            else AttentionPack.DEFAULT
-        )
+        # Both all-MXFP4 kernels consume an FP6 P operand, so both need the FP6-P V order.
+        v_pack = AttentionPack.V_FOR_FP6_P
 
         def _quantize_mxfp4():
             quant_q, quant_k = q_bshd, k_bshd
@@ -1336,16 +1455,8 @@ def _quantize_mxfp4():
                 quant_q, quant_k = cancel_internal_qk_rotation(quant_q, quant_k)
             if is_f4f4:
                 return _production_quantize_f4f4(
-                    quant_q, quant_k, v_bshd, softmax_scale, use_dense_p_pack
+                    quant_q, quant_k, v_bshd, softmax_scale
                 )
-            if sparse_mxfp4:
-                q_fp4, q_scale = quantize_mxfp4_q(
-                    quant_q, mha_v4_q_multiplier(softmax_scale)
-                )
-                k_raw, k_scale = quantize_mxfp4_k(quant_k)
-                k_fp4 = mxfp4_k_view(k_raw, k_scale)
-                v_fp8, v_scale = quantize_v_fp8(v_bshd)
-                return q_fp4, q_scale, k_fp4, k_scale, v_fp8, v_scale
             return _production_quantize_mxfp4(quant_q, quant_k, v_bshd, softmax_scale)
 
         def _kernel_mxfp4(q_fp4, q_descale, k_fp4, k_descale, v_quantized, v_descale):
@@ -1389,6 +1500,9 @@ def _kernel_mxfp4(q_fp4, q_descale, k_fp4, k_descale, v_quantized, v_descale):
                 "MXFP6 Hadamard preprocessing requires block_r=128 "
                 "and does not support --qsmooth"
             )
+        # Every MXFP6-Q row ships an FP6-P object in both modes. The quantizer and the launcher
+        # must agree on this.
+        uses_fp6_p = is_mxfp6 or is_f6f4
 
         def _quantize_mxfp6():
             quant_q, quant_k = q_bshd, k_bshd
@@ -1404,7 +1518,7 @@ def _quantize_mxfp6():
                     if is_f6f4
                     else AttentionFormat.MXFP6 if is_mxfp6 else None
                 ),
-                fp6_p=(is_f6f4 or is_mxfp6) and block_lut is None,
+                fp6_p=uses_fp6_p,
             )
 
         v_format = (
@@ -1429,9 +1543,7 @@ def _kernel_mxfp6(q_fp6, q_descale, k_fp6, k_descale, v_quantized, v_descale):
                 v_format,
                 *scale_modes,
                 v_pack=(
-                    AttentionPack.V_FOR_FP6_P
-                    if (is_f6f4 or is_mxfp6) and block_lut is None
-                    else AttentionPack.DEFAULT
+                    AttentionPack.V_FOR_FP6_P if uses_fp6_p else AttentionPack.DEFAULT
                 ),
                 softmax_scale=softmax_scale,
             )
@@ -1494,7 +1606,13 @@ def compute_accuracy_metrics(
 
 
 def fp8_max_diff_percentage(args: argparse.Namespace) -> float:
-    if args.input_distribution in ("transformer", "sink"):
+    # The coherent distributions and the LLM one are both harder on FP8 than iid inputs.
+    if args.input_distribution in (
+        "transformer",
+        "sink",
+        "diffusion",
+        "kcommon",
+    ) or str(args.input_distribution).startswith("dump:"):
         return 2.0
     return 0.5
 
@@ -1544,10 +1662,9 @@ def make_reference_output(
     )
     ref = args.ref
 
-    # The torch reference (attention_ref) materializes a full [b, hq, sq, sk] fp32 scores tensor and
-    # softmaxes it; past ~32 GiB that path becomes numerically UNRELIABLE -- the cosine collapses
-    # even for a correct kernel (measured ~0.45 at sq=sk=75520, while sq=sk=32768 ~21 GiB is fine).
-    # Warn and point the user at --ref aiter_bf16, which streams the scores and stays accurate.
+    # attention_ref materializes a full [b, hq, sq, sk] fp32 scores tensor; past ~32 GiB it becomes
+    # numerically unreliable and cosine collapses even for a correct kernel (~0.45 at sq=sk=75520,
+    # while 32768 is fine). Point the user at --ref aiter_bf16, which streams the scores instead.
     if ref == "torch":
         b_, sq_, hq_, _ = q_bshd.shape
         sk_ = k_bshd.shape[1]
@@ -1629,13 +1746,9 @@ def benchmark_payload_bytes(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
-    *,
-    sparse: bool = False,
 ) -> tuple[float, float, float]:
     if args.e2e:
         return float(q.element_size()), float(k.element_size()), float(v.element_size())
-    if args.kernel == "mha4_mxfp4" and sparse:
-        return 0.5, 0.5, 1.0
     return KERNEL_SPECS[args.kernel].payload_bytes
 
 
@@ -1697,7 +1810,7 @@ def benchmark_single_case(
 
     mem = compute_memory_bytes(
         shape,
-        *benchmark_payload_bytes(args, q, k, v, sparse=block_lut is not None),
+        *benchmark_payload_bytes(args, q, k, v),
     )
 
     sparse_flops = None
@@ -2251,24 +2364,22 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--causal", action="store_true", help="Enable causal attention")
     parser.add_argument(
         "--input-distribution",
-        type=str,
-        default="transformer",
-        choices=[
-            "zero",
-            "normal",
-            "transformer",
-            "sink",
-            "underflow",
-            "latepeak",
-            "maxstair",
-        ],
+        type=_input_distribution,
+        default="diffusion",
+        metavar="{zero,normal,diffusion,padded,transformer,sink,underflow,latepeak,maxstair,kcommon,dump:PATH}",
         help=(
-            "Distribution used for generated Q/K/V tensors. 'zero' sets all Q/K/V values "
-            "to zero; 'sink' is a realistic "
-            "StreamingLLM attention sink pattern; 'underflow'/'latepeak' are "
-            "adversarial fp8 tile-skip / frozen-max rollback regression tripwires; "
-            "'maxstair' raises the max every KV tile and triggers rollback for alternating "
-            "query-row groups."
+            "Distribution used for generated Q/K/V tensors. 'diffusion' (default) is calibrated "
+            "to dumped Wan, HunyuanVideo 1.5 and Flux2 attention; 'transformer' is the LLM "
+            "activation model, which differs mainly in carrying almost no token-common "
+            "component; 'padded' is diffusion with empty Ulysses-padding heads, which real "
+            "sharded models produce and which no other distribution covers; 'zero' "
+            "sets all Q/K/V values to zero; 'sink' is a realistic StreamingLLM attention sink "
+            "pattern; 'underflow'/'latepeak' are adversarial fp8 tile-skip / frozen-max rollback "
+            "regression tripwires; 'maxstair' raises the max every KV tile and triggers rollback "
+            "for alternating query-row groups; 'kcommon' adds a shared K direction far past what "
+            "real models reach, which inflates quantization noise without changing any softmax "
+            "statistic; 'dump:PATH' replays Q/K/V captured from a real model, which no synthetic "
+            "distribution here reproduces the quantization error of."
         ),
     )
     parser.add_argument(
diff --git a/op_tests/test_mha_v4.py b/op_tests/test_mha_v4.py
index 342558efc86..ab43ee46948 100644
--- a/op_tests/test_mha_v4.py
+++ b/op_tests/test_mha_v4.py
@@ -4,10 +4,9 @@
 import argparse
 import itertools
 import math
-import os
 import subprocess
 import sys
-from typing import NamedTuple
+import textwrap
 
 import pandas as pd
 import pytest
@@ -15,20 +14,19 @@
 import torch._dynamo
 
 import aiter
+import aiter.ops.mha_v4 as mha_v4_module
 from aiter import dtypes
-from aiter.jit.core import AITER_ROOT_DIR
 from aiter.jit.utils.chip_info import get_gfx
 from aiter.ops.mha_v4 import (
     AttentionFormat,
     AttentionPack,
     AttentionScaleMode,
+    _k_mean,
     _RawRecipeKind,
     _resolve_raw_recipe,
     mha_v4,
     mha_v4_kv_tile,
-    mha_v4_mxfp8,
     mha_v4_packed,
-    mha_v4_sparse_work_table,
     native_fp8_format,
     scale_modes_for_formats,
 )
@@ -54,14 +52,12 @@
     quantize_mxfp6_q,
     quantize_mxfp8_k,
     quantize_mxfp8_q,
-    quantize_v_mxfp4,
     quantize_v_mxfp4_fp6_p,
     quantize_v_mxfp6,
     quantize_v_mxfp6_fp6_p,
     rotate_activation_hd128,
     rotate_activation_mxfp6_quant,
 )
-from aiter.ops.triton.attention.utils import block_attn_mask_to_ragged_lut
 from aiter.ops.triton.quant.mxfp6_fmha_pack import (
     _v_direct_kvtab,
     fp6_k_raw_buffer_sizes,
@@ -98,69 +94,6 @@ def _rotate_hd128_reference(value):
     return (rotated / 128**0.5).to(value.dtype)
 
 
-def _reference_mxfp4_v(value):
-    batch, sequence, heads, _ = value.shape
-    padded_sequence = fp4_v_padded_sequence(sequence)
-    tiles = padded_sequence // 128
-    padded = torch.nn.functional.pad(
-        value.float(), (0, 0, 0, 0, 0, padded_sequence - sequence)
-    )
-    padded = padded.permute(0, 2, 1, 3)
-
-    column = torch.arange(64, device=value.device)
-    lane = column % 32
-    permutation = 4 * (lane // 8) + 16 * ((lane // 4) % 2) + lane % 4
-    tau64 = 32 * (column // 32) + permutation
-    kperm = torch.empty(64, dtype=torch.long, device=value.device)
-    kperm[tau64] = column
-
-    raw = torch.zeros(
-        fp4_v_raw_buffer_size(batch, sequence, heads),
-        dtype=torch.uint8,
-        device=value.device,
-    )
-    payload = raw[:-64].view(batch, heads, tiles * 8192)
-    scale = torch.empty(
-        (batch, heads, tiles * 512), dtype=torch.uint8, device=value.device
-    )
-    for tile in range(tiles):
-        for channel_block in range(4):
-            for token_half in range(2):
-                unit = 2 * channel_block + token_half
-                tokens = tile * 128 + token_half * 64 + kperm
-                channels = slice(channel_block * 32, (channel_block + 1) * 32)
-                block = padded[:, :, tokens, channels]
-                exponents = []
-                normalized = torch.empty_like(block)
-                for token_block in range(2):
-                    columns = slice(token_block * 32, (token_block + 1) * 32)
-                    amax = block[:, :, columns].abs().amax(dim=2)
-                    exponent = torch.ceil(
-                        torch.log2(torch.clamp_min(amax, 1e-12) / 6.0)
-                    )
-                    exponents.append(exponent)
-                    normalized[:, :, columns] = block[:, :, columns] / torch.exp2(
-                        exponent[:, :, None]
-                    )
-
-                code = _e2m1_code_ties_low(normalized)
-                packed = code[..., 0::2] | (code[..., 1::2] << 4)
-                payload[
-                    :, :, tile * 8192 + unit * 1024 : tile * 8192 + (unit + 1) * 1024
-                ] = packed.flatten(2)
-
-                scale_base = tile * 512 + token_half * 256
-                for token_block, exponent in enumerate(exponents):
-                    encoded = (exponent + 127).clamp(0, 255).to(torch.uint8)
-                    for pair in range(16):
-                        offset = (
-                            scale_base + token_block * 128 + 8 * pair + channel_block
-                        )
-                        scale[:, :, offset] = encoded[:, :, 2 * pair]
-                        scale[:, :, offset + 4] = encoded[:, :, 2 * pair + 1]
-    return raw, scale
-
-
 @pytest.fixture(autouse=True)
 def isolate_dynamo_cache():
     """Keep each test's ``torch.compile`` behaviour independent of the tests that ran before it.
@@ -254,14 +187,22 @@ def test_mha_v4_bf16fp8_scale_recipe():
             AttentionFormat.MXFP6,
             True,
             _RawRecipeKind.FP8,
-            AttentionPack.DEFAULT,
+            AttentionPack.V_FOR_FP6_P,
         ),
+        # All-MXFP4 consumes FP6 probabilities against MXFP4 V, so both modes need the repacked V.
         (
             AttentionFormat.MXFP4,
             AttentionFormat.MXFP4,
             False,
             _RawRecipeKind.MXFP4,
-            AttentionPack.DEFAULT,
+            AttentionPack.V_FOR_FP6_P,
+        ),
+        (
+            AttentionFormat.MXFP4,
+            AttentionFormat.MXFP4,
+            True,
+            _RawRecipeKind.MXFP4,
+            AttentionPack.V_FOR_FP6_P,
         ),
         (
             AttentionFormat.MXFP6,
@@ -275,7 +216,15 @@ def test_mha_v4_bf16fp8_scale_recipe():
             AttentionFormat.MXFP4,
             True,
             _RawRecipeKind.MXFP6,
-            AttentionPack.DEFAULT,
+            AttentionPack.V_FOR_FP6_P,
+        ),
+        # MXFP6 Q/K/V ships an FP6-P object in both modes, so sparse keeps the repacked V.
+        (
+            AttentionFormat.MXFP6,
+            AttentionFormat.MXFP6,
+            True,
+            _RawRecipeKind.MXFP6,
+            AttentionPack.V_FOR_FP6_P,
         ),
     ],
 )
@@ -295,31 +244,30 @@ def test_mha_v4_resolves_raw_recipe(q_format, v_format, sparse, kind, v_pack):
 
 
 @pytest.mark.parametrize(
-    ("q_format", "v_format", "message"),
+    ("v_format", "kind"),
     [
         (
             AttentionFormat.BF16,
-            AttentionFormat.BF16,
-            "does not have a BF16 manifest row",
+            _RawRecipeKind.BF16,
         ),
         (
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP6,
-            "MXFP6 Q/K/V",
+            native_fp8_format(),
+            _RawRecipeKind.BF16_FP8,
         ),
     ],
 )
-def test_mha_v4_rejects_unavailable_sparse_recipe(q_format, v_format, message):
-    with pytest.raises(NotImplementedError, match=message):
-        _resolve_raw_recipe(
-            q_format,
-            q_format,
-            v_format,
-            None,
-            None,
-            None,
-            sparse=True,
-        )
+def test_mha_v4_resolves_bf16_sparse_recipe(v_format, kind):
+    recipe = _resolve_raw_recipe(
+        AttentionFormat.BF16,
+        AttentionFormat.BF16,
+        v_format,
+        None,
+        None,
+        None,
+        sparse=True,
+    )
+    assert recipe.kind == kind
+    assert recipe.v_pack == AttentionPack.DEFAULT
 
 
 @pytest.mark.parametrize("sparse", [False, True])
@@ -581,6 +529,41 @@ def test_mha_v4_rotated_fp8_quantization_rejects_noncontiguous_input():
         quantize_fp8_rotated(value)
 
 
+@pytest.mark.parametrize(
+    "quantize",
+    [
+        "quantize_fp8_rotated",
+        "quantize_mxfp8_k",
+        "quantize_mxfp4_k",
+        "quantize_mxfp6_k",
+    ],
+)
+def test_mha_v4_k_quantizers_reject_an_off_device_mean(quantize):
+    """A host mean reached the kernel as a device pointer and faulted the GPU.
+
+    Out of process because these checks abort rather than raise, as every AITER_CHECK in that
+    translation unit does. All four quantizers share one validator, so all four are covered.
+    """
+    source = textwrap.dedent(f"""
+        import torch
+        from aiter.ops.mha_v4_quant import {quantize} as quantize
+
+        value = torch.randn((1, 128, 2, 128), device="cuda", dtype=torch.bfloat16)
+        quantize(value, torch.zeros((1, 2, 128), dtype=torch.float32))
+        torch.cuda.synchronize()
+        """)
+    finished = subprocess.run(
+        [sys.executable, "-c", source],
+        capture_output=True,
+        text=True,
+        timeout=1800,
+        check=False,
+    )
+
+    assert finished.returncode != 0
+    assert "same GPU as input" in finished.stderr, finished.stderr[-2000:]
+
+
 @pytest.mark.skipif(
     get_gfx() not in ("gfx942", "gfx950"),
     reason="gfx942/gfx950 activation rotation",
@@ -606,8 +589,9 @@ def test_mha_v4_fp8_raw_recipe_matches_rotated_packed():
     v = torch.randn_like(q)
     fp8_format = native_fp8_format()
 
+    # mha_v4 smooths K before quantizing; packed callers pass the same mean themselves.
     q_quantized, q_descale = quantize_fp8_rotated(q)
-    k_quantized, k_descale = quantize_fp8_rotated(k)
+    k_quantized, k_descale = quantize_fp8_rotated(k, _k_mean(k, _RawRecipeKind.FP8))
     v_quantized, v_descale = quantize_fp8(v)
     expected = mha_v4_packed(
         q_quantized,
@@ -632,6 +616,117 @@ def test_mha_v4_fp8_raw_recipe_matches_rotated_packed():
     assert torch.equal(compiled, expected)
 
 
+@pytest.mark.skipif(
+    get_gfx() not in ("gfx942", "gfx950"),
+    reason="gfx942/gfx950 FP8 recipe validation",
+)
+@pytest.mark.parametrize("recipe", ["fp8", "mxfp8", "mxfp4"])
+def test_mha_v4_quantized_tolerates_k_common_mode(recipe):
+    """A direction shared by every key must not cost accuracy.
+
+    Softmax is shift invariant in such a component, so the reference barely moves; only the
+    quantizers care. Without the mean subtraction the error grows several-fold here, so this
+    is the tripwire for silently dropping it.
+    """
+    torch.manual_seed(17)
+    q = torch.randn((1, 1024, 4, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(q)
+    base_k = torch.randn_like(q)
+    direction = torch.randn((1, 1, 4, 128), device="cuda", dtype=torch.bfloat16)
+
+    if recipe == "fp8":
+        formats = (native_fp8_format(),) * 3
+        scale_modes = {}
+    elif recipe == "mxfp8":
+        # MXFP8 is selected by the E8M0 scale modes, not by a distinct format.
+        formats = (native_fp8_format(),) * 3
+        scale_modes = {
+            "q_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+            "k_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+            "v_scale_mode": AttentionScaleMode.F32_PER_TENSOR,
+        }
+    else:
+        formats = (AttentionFormat.MXFP4,) * 3
+        scale_modes = {}
+
+    errors = []
+    for common in (0.0, 16.0):
+        k = base_k + common * direction
+        reference = _dense_reference(q.float(), k.float(), v.float(), k.shape[1])
+        out = mha_v4(q, k, v, *formats, **scale_modes)
+        errors.append(
+            ((out.float() - reference).norm() / reference.norm()).item(),
+        )
+
+    assert errors[1] < 1.5 * errors[0], (
+        f"{recipe} degrades under a shared K direction: "
+        f"{errors[0]:.4f} -> {errors[1]:.4f}; is K smoothing still applied?"
+    )
+
+
+@pytest.mark.skipif(
+    get_gfx() not in ("gfx942", "gfx950"),
+    reason="gfx942/gfx950 FP8 recipe validation",
+)
+@pytest.mark.parametrize("recipe", ["fp8", "mxfp8", "mxfp4"])
+def test_mha_v4_lse_survives_k_common_mode(recipe):
+    """K smoothing must not leak into the exported LSE.
+
+    Smoothing runs the kernel against k - k_mean, shifting every score by the per-query constant
+    q @ k_mean. Output cannot see it -- a shift shared by all keys cancels in the softmax -- so
+    only the LSE carries it. Chunked consumers weight each chunk by exp(lse) and derive their own
+    k_mean per chunk, so an uncorrected shift mis-weights the chunks; ring attention lost a third
+    of its output norm this way while every output test stayed green.
+    """
+    torch.manual_seed(17)
+    q = torch.randn((1, 1024, 4, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(q)
+    base_k = torch.randn_like(q)
+    direction = torch.randn((1, 1, 4, 128), device="cuda", dtype=torch.bfloat16)
+
+    if recipe == "fp8":
+        formats = (native_fp8_format(),) * 3
+        scale_modes = {}
+    elif recipe == "mxfp8":
+        formats = (native_fp8_format(),) * 3
+        scale_modes = {
+            "q_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+            "k_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+            "v_scale_mode": AttentionScaleMode.F32_PER_TENSOR,
+        }
+    else:
+        formats = (AttentionFormat.MXFP4,) * 3
+        scale_modes = {}
+
+    softmax_scale = q.shape[-1] ** -0.5
+    errors = []
+    for common in (0.0, 16.0):
+        k = base_k + common * direction
+        scores = (
+            q.float().permute(0, 2, 1, 3) @ k.float().permute(0, 2, 3, 1)
+        ) * softmax_scale
+        reference = torch.logsumexp(scores, dim=-1)
+        _, lse = mha_v4(
+            q,
+            k,
+            v,
+            *formats,
+            return_lse=True,
+            softmax_scale=softmax_scale,
+            **scale_modes,
+        )
+        errors.append((lse.float() - reference).abs().max().item())
+
+    assert (
+        _k_mean(base_k + 16.0 * direction, _RawRecipeKind.FP8).abs().max() > 0
+    ), "K smoothing did not engage, so this case cannot detect the leak"
+    assert errors[1] < errors[0] + 0.5, (
+        f"{recipe} LSE degrades under a shared K direction: "
+        f"{errors[0]:.4f} -> {errors[1]:.4f} nats; is the k_mean shift still "
+        "added back into the LSE?"
+    )
+
+
 @pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP8 quantization")
 @pytest.mark.parametrize("case", ["random", "zero", "powers", "extreme"])
 def test_mha_v4_mxfp8_q_matches_unfused_pipeline(case):
@@ -741,37 +836,6 @@ def test_mha_v4_mxfp4_v_backing_storage_covers_logical_view(batch, sequence, hea
     assert max_logical_offset < raw_size
 
 
-@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP4 V validation")
-@pytest.mark.parametrize("sequence", [1, 63, 64, 127, 128, 129, 255, 257])
-def test_mha_v4_mxfp4_v_pack_matches_reference(sequence):
-    torch.manual_seed(sequence)
-    value = torch.randn((2, sequence, 3, 128), device="cuda", dtype=torch.bfloat16)
-    raw, scale = quantize_v_mxfp4(value)
-    raw_again, scale_again = quantize_v_mxfp4(value)
-    expected_raw, expected_scale = _reference_mxfp4_v(value)
-
-    assert raw.shape == (fp4_v_raw_buffer_size(2, sequence, 3),)
-    assert scale.shape == (2, 3, ((sequence + 127) // 128) * 512)
-    assert raw.dtype == scale.dtype == torch.uint8
-    assert torch.equal(raw, expected_raw)
-    assert torch.equal(scale, expected_scale)
-    assert torch.equal(raw, raw_again)
-    assert torch.equal(scale, scale_again)
-    # The HIP producer replaced a Triton packer; keep the retired one as a second oracle.
-    triton_raw, triton_scale = pack_v_mxfp4_colmajor_raw(value)
-    assert torch.equal(raw, triton_raw)
-    assert torch.equal(scale, triton_scale)
-    assert torch.count_nonzero(raw[-64:]) == 0
-    logical = mxfp4_v_view(raw, scale, sequence)
-    assert logical.shape == value.shape
-    assert logical.stride() == (
-        3 * fp4_v_padded_sequence(sequence) * 64,
-        64,
-        fp4_v_padded_sequence(sequence) * 64,
-        1,
-    )
-
-
 @pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP4 V validation")
 @pytest.mark.parametrize("sequence", [1, 63, 64, 127, 128, 129, 255, 257])
 def test_mha_v4_mxfp4_fp6_p_pack_matches_permuted_canonical(sequence):
@@ -918,8 +982,8 @@ def test_mha_v4_q_scale_backing_storage_covers_query_tile(quantize, sequence):
 @pytest.mark.parametrize("sequence", [1, 128, 129, 257, 512])
 @pytest.mark.parametrize(
     "quantize",
-    [quantize_v_mxfp4, quantize_v_mxfp4_fp6_p],
-    ids=["canonical", "fp6_p"],
+    [quantize_v_mxfp4_fp6_p],
+    ids=["fp6_p"],
 )
 def test_mha_v4_mxfp4_v_scale_backing_storage_covers_lookahead_tiles(
     quantize, sequence
@@ -946,7 +1010,9 @@ def test_mha_v4_mxfp4_v_scale_backing_storage_covers_lookahead_tiles(
 
 def test_mha_v4_rejects_unsupported_contracts():
     q = torch.empty((1, 128, 2, 128), device="cuda", dtype=torch.bfloat16)
-    with pytest.raises(NotImplementedError, match="do not produce LSE"):
+    # Dense LSE is supported for every shipped format row; the sorted-sparse path is not.
+    block_mask = torch.ones((1, 2, 1, 1), device="cuda", dtype=torch.bool)
+    with pytest.raises(NotImplementedError, match="sorted-sparse path"):
         mha_v4(
             q,
             q,
@@ -955,6 +1021,7 @@ def test_mha_v4_rejects_unsupported_contracts():
             AttentionFormat.FP8,
             AttentionFormat.FP8,
             return_lse=True,
+            block_mask=block_mask,
         )
     with pytest.raises(ValueError, match="matching Q and K formats"):
         mha_v4(
@@ -967,6 +1034,143 @@ def test_mha_v4_rejects_unsupported_contracts():
         )
 
 
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MHA v4 validation")
+@pytest.mark.parametrize(
+    ("q_format", "v_format", "scale_modes"),
+    [
+        (AttentionFormat.BF16, AttentionFormat.BF16, None),
+        (AttentionFormat.BF16, AttentionFormat.FP8, None),
+        (AttentionFormat.INT8, AttentionFormat.FP8, None),
+        (AttentionFormat.FP8, AttentionFormat.FP8, None),
+        (AttentionFormat.FP8, AttentionFormat.MXFP6, None),
+        (AttentionFormat.MXFP4, AttentionFormat.MXFP4, None),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.FP8, None),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP6, None),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP4, None),
+        (
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            (
+                AttentionScaleMode.E8M0_PER_1X32,
+                AttentionScaleMode.E8M0_PER_1X32,
+                AttentionScaleMode.F32_PER_TENSOR,
+            ),
+        ),
+    ],
+)
+def test_mha_v4_dense_lse_matches_reference(q_format, v_format, scale_modes):
+    """A wrong LSE does not fail output validation, so it needs its own reference check.
+
+    The quantized rows carry a small systematic bias from the approximate exp2; the bound here is
+    wide enough to pass that but far tighter than the failure modes it guards, which are a missing
+    log2(L) term (error grows like ln(Sk)) and an unapplied P-pack divisor (a constant ln2 or 2ln2).
+    """
+    torch.manual_seed(31)
+    q = torch.randn((1, 512, 5, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    softmax_scale = 128**-0.5
+    qsm, ksm, vsm = scale_modes if scale_modes else (None, None, None)
+    kwargs = {
+        "softmax_scale": softmax_scale,
+        "q_scale_mode": qsm,
+        "k_scale_mode": ksm,
+        "v_scale_mode": vsm,
+    }
+
+    out_only = mha_v4(q, k, v, q_format, q_format, v_format, **kwargs)
+    out, lse = mha_v4(q, k, v, q_format, q_format, v_format, return_lse=True, **kwargs)
+
+    scores = torch.matmul(
+        q.float().permute(0, 2, 1, 3),
+        k.float().permute(0, 2, 1, 3).transpose(-1, -2),
+    )
+    reference = torch.logsumexp(scores * softmax_scale, dim=-1)
+
+    assert lse.shape == reference.shape
+    assert lse.dtype == torch.float32
+    assert torch.isfinite(lse).all()
+    # Asking for the LSE must not perturb O: the epilogue runs either way and only the store is
+    # gated, so its scratch registers must not touch anything O still needs.
+    assert torch.equal(out_only, out)
+    error = (lse - reference).abs()
+    assert error.max().item() < 0.25, error.max().item()
+    assert error.mean().item() < 0.05, error.mean().item()
+
+
+@pytest.mark.parametrize(
+    ("q_format", "v_format", "max_spread"),
+    [
+        (AttentionFormat.BF16, AttentionFormat.BF16, 0.05),
+        (AttentionFormat.BF16, AttentionFormat.FP8, 0.05),
+        (AttentionFormat.INT8, AttentionFormat.FP8, 0.05),
+        (AttentionFormat.FP8, AttentionFormat.FP8, 0.75),
+        (AttentionFormat.FP8, AttentionFormat.MXFP6, 0.75),
+        (AttentionFormat.MXFP4, AttentionFormat.MXFP4, 0.35),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.FP8, 0.35),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP6, 0.35),
+        (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP4, 0.35),
+    ],
+)
+def test_mha_v4_lse_bias_is_constant_across_key_chunks(q_format, v_format, max_spread):
+    """Ring weights each chunk by exp(lse), so only a bias identical across chunks cancels.
+
+    The absolute bias is allowed to be nonzero and the test above already bounds it; what this
+    one pins is that it does not move from chunk to chunk. Keys escalate along the sequence on
+    purpose: random ones keep every row diffuse, the frozen-max conversion gate never trips, and
+    a rollback missing from the exported max is then invisible. With this input the recipes span
+    0.00 to 0.24 nats, while that defect measured 4.55.
+    """
+    torch.manual_seed(31)
+    batch, sequence, heads, head_dim, chunks = 1, 2048, 5, 128, 4
+    query = torch.randn(
+        (batch, sequence, heads, head_dim), device="cuda", dtype=torch.bfloat16
+    )
+    value = torch.randn_like(query)
+    ramp = torch.linspace(1.0, 6.0, sequence, device="cuda", dtype=torch.float32)
+    key = (torch.randn_like(query).float() * ramp.view(1, -1, 1, 1)).to(torch.bfloat16)
+    softmax_scale = head_dim**-0.5
+
+    span = sequence // chunks
+    biases = []
+    for start in range(0, sequence, span):
+        key_chunk = key[:, start : start + span]
+        _, lse = mha_v4(
+            query,
+            key_chunk,
+            value[:, start : start + span],
+            q_format,
+            q_format,
+            v_format,
+            softmax_scale=softmax_scale,
+            return_lse=True,
+        )
+        scores = query.float().permute(0, 2, 1, 3) @ key_chunk.float().permute(
+            0, 2, 3, 1
+        )
+        reference = torch.logsumexp(scores * softmax_scale, dim=-1)
+        biases.append((lse.float() - reference).mean().item())
+
+    spread = max(biases) - min(biases)
+    assert spread < max_spread, f"per-chunk bias {biases} spans {spread:.4f} nats"
+
+
+def test_mha_v4_lse_is_gated_off_gfx950(monkeypatch):
+    """gfx942 carries the epilogue, but its exported value has never been measured.
+
+    A wrong LSE passes every output test, because O never reads it, so presence of the store is
+    not evidence of correctness. Drop the gate once MI300 is compared against torch.logsumexp.
+    """
+    monkeypatch.setattr(mha_v4_module, "get_gfx", lambda: "gfx942")
+    q = torch.randn((1, 128, 4, 128), device="cuda", dtype=torch.bfloat16)
+    formats = (AttentionFormat.BF16, AttentionFormat.BF16, AttentionFormat.BF16)
+
+    with pytest.raises(NotImplementedError, match="not validated on gfx942"):
+        mha_v4(q, q, q, *formats, return_lse=True)
+
+    assert torch.isfinite(mha_v4(q, q, q, *formats)).all()
+
+
 @pytest.mark.parametrize(
     "q_format",
     [
@@ -1160,7 +1364,7 @@ def test_mha_v4_packed_rejects_unbacked_mx_scales():
 
     mxfp4_q, mxfp4_q_scale = quantize_mxfp4_q(value, 1.0)
     mxfp4_raw, mxfp4_k_scale = quantize_mxfp4_k(value)
-    mxfp4_v_raw, mxfp4_v_scale = quantize_v_mxfp4(value)
+    mxfp4_v_raw, mxfp4_v_scale = quantize_v_mxfp4_fp6_p(value)
     with pytest.raises(RuntimeError, match="speculative tile gather"):
         mha_v4_packed(
             mxfp4_q,
@@ -1212,6 +1416,54 @@ def test_mha_v4_zero_inputs_are_finite(q_format, v_format):
     assert torch.isfinite(out).all()
 
 
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MHA v4 validation")
+@pytest.mark.parametrize(
+    ("q_format", "v_format", "scale_modes"),
+    [
+        (AttentionFormat.BF16, AttentionFormat.BF16, {}),
+        (AttentionFormat.BF16, AttentionFormat.FP8, {}),
+        (AttentionFormat.INT8, AttentionFormat.FP8, {}),
+        (AttentionFormat.FP8, AttentionFormat.FP8, {}),
+        (
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            {
+                "q_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+                "k_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+                "v_scale_mode": AttentionScaleMode.F32_PER_TENSOR,
+            },
+        ),
+        (AttentionFormat.FP8, AttentionFormat.MXFP6, {}),
+        (AttentionFormat.MXFP6, AttentionFormat.FP8, {}),
+        (AttentionFormat.MXFP6, AttentionFormat.MXFP6, {}),
+        (AttentionFormat.MXFP6, AttentionFormat.MXFP4, {}),
+        (AttentionFormat.MXFP4, AttentionFormat.MXFP4, {}),
+    ],
+)
+def test_mha_v4_empty_heads_are_finite(q_format, v_format, scale_modes):
+    """Sequence-parallel head padding leaves whole (batch, head) slices zero.
+
+    Only reachable with live heads alongside them: f6f8 returned NaN for exactly half its output
+    here, because its per-channel FP8 V quantizer divided by a zero amax while the populated heads
+    kept the tensor looking healthy.
+    """
+    torch.manual_seed(0)
+    q = torch.randn((1, 512, 4, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    for tensor in (q, k, v):
+        tensor[:, :, 2:] = 0
+
+    out = mha_v4(q, k, v, q_format, q_format, v_format, **scale_modes)
+    torch.cuda.synchronize()
+    assert torch.isfinite(out).all(), (
+        f"{q_format.name}/{v_format.name} produced "
+        f"{int(torch.isnan(out).sum())} NaN on padded heads"
+    )
+    assert torch.count_nonzero(out[:, :, 2:]) == 0
+    assert torch.count_nonzero(out[:, :, :2]) > 0
+
+
 @pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 BF16-FP8 validation")
 @pytest.mark.parametrize(("sequence_q", "sequence_k"), [(129, 257), (257, 193)])
 def test_mha_v4_bf16fp8_matches_dequantized_reference(sequence_q, sequence_k):
@@ -1433,32 +1685,6 @@ def test_mha_v4_raw_compile_parity(q_format, v_format):
     assert churn.numel() == 16 * 1024 * 1024
 
 
-@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP8 validation")
-def test_mha_v4_mxfp8_deprecated_alias_matches_mha_v4():
-    torch.manual_seed(41)
-    q = torch.randn((1, 257, 5, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn_like(q)
-    v = torch.randn_like(q)
-    fp8_format = native_fp8_format()
-
-    expected = mha_v4(
-        q,
-        k,
-        v,
-        fp8_format,
-        fp8_format,
-        fp8_format,
-        q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-        k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-        v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
-    )
-    with pytest.deprecated_call():
-        actual = mha_v4_mxfp8(q, k, v)
-    torch.cuda.synchronize()
-
-    assert torch.equal(actual, expected)
-
-
 @pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP8 validation")
 def test_mha_v4_raw_mxfp8_compile_parity():
     torch.manual_seed(41)
@@ -1522,1090 +1748,328 @@ def test_mha_v4_raw_mxfp4_v_supports_unaligned_sequence(q_format):
 
     assert torch.equal(eager, compiled)
     assert torch.isfinite(compiled).all()
+    # Determinism and finiteness alone pass on a wrong-but-stable result, so pin the value too.
+    assert _cosine_against_attention(eager, q, k, v) > _MX_UNALIGNED_COSINE
 
 
-def _mha_v4_sparse_co_available() -> bool:
-    gfx = get_gfx()
-    asm_dir = os.environ.get("AITER_ASM_DIR", os.path.join(AITER_ROOT_DIR, "hsa"))
-    if gfx == "gfx942":
-        return os.path.isfile(
-            os.path.join(
-                asm_dir, "gfx942", "fmha_v4_fwd", "MI300", "fwd_hd128_fp8_sparse.co"
-            )
-        )
-    return os.path.isfile(
-        os.path.join(asm_dir, "gfx950", "fmha_v4_fwd", "fwd_hd128_fp8_sparse.co")
-    )
+def _cosine_against_attention(actual, q, k, v):
+    """Cosine of a BSHD MHA v4 result against full-precision attention."""
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2).float(),
+        k.transpose(1, 2).float(),
+        v.transpose(1, 2).float(),
+    ).transpose(1, 2)
+    return torch.nn.functional.cosine_similarity(
+        actual.float().flatten(), reference.flatten(), dim=0
+    ).item()
 
 
-_MHA_V4_SPARSE_ARCH = get_gfx() in ("gfx942", "gfx950")
+# Loose enough to absorb MXFP4 quantization error, which costs about 0.02 on its own at a
+# tile-aligned length, and tight enough that a mishandled partial tile cannot hide.
+_MX_UNALIGNED_COSINE = 0.95
 
+_MX_V_RECIPES = [
+    pytest.param(native_fp8_format(), AttentionFormat.MXFP6, id="f8f6"),
+    pytest.param(AttentionFormat.MXFP6, AttentionFormat.MXFP6, id="mxfp6"),
+    pytest.param(AttentionFormat.MXFP6, AttentionFormat.MXFP4, id="f6f4"),
+    pytest.param(AttentionFormat.MXFP4, AttentionFormat.MXFP4, id="f4f4"),
+]
 
-def test_mha_v4_packed_rejects_partial_lut():
-    dummy = torch.empty(0)
-    with pytest.raises(ValueError, match="all be set or all omitted"):
-        mha_v4_packed(
-            dummy,
-            dummy,
-            dummy,
-            dummy,
-            dummy,
-            dummy,
-            AttentionFormat.INT8,
-            AttentionFormat.INT8,
-            AttentionFormat.FP8,
-            AttentionScaleMode.F32_PER_TENSOR,
-            AttentionScaleMode.F32_PER_TENSOR,
-            AttentionScaleMode.F32_PER_TENSOR,
-            kv_block_indices=dummy,
-        )
 
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX V validation")
+@pytest.mark.parametrize(("qk_format", "v_format"), _MX_V_RECIPES)
+@pytest.mark.parametrize("tail", [1, 8, 32, 64, 96, 127])
+def test_mha_v4_mx_v_partial_kv_tile_matches_attention(qk_format, v_format, tail):
+    """Every MX V recipe must stay correct when the last KV tile is partly filled.
 
-def test_mha_v4_rejects_wrong_block_mask_shape():
-    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16)
-    mask = torch.ones((1, 2, 1, 1), dtype=torch.bool)
-    with pytest.raises(ValueError, match="block_mask must have shape"):
-        mha_v4(
-            q,
-            q,
-            q,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            block_mask=mask,
-        )
+    The KV length is one full 128-token tile plus `tail`, so the only thing varying is
+    partial-tile occupancy. Recipes with a per-tensor FP8 V are flat across `tail`, so any
+    dependence here belongs to the MX V path.
+    """
+    sequence = 128 + tail
+    torch.manual_seed(1234)
+    q = torch.randn((1, sequence, 5, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
 
+    actual = mha_v4(q, k, v, qk_format, qk_format, v_format)
+    torch.cuda.synchronize()
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse schema")
-def test_mha_v4_sparse_schema_mutates_only_out():
-    dense = str(torch.ops.aiter.mha_v4_fwd_launch.default._schema)
-    assert "Tensor kv_block_indices" not in dense
-    assert "Tensor(a6!) out" in dense
-
-    sparse = str(torch.ops.aiter.mha_v4_fwd_sparse_launch.default._schema)
-    assert "Tensor kv_block_indices" in sparse
-    assert "Tensor lut_start" in sparse
-    assert "Tensor lut_count" in sparse
-    assert "Tensor(a6!) out" in sparse
-    assert sparse.endswith("-> ()")
-
-
-def _work_table_counts(total, pattern):
-    """LUT lengths covering the tie structures the ordering has to get right."""
-    if pattern == "uniform":
-        return torch.full((total,), 7, device="cuda", dtype=torch.int32)
-    if pattern == "zeros":
-        return torch.zeros((total,), device="cuda", dtype=torch.int32)
-    if pattern == "random":
-        return torch.randint(0, 64, (total,), device="cuda", dtype=torch.int32)
-    if pattern == "wide_random":
-        return torch.randint(0, 8192, (total,), device="cuda", dtype=torch.int32)
-    if pattern == "two_values":
-        alternating = torch.arange(total, device="cuda") % 3 == 0
-        return torch.where(alternating, 9, 4).to(torch.int32)
-    if pattern == "descending":
-        return torch.arange(total, 0, -1, device="cuda", dtype=torch.int32)
-    return torch.arange(1, total + 1, device="cuda", dtype=torch.int32)
-
-
-def _unpack_work_table(table, nhead, q_tiles):
-    q_idx = (table & 0xFFFF).long()
-    h_idx = ((table >> 16) & 0xFF).long()
-    b_idx = ((table >> 24) & 0xFF).long()
-    return (b_idx * nhead + h_idx) * q_tiles + q_idx
-
-
-# The table has one entry per (batch, head, query tile). 8192 is the point where the builder hands
-# the sort to ATen, so straddle it, and include sizes that are not multiples of a wave or workgroup.
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.parametrize(
-    ("batch", "nhead", "q_tiles"),
-    [
-        (1, 1, 1),
-        (1, 8, 3),
-        (2, 5, 7),
-        (1, 16, 32),
-        (1, 32, 32),
-        (1, 5, 296),
-        (4, 16, 64),
-        (8, 16, 64),
-        (8, 32, 64),
-    ],
-)
-@pytest.mark.parametrize(
-    "pattern",
-    [
-        "uniform",
-        "zeros",
-        "random",
-        "wide_random",
-        "two_values",
-        "descending",
-        "ascending",
-    ],
-)
-def test_mha_v4_sparse_work_table_is_longest_lut_first(batch, nhead, q_tiles, pattern):
-    torch.manual_seed(7)
-    total = batch * nhead * q_tiles
-    counts = _work_table_counts(total, pattern)
+    assert torch.isfinite(actual).all()
+    assert _cosine_against_attention(actual, q, k, v) > _MX_UNALIGNED_COSINE
 
-    table = mha_v4_sparse_work_table(counts, batch, nhead, q_tiles)
-    visited = _unpack_work_table(table, nhead, q_tiles)
 
-    # Every tile exactly once. This is the part a wrong table would turn into a wrong result.
-    assert torch.equal(visited.sort().values, torch.arange(total, device="cuda"))
+def run_torch_mha_v4(q, k, v, softmax_scale):
+    """Compute dense BSHD attention in FP32 for benchmark validation."""
+    scores = (
+        torch.matmul(
+            q.transpose(1, 2).float(), k.transpose(1, 2).float().transpose(-1, -2)
+        )
+        * softmax_scale
+    )
+    return torch.matmul(
+        torch.softmax(scores, dim=-1), v.transpose(1, 2).float()
+    ).transpose(1, 2)
 
-    # Longest LUT first, so no heavy tile straggles behind the rest.
-    ordered = counts[visited]
-    assert bool((ordered[:-1] >= ordered[1:]).all())
 
-    # Ties keep raster order, which is what leaves uniform counts spatially coherent. A stable
-    # reference sort pins the whole permutation, not just the two properties above.
-    expected = torch.argsort(counts, descending=True, stable=True)
-    assert torch.equal(visited, expected)
+@benchmark()
+def benchmark_mha_v4(batch, sequence_q, sequence_k, heads, dtype):
+    """Benchmark every dense MHA v4 recipe against one Torch reference.
 
+    All recipes are timed side by side because they differ only in how Q/K/V are quantized, so
+    the interesting number is what each costs at the same shape and what it costs in accuracy.
+    The quantized rows time their quantizers too, since that is what mha_v4 runs per call.
+    """
+    head_dim = 128
+    softmax_scale = head_dim**-0.5
+    torch.manual_seed(batch + sequence_q + sequence_k + heads)
+    q = torch.randn((batch, sequence_q, heads, head_dim), device="cuda", dtype=dtype)
+    k = torch.randn((batch, sequence_k, heads, head_dim), device="cuda", dtype=dtype)
+    v = torch.randn_like(k)
+    reference = run_torch_mha_v4(q, k, v, softmax_scale)
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.parametrize(
-    "batch,nhead,q_tiles", [(1, 16, 32), (1, 5, 296), (8, 16, 64), (8, 32, 64)]
-)
-def test_mha_v4_sparse_work_table_leaves_uniform_counts_in_raster_order(
-    batch, nhead, q_tiles
-):
-    """Top-k sparsity gives every tile the same LUT length, and that case must not be shuffled."""
-    total = batch * nhead * q_tiles
-    counts = torch.full((total,), 5, device="cuda", dtype=torch.int32)
-
-    table = mha_v4_sparse_work_table(counts, batch, nhead, q_tiles)
-
-    visited = _unpack_work_table(table, nhead, q_tiles)
-    assert torch.equal(visited, torch.arange(total, device="cuda"))
-
-
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-@pytest.mark.parametrize(
-    "launch",
-    [
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                native_fp8_format(),
-                native_fp8_format(),
-                native_fp8_format(),
-                block_mask=mask,
-            ),
-            id="fp8",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                AttentionFormat.INT8,
-                AttentionFormat.INT8,
-                native_fp8_format(),
-                block_mask=mask,
-            ),
-            id="i8fp8",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                native_fp8_format(),
-                native_fp8_format(),
-                native_fp8_format(),
-                block_mask=mask,
-                q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-                k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-                v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
-            ),
-            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-            id="mxfp8",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                native_fp8_format(),
-                native_fp8_format(),
-                AttentionFormat.MXFP6,
-                block_mask=mask,
-            ),
-            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-            id="f8f6",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                AttentionFormat.MXFP6,
-                AttentionFormat.MXFP6,
-                native_fp8_format(),
-                block_mask=mask,
-            ),
-            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-            id="f6f8",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
-                q,
-                k,
-                v,
-                AttentionFormat.MXFP6,
-                AttentionFormat.MXFP6,
-                AttentionFormat.MXFP4,
-                block_mask=mask,
-            ),
-            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-            id="f6f4",
-        ),
-        pytest.param(
-            lambda q, k, v, mask: mha_v4(
+    fp8 = native_fp8_format()
+    recipes = {
+        "bf16": (AttentionFormat.BF16, AttentionFormat.BF16),
+        "bf16fp8": (AttentionFormat.BF16, fp8),
+        "i8fp8": (AttentionFormat.INT8, fp8),
+        "fp8": (fp8, fp8),
+        "f8f6": (fp8, AttentionFormat.MXFP6),
+        "f6f8": (AttentionFormat.MXFP6_E2M3, fp8),
+        "mxfp6": (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP6),
+        "f6f4": (AttentionFormat.MXFP6_E2M3, AttentionFormat.MXFP4),
+        "mxfp4": (AttentionFormat.MXFP4, AttentionFormat.MXFP4),
+    }
+    candidates = {
+        name: (
+            lambda q_format=q_format, v_format=v_format: mha_v4(
                 q,
                 k,
                 v,
-                AttentionFormat.MXFP4,
-                AttentionFormat.MXFP4,
-                AttentionFormat.MXFP4 if mask is None else native_fp8_format(),
-                block_mask=mask,
-            ),
-            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-            id="mxfp4",
-        ),
-    ],
-)
-def test_mha_v4_sparse_all_true_mask_matches_dense(launch):
-    torch.manual_seed(41)
-    q = torch.randn((1, 511, 5, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn((1, 512, 5, 128), device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    mask = torch.ones(
-        (1, 5, 2, 512 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
-    )
-    dense = launch(q, k, v, None)
-    sparse = launch(q, k, v, mask)
-    torch.cuda.synchronize()
-    _assert_sparse_matches_dense(sparse, dense)
-
-
-def _assert_sparse_matches_dense(sparse, dense, message=None):
-    """Compare code objects that use different softmax reduction schedules."""
-    cosine = torch.nn.functional.cosine_similarity(
-        sparse.float().flatten(), dense.float().flatten(), dim=0
-    )
-    assert cosine > 0.99, message
-    assert torch.isfinite(sparse).all()
-
-
-@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_f4f4_sparse_all_true_mask_matches_dense():
-    """Retained FP8-P sparse F4F4 remains close to dense FP6-P on an all-true mask."""
-    torch.manual_seed(41)
-    q = torch.randn((1, 511, 5, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn((1, 512, 5, 128), device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    mask = torch.ones(
-        (1, 5, 2, 512 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
-    )
-    args = (AttentionFormat.MXFP4,) * 3
-    dense = mha_v4(q, k, v, *args)
-    sparse = mha_v4(q, k, v, *args, block_mask=mask)
-    torch.cuda.synchronize()
-
-    _assert_sparse_matches_dense(sparse, dense)
-
-
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.parametrize(
-    "v_format",
-    [
-        pytest.param(AttentionFormat.BF16, id="bf16"),
-        pytest.param(native_fp8_format(), id="bf16fp8"),
-    ],
-)
-def test_mha_v4_sparse_dense_only_formats_reject_block_mask(v_format):
-    q = torch.zeros((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16)
-    mask = torch.ones(
-        (1, 2, 1, 256 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
-    )
-    with pytest.raises(NotImplementedError, match="does not have a BF16 manifest row"):
-        mha_v4(
-            q,
-            q,
-            q,
-            AttentionFormat.BF16,
-            AttentionFormat.BF16,
-            v_format,
-            block_mask=mask,
+                q_format,
+                q_format,
+                v_format,
+                softmax_scale=softmax_scale,
+            )
         )
+        for name, (q_format, v_format) in recipes.items()
+    }
 
-
-@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP6 validation")
-def test_mha_v4_mxfp6_rejects_block_mask():
-    q = torch.zeros((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16)
-    mask = torch.ones(
-        (1, 2, 1, 256 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
-    )
-    with pytest.raises(NotImplementedError, match="MXFP6 Q/K/V"):
-        mha_v4(
-            q,
-            q,
-            q,
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP6,
-            block_mask=mask,
+    flops = 4 * batch * heads * sequence_q * sequence_k * head_dim
+    elements = batch * heads * head_dim * (sequence_q * 2 + sequence_k * 2)
+    nbytes = elements * q.element_size()
+    ret = {"gfx": get_gfx()}
+    for name, candidate in candidates.items():
+        output, us = run_perftest(candidate)
+        # Quantized recipes are expected to land well outside a bitwise bound; the err column is
+        # what carries their accuracy, so compare loosely and let the number speak.
+        err = checkAllclose(
+            reference,
+            output.to(dtypes.fp32),
+            rtol=1e-1,
+            atol=1e-1,
+            msg=f"{name}: dense",
         )
+        ret[f"{name} us"] = us
+        ret[f"{name} TFLOPS"] = flops / us / 1e6
+        ret[f"{name} TB/s"] = nbytes / us / 1e6
+        ret[f"{name} err"] = err
+    return ret
 
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_sparse_block_mask_compiles_without_graph_breaks():
-    """The mask path derives its geometry from host state, which Dynamo cannot trace.
-
-    mha_v4_kv_tile() reads the manifest and get_gfx() shells out to rocminfo, so both sit behind
-    torch_compile_guard. Without that the sparse mask path costs graph breaks per trace and fails
-    under fullgraph, which no other test in this file would notice.
-    """
-    torch.manual_seed(41)
-    kv_tile = mha_v4_kv_tile()
-    q = torch.randn((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn((1, 4 * kv_tile, 2, 128), device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    mask = torch.ones((1, 2, 1, 4), device="cuda", dtype=torch.bool)
-    fp8_format = native_fp8_format()
-
-    def call():
-        return mha_v4(q, k, v, fp8_format, fp8_format, fp8_format, block_mask=mask)
-
-    explained = torch._dynamo.explain(call)()
-    assert explained.break_reasons == [], [
-        str(reason.reason) for reason in explained.break_reasons
-    ]
-
-    eager = call()
-    compiled = torch.compile(call, fullgraph=True)()
-    torch.cuda.synchronize()
-
-    assert torch.equal(eager, compiled)
-
-
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-@pytest.mark.parametrize(
-    ("q_format", "v_format"),
-    [
-        pytest.param(
-            native_fp8_format(),
-            native_fp8_format(),
-            id="fp8",
-        ),
-        pytest.param(
-            AttentionFormat.MXFP4,
-            native_fp8_format(),
-            marks=pytest.mark.skipif(
-                get_gfx() != "gfx950", reason="gfx950 MXFP4 sparse"
-            ),
-            id="mxfp4",
-        ),
-    ],
-)
-def test_mha_v4_sparse_gqa_all_true_mask_matches_repeated_kv(q_format, v_format):
-    torch.manual_seed(41)
-    query_heads = 8
-    kv_heads = 2
-    gqa_ratio = query_heads // kv_heads
-    q = torch.randn((1, 256, query_heads, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn((1, 256, kv_heads, 128), device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    kv_tiles = 256 // mha_v4_kv_tile()
-    mask = torch.ones((1, query_heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
-    k_repeated = k.repeat_interleave(gqa_ratio, dim=2)
-    v_repeated = v.repeat_interleave(gqa_ratio, dim=2)
-
-    gqa_sparse = mha_v4(q, k, v, q_format, q_format, v_format, block_mask=mask)
-    mha_sparse = mha_v4(
-        q,
-        k_repeated,
-        v_repeated,
-        q_format,
-        q_format,
-        v_format,
-        block_mask=mask,
-    )
-    torch.cuda.synchronize()
-
-    assert torch.equal(gqa_sparse, mha_sparse)
-    if q_format != AttentionFormat.MXFP4:
-        gqa_dense = mha_v4(q, k, v, q_format, q_format, v_format)
-        mha_dense = mha_v4(q, k_repeated, v_repeated, q_format, q_format, v_format)
-        assert torch.equal(gqa_dense, mha_dense)
-        _assert_sparse_matches_dense(gqa_sparse, gqa_dense)
-    assert torch.equal(gqa_sparse, mha_sparse)
-
-
-class _Operand(NamedTuple):
-    """A quantized MHA v4 operand and the descale it was produced with."""
-
-    quantized: torch.Tensor
-    descale: torch.Tensor
-
-
-def _sparse_fp8_operands(sequence_k, heads=2, sequence_q=256, batch=1, seed=0):
-    """Quantize once so sparse and reference runs share descales exactly.
+def main():
+    if get_gfx() != "gfx950":
+        aiter.logger.warning(
+            "MHA v4 dense benchmark unsupported on %s; skipping", get_gfx()
+        )
+        return
 
-    Re-quantizing a KV slice would pick a different per-tensor amax, which shifts every
-    value and hides whether the kernel read the KV blocks the LUT named.
-    """
-    torch.manual_seed(seed)
-    q = torch.randn(
-        (batch, sequence_q, heads, 128), device="cuda", dtype=torch.bfloat16
-    )
-    k = torch.randn(
-        (batch, sequence_k, heads, 128), device="cuda", dtype=torch.bfloat16
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawTextHelpFormatter,
+        description="Benchmark the dense MHA v4 recipes",
     )
-    v = torch.randn_like(k)
-    return (
-        _Operand(*quantize_fp8_rotated(q)),
-        _Operand(*quantize_fp8_rotated(k)),
-        _Operand(*quantize_fp8(v)),
+    parser.add_argument("-b", "--batch", type=int, nargs="*", default=[1])
+    parser.add_argument("--sequence-q", type=int, nargs="*", default=[1024])
+    parser.add_argument("--sequence-k", type=int, nargs="*", default=[1024, 4096])
+    parser.add_argument("--heads", type=int, nargs="*", default=[8])
+    parser.add_argument(
+        "-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default=[dtypes.bf16]
     )
+    args = parser.parse_args()
 
-
-def _sparse_fp8_launch(q, k, v, block_mask=None):
-    lut = {}
-    if block_mask is not None:
-        indices, start, count = block_attn_mask_to_ragged_lut(
-            block_mask,
-            num_heads=block_mask.shape[1],
-            return_none_if_dense=False,
+    rows = []
+    for batch, sequence_q, sequence_k, heads, dtype in itertools.product(
+        args.batch, args.sequence_q, args.sequence_k, args.heads, args.dtype
+    ):
+        if dtype != dtypes.bf16:
+            aiter.logger.warning("MHA v4 dense benchmark skips dtype %s", dtype)
+            continue
+        rows.append(benchmark_mha_v4(batch, sequence_q, sequence_k, heads, dtype))
+    if rows:
+        frame = pd.DataFrame(rows)
+        aiter.logger.info(
+            "MHA v4 dense summary (markdown):\n%s",
+            frame.to_markdown(index=False),
         )
-        lut = {
-            "kv_block_indices": indices,
-            "lut_start": start,
-            "lut_count": count,
-        }
-    fp8_format = native_fp8_format()
-    return mha_v4_packed(
-        q.quantized,
-        k.quantized,
-        v.quantized,
-        q.descale,
-        k.descale,
-        v.descale,
-        fp8_format,
-        fp8_format,
-        fp8_format,
-        AttentionScaleMode.F32_PER_TENSOR,
-        AttentionScaleMode.F32_PER_TENSOR,
-        AttentionScaleMode.F32_PER_TENSOR,
-        **lut,
-    )
 
 
-def _gather_kv_tiles(operand, tiles):
-    """Concatenate the named KV tiles, leaving the quantized bytes and descale untouched."""
-    kv_tile = mha_v4_kv_tile()
-    gathered = torch.cat(
-        [operand.quantized[:, tile * kv_tile : (tile + 1) * kv_tile] for tile in tiles],
-        dim=1,
-    )
-    return _Operand(gathered.contiguous(), operand.descale)
+def _dense_reference(q, k, v, valid):
+    """Attention for one batch over its first `valid` keys, in BSHD."""
+    return torch.nn.functional.scaled_dot_product_attention(
+        q.permute(0, 2, 1, 3),
+        k[:, :valid].permute(0, 2, 1, 3),
+        v[:, :valid].permute(0, 2, 1, 3),
+    ).permute(0, 2, 1, 3)
 
 
-def _tile_mask(heads, kv_tiles, tiles, q_tiles=1, batch=1):
-    mask = torch.zeros(
-        (batch, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool
-    )
-    for tile in tiles:
-        mask[:, :, :, tile] = True
-    return mask
-
+@pytest.mark.parametrize("lengths", [(300, 137), (512, 64), (1, 511), (65, 65)])
+def test_mha_v4_seqlens_k_attends_over_each_batch_length(lengths):
+    torch.manual_seed(41)
+    batch, sequence, heads = len(lengths), 512, 4
+    q = torch.randn((batch, sequence, heads, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    seqlens_k = torch.tensor(lengths, device="cuda", dtype=torch.int32)
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-@pytest.mark.parametrize("tiles", [(0,), (1,), (3,), (0, 2), (1, 2, 3)])
-def test_mha_v4_sparse_reads_only_the_kv_tiles_the_lut_names(tiles):
-    """A kernel that ignored kv_block_indices would pass every all-True test."""
-    heads = 2
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
-
-    mask = _tile_mask(heads, kv_tiles, tiles)
-    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
-    # Dense over exactly the selected tiles: same quantized bytes, same descales, so the
-    # only difference is which KV blocks take part.
-    dense = _sparse_fp8_launch(
-        q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+    out = mha_v4(
+        q,
+        k,
+        v,
+        AttentionFormat.BF16,
+        AttentionFormat.BF16,
+        AttentionFormat.BF16,
+        seqlens_k=seqlens_k,
     )
     torch.cuda.synchronize()
 
-    _assert_sparse_matches_dense(sparse, dense)
-
-
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_sparse_distinct_kv_tiles_give_distinct_results():
-    """Guards the reference itself: selecting different tiles must change the output."""
-    heads = 2
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
-
-    outputs = [
-        _sparse_fp8_launch(q, k, v, block_mask=_tile_mask(heads, kv_tiles, (tile,)))
-        for tile in range(kv_tiles)
-    ]
-    torch.cuda.synchronize()
-
-    for tile in range(1, kv_tiles):
-        assert not torch.equal(outputs[0], outputs[tile]), (
-            f"kv tile 0 and kv tile {tile} produced identical output, so the kernel is "
-            "not reading kv_block_indices"
+    for b, valid in enumerate(lengths):
+        reference = _dense_reference(q[b : b + 1], k[b : b + 1], v[b : b + 1], valid)
+        cosine = torch.nn.functional.cosine_similarity(
+            out[b : b + 1].float().flatten(), reference.float().flatten(), dim=0
         )
+        assert cosine > 0.99, f"batch {b} of {lengths}"
 
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_sparse_gives_each_head_its_own_kv_tiles():
-    """4-D masks may give heads different KV lists; each head must follow its own row."""
-    heads = 3
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    per_head = ((0,), (3,), (1, 2))
-    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
-
-    mask = torch.zeros((1, heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
-    for head, tiles in enumerate(per_head):
-        for tile in tiles:
-            mask[:, head, :, tile] = True
-    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
-    torch.cuda.synchronize()
-
-    for head, tiles in enumerate(per_head):
-        dense = _sparse_fp8_launch(
-            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
-        )
-        torch.cuda.synchronize()
-        _assert_sparse_matches_dense(
-            sparse[:, :, head],
-            dense[:, :, head],
-            f"head {head} did not attend to tiles {tiles}",
-        )
-
+def test_mha_v4_seqlens_k_at_full_length_is_the_dense_result():
+    """A null slot and a slot naming the whole key length must run the same code path."""
+    torch.manual_seed(41)
+    q = torch.randn((2, 512, 4, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    formats = (AttentionFormat.BF16, AttentionFormat.BF16, AttentionFormat.BF16)
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_sparse_follows_the_lut_across_query_tiles():
-    """Multiple query tiles exercise the work table on a real launch, not just its ordering."""
-    heads = 2
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    q_tiles = 2
-    q, k, v = _sparse_fp8_operands(
-        sequence_k=kv_tiles * kv_tile, heads=heads, sequence_q=256 * q_tiles
+    dense = mha_v4(q, k, v, *formats)
+    full = mha_v4(
+        q,
+        k,
+        v,
+        *formats,
+        seqlens_k=torch.full((2,), 512, device="cuda", dtype=torch.int32),
     )
-
-    mask = torch.zeros((1, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool)
-    mask[:, :, 0, 0] = True
-    mask[:, :, 1, 3] = True
-    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
     torch.cuda.synchronize()
+    assert torch.equal(dense, full)
 
-    for q_tile, tiles in ((0, (0,)), (1, (3,))):
-        dense = _sparse_fp8_launch(
-            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
-        )
-        torch.cuda.synchronize()
-        rows = slice(q_tile * 256, (q_tile + 1) * 256)
-        _assert_sparse_matches_dense(
-            sparse[:, rows],
-            dense[:, rows],
-            f"query tile {q_tile} did not attend to tiles {tiles}",
-        )
-
-
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-@pytest.mark.parametrize("tail_rows", [64, 128, 200])
-def test_mha_v4_sparse_partial_query_tile_follows_the_lut(tail_rows):
-    """A trailing partial query tile has to select and zero the same way a full one does.
-
-    Production sequence lengths are not multiples of the 256-row query tile, and the empty-row
-    no-op is built on the same tail masking the partial tile uses, so the two belong in one
-    case: the short tile must still read the KV blocks its row names, and an all-False row on
-    that tile must come back zero rather than reading the masked-off remainder.
-    """
-    heads = 2
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    q_tiles = 3
-    per_tile = ((0, (0,)), (1, (1, 2)), (2, (3,)))
-    q, k, v = _sparse_fp8_operands(
-        sequence_k=kv_tiles * kv_tile,
-        heads=heads,
-        sequence_q=256 * (q_tiles - 1) + tail_rows,
-    )
 
-    mask = torch.zeros((1, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool)
-    for q_tile, tiles in per_tile:
-        for tile in tiles:
-            mask[:, :, q_tile, tile] = True
-    mask[:, 1, q_tiles - 1, :] = False  # head 1's partial-tile row selects nothing
-    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
-    torch.cuda.synchronize()
+def test_mha_v4_rejects_unusable_seqlens_k():
+    q = torch.randn((2, 256, 4, 128), device="cuda", dtype=torch.bfloat16)
+    formats = (AttentionFormat.BF16, AttentionFormat.BF16, AttentionFormat.BF16)
 
-    for q_tile, tiles in per_tile:
-        dense = _sparse_fp8_launch(
-            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+    with pytest.raises(ValueError, match="int32"):
+        mha_v4(
+            q, q, q, *formats, seqlens_k=torch.ones(2, device="cuda", dtype=torch.int64)
         )
-        torch.cuda.synchronize()
-        rows = slice(q_tile * 256, min((q_tile + 1) * 256, sparse.shape[1]))
-        live_heads = 1 if q_tile == q_tiles - 1 else heads
-        _assert_sparse_matches_dense(
-            sparse[:, rows, :live_heads],
-            dense[:, rows, :live_heads],
-            f"query tile {q_tile} did not attend to tiles {tiles}",
+    with pytest.raises(ValueError, match="one entry per batch"):
+        mha_v4(
+            q, q, q, *formats, seqlens_k=torch.ones(1, device="cuda", dtype=torch.int32)
         )
-
-    tail = slice((q_tiles - 1) * 256, sparse.shape[1])
-    assert torch.equal(
-        sparse[:, tail, 1], torch.zeros_like(sparse[:, tail, 1])
-    ), "empty row on the partial query tile is not zero"
-    # Without this the case would also pass on a kernel that skipped the short tile entirely,
-    # since both sides of the comparison above would then be zero.
-    assert (
-        sparse[:, tail, 0].abs().max() > 0
-    ), "live partial-tile row came back degenerate"
-    assert torch.isfinite(sparse).all(), "partial query tile leaked NaN or infinity"
-
-
-def _gfx950_only(launch, label):
-    return pytest.param(
-        launch,
-        id=label,
-        marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
-    )
-
-
-_EMPTY_ROW_LAUNCHES = [
-    pytest.param(
-        lambda q, k, v, m: mha_v4(
+    with pytest.raises(NotImplementedError, match="per-batch key lengths"):
+        mha_v4(
+            q,
             q,
-            k,
-            v,
-            native_fp8_format(),
-            native_fp8_format(),
-            native_fp8_format(),
-            block_mask=m,
-        ),
-        id="fp8",
-    ),
-    pytest.param(
-        lambda q, k, v, m: mha_v4(
             q,
-            k,
-            v,
             AttentionFormat.INT8,
             AttentionFormat.INT8,
             native_fp8_format(),
-            block_mask=m,
-        ),
-        id="i8fp8",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            native_fp8_format(),
-            native_fp8_format(),
-            native_fp8_format(),
-            block_mask=m,
-            q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-            k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
-            v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
-        ),
-        "mxfp8",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            native_fp8_format(),
-            native_fp8_format(),
-            AttentionFormat.MXFP6,
-            block_mask=m,
-        ),
-        "f8f6",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP6,
-            native_fp8_format(),
-            block_mask=m,
-        ),
-        "f6f8",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP6,
-            AttentionFormat.MXFP4,
-            block_mask=m,
-        ),
-        "f6f4",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            AttentionFormat.MXFP4,
-            AttentionFormat.MXFP4,
-            native_fp8_format(),
-            block_mask=m,
-        ),
-        "mxfp4",
-    ),
-    _gfx950_only(
-        lambda q, k, v, m: mha_v4(
-            q,
-            k,
-            v,
-            AttentionFormat.MXFP4,
-            AttentionFormat.MXFP4,
-            AttentionFormat.MXFP4,
-            block_mask=m,
-        ),
-        "f4f4",
-    ),
-]
-
+            block_mask=torch.ones(
+                (2, 4, 1, 256 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
+            ),
+            seqlens_k=torch.ones(2, device="cuda", dtype=torch.int32),
+        )
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-@pytest.mark.parametrize("launch", _EMPTY_ROW_LAUNCHES)
-def test_mha_v4_sparse_empty_row_writes_zeros(launch):
-    """An all-False row selects no KV block, so its output tile must be zero, not garbage.
 
-    Its lut_start also sits one past the last kv_block_indices entry, which is what used to walk
-    the prologue's unguarded reads into a multi-gigabyte scalar offset and fault the kernel.
-    """
-    heads = 2
-    kv_tiles = 4
-    selected = (0, 1)
-    torch.manual_seed(0)
-    q = torch.randn((1, 256, heads, 128), device="cuda", dtype=torch.bfloat16)
-    k = torch.randn(
-        (1, kv_tiles * mha_v4_kv_tile(), heads, 128),
-        device="cuda",
-        dtype=torch.bfloat16,
+def _launch_bf16_dense(q, out, seqlens_k=None, lse=None):
+    """Drive the dense launcher directly, past the screening mha_v4_packed does first."""
+    bf16 = int(AttentionFormat.BF16)
+    torch.ops.aiter.mha_v4_fwd_launch(
+        q,
+        q,
+        q,
+        q,
+        q,
+        q,
+        out,
+        bf16,
+        bf16,
+        bf16,
+        int(AttentionPack.DEFAULT),
+        0,
+        0,
+        0,
+        128**-0.5,
+        seqlens_k,
+        lse,
     )
-    v = torch.randn_like(k)
-
-    mask = torch.zeros((1, heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
-    for tile in selected:
-        mask[:, 0, :, tile] = True  # head 0 selects two tiles; head 1 stays all-False
-    out = launch(q, k, v, mask)
-    torch.cuda.synchronize()
-
-    assert torch.equal(
-        out[:, :, 1], torch.zeros_like(out[:, :, 1])
-    ), "empty row is not zero"
-    assert torch.isfinite(out).all(), "empty row leaked NaN or infinity"
-    assert out[:, :, 0].abs().max() > 0, "live row came back degenerate"
-
-    # The empty row must not perturb the row that does select tiles: give head 1 a tile and head 0's
-    # output has to stay bit-identical, since each workgroup owns one (batch, head, query tile).
-    mask[:, 1, :, 0] = True
-    populated = launch(q, k, v, mask)
-    torch.cuda.synchronize()
-    assert torch.equal(
-        out[:, :, 0], populated[:, :, 0]
-    ), "empty row disturbed the live row"
 
 
-def test_mha_v4_rejects_non_bool_block_mask():
-    """Counts come from a sum but the fill uses truthiness, so non-bool masks disagree."""
-    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16)
-    mask = torch.ones((1, 2, 1, 256 // mha_v4_kv_tile()), dtype=torch.int32)
-    with pytest.raises(ValueError, match="block_mask must be a bool tensor"):
-        mha_v4(
-            q,
-            q,
-            q,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            block_mask=mask,
-        )
+def test_mha_v4_launch_rejects_off_device_seqlens_k():
+    """The launcher hands this pointer straight to the GPU, so it guards independently of Python."""
+    q = torch.randn((2, 256, 4, 128), device="cuda", dtype=torch.bfloat16)
+    out = torch.empty_like(q)
 
+    with pytest.raises(RuntimeError, match="same device as Q"):
+        _launch_bf16_dense(q, out, seqlens_k=torch.full((2,), 256, dtype=torch.int32))
 
-@pytest.mark.skipif(
-    torch.cuda.device_count() < 2, reason="needs two GPUs to mismatch devices"
-)
-def test_mha_v4_rejects_block_mask_on_another_device():
-    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16, device="cuda:0")
-    mask = torch.ones(
-        (1, 2, 1, 256 // mha_v4_kv_tile()), dtype=torch.bool, device="cuda:1"
-    )
-    with pytest.raises(ValueError, match="block_mask must be on the same device"):
-        mha_v4(
-            q,
-            q,
-            q,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            AttentionFormat.FP8,
-            block_mask=mask,
-        )
 
+def test_mha_v4_launch_rejects_off_device_lse():
+    """mha_v4_packed takes an LSE buffer from the caller and never checks where it lives."""
+    q = torch.randn((2, 256, 4, 128), device="cuda", dtype=torch.bfloat16)
+    out = torch.empty_like(q)
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
-def test_mha_v4_sparse_rejects_empty_kv_block_indices():
-    """Rows may be empty, but the ASM still dereferences the row base, so the buffer cannot be."""
-    heads = 2
-    kv_tile = mha_v4_kv_tile()
-    kv_tiles = 4
-    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
-    rows = heads  # batch 1, one query tile
-    device = q.quantized.device
-    fp8_format = native_fp8_format()
-    with pytest.raises(RuntimeError, match="must be non-empty"):
-        mha_v4_packed(
-            q.quantized,
-            k.quantized,
-            v.quantized,
-            q.descale,
-            k.descale,
-            v.descale,
-            fp8_format,
-            fp8_format,
-            fp8_format,
-            AttentionScaleMode.F32_PER_TENSOR,
-            AttentionScaleMode.F32_PER_TENSOR,
-            AttentionScaleMode.F32_PER_TENSOR,
-            kv_block_indices=torch.zeros(0, dtype=torch.int32, device=device),
-            lut_start=torch.zeros(rows, dtype=torch.int32, device=device),
-            lut_count=torch.ones(rows, dtype=torch.int32, device=device),
-        )
+    with pytest.raises(RuntimeError, match="same device as Q"):
+        _launch_bf16_dense(q, out, lse=torch.empty((2, 4, 256), dtype=torch.float32))
 
 
-@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
-@pytest.mark.skipif(
-    not _mha_v4_sparse_co_available(),
-    reason="sorted-sparse MHA v4 code object is not deployed",
-)
 @pytest.mark.parametrize(
-    "mutation,message",
+    "formats",
     [
-        pytest.param(
-            "indices.fill_(9999)",
-            "outside",
-            id="index_out_of_range",
-        ),
-        pytest.param(
-            "start.fill_(-1)",
-            "negative",
-            id="negative_start",
-        ),
+        (AttentionFormat.MXFP4, AttentionFormat.MXFP4, AttentionFormat.MXFP4),
+        (AttentionFormat.MXFP6, AttentionFormat.MXFP6, AttentionFormat.MXFP6),
+        (AttentionFormat.INT8, AttentionFormat.INT8, None),
+        (None, None, None),
     ],
 )
-def test_mha_v4_sparse_validation_rejects_malformed_lut(mutation, message):
-    """Enable opt-in validation before AITER loads, without slowing the parent test process."""
-    probe = f"""
-from op_tests.test_mha_v4 import (
-    AttentionFormat,
-    AttentionScaleMode,
-    _sparse_fp8_operands,
-    _tile_mask,
-    block_attn_mask_to_ragged_lut,
-    mha_v4_kv_tile,
-    mha_v4_packed,
-    native_fp8_format,
-)
-
-heads = 2
-kv_tiles = 4
-q, k, v = _sparse_fp8_operands(
-    sequence_k=kv_tiles * mha_v4_kv_tile(), heads=heads
-)
-mask = _tile_mask(heads, kv_tiles, (0, 1))
-indices, start, count = block_attn_mask_to_ragged_lut(
-    mask, num_heads=heads, return_none_if_dense=False
-)
-{mutation}
-fp8_format = native_fp8_format()
-mha_v4_packed(
-    q.quantized,
-    k.quantized,
-    v.quantized,
-    q.descale,
-    k.descale,
-    v.descale,
-    fp8_format,
-    fp8_format,
-    fp8_format,
-    AttentionScaleMode.F32_PER_TENSOR,
-    AttentionScaleMode.F32_PER_TENSOR,
-    AttentionScaleMode.F32_PER_TENSOR,
-    kv_block_indices=indices,
-    lut_start=start,
-    lut_count=count,
-)
-"""
-    env = {**os.environ, "AITER_MHA_V4_VALIDATE_LUT": "1"}
-    result = subprocess.run(
-        [sys.executable, "-c", probe],
-        cwd=AITER_ROOT_DIR,
-        env=env,
-        text=True,
-        capture_output=True,
-        check=False,
-    )
-    assert result.returncode != 0
-    assert message in result.stderr
-
-
-def run_torch_mha_v4(q, k, v, softmax_scale):
-    """Compute dense BSHD attention in FP32 for benchmark validation."""
-    scores = (
-        torch.matmul(
-            q.transpose(1, 2).float(), k.transpose(1, 2).float().transpose(-1, -2)
-        )
-        * softmax_scale
-    )
-    return torch.matmul(
-        torch.softmax(scores, dim=-1), v.transpose(1, 2).float()
-    ).transpose(1, 2)
+def test_mha_v4_rejects_seqlens_k_on_recipes_that_ignore_it(formats):
+    """Only the BF16 Q/K objects read the slot; the rest must fail, not silently pad-attend.
 
-
-@benchmark()
-def benchmark_mha_v4(batch, sequence_q, sequence_k, heads, dtype):
-    """Benchmark the public dense BF16 MHA v4 path against a Torch reference."""
-    head_dim = 128
-    softmax_scale = head_dim**-0.5
-    torch.manual_seed(batch + sequence_q + sequence_k + heads)
-    q = torch.randn((batch, sequence_q, heads, head_dim), device="cuda", dtype=dtype)
-    k = torch.randn((batch, sequence_k, heads, head_dim), device="cuda", dtype=dtype)
-    v = torch.randn_like(k)
-    reference = run_torch_mha_v4(q, k, v, softmax_scale)
-    candidates = {
-        "mha_v4": lambda: mha_v4(
+    MXFP4 and MXFP6 matter most: they return through their own launchers, which never forward
+    seqlens_k, so a missing gate here is invisible rather than merely unsupported.
+    """
+    q_format, k_format, v_format = formats
+    fp8 = native_fp8_format()
+    q_format = q_format or fp8
+    k_format = k_format or fp8
+    v_format = v_format or fp8
+    q = torch.randn((2, 256, 4, 128), device="cuda", dtype=torch.bfloat16)
+
+    with pytest.raises(NotImplementedError, match="per-batch key lengths"):
+        mha_v4(
             q,
-            k,
-            v,
-            AttentionFormat.BF16,
-            AttentionFormat.BF16,
-            AttentionFormat.BF16,
-            softmax_scale=softmax_scale,
-        )
-    }
-    flops = 4 * batch * heads * sequence_q * sequence_k * head_dim
-    elements = batch * heads * head_dim * (sequence_q * 2 + sequence_k * 2)
-    nbytes = elements * q.element_size()
-    ret = {"gfx": get_gfx()}
-    for name, candidate in candidates.items():
-        output, us = run_perftest(candidate)
-        err = checkAllclose(
-            reference,
-            output.to(dtypes.fp32),
-            rtol=2e-2,
-            atol=2e-2,
-            msg=f"{name}: dense BF16",
-        )
-        ret[f"{name} us"] = us
-        ret[f"{name} TFLOPS"] = flops / us / 1e6
-        ret[f"{name} TB/s"] = nbytes / us / 1e6
-        ret[f"{name} err"] = err
-    return ret
-
-
-def main():
-    if get_gfx() != "gfx950":
-        aiter.logger.warning(
-            "MHA v4 BF16 benchmark unsupported on %s; skipping", get_gfx()
-        )
-        return
-
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.RawTextHelpFormatter,
-        description="Benchmark dense BF16 MHA v4",
-    )
-    parser.add_argument("-b", "--batch", type=int, nargs="*", default=[1])
-    parser.add_argument("--sequence-q", type=int, nargs="*", default=[128, 256])
-    parser.add_argument("--sequence-k", type=int, nargs="*", default=[128, 256])
-    parser.add_argument("--heads", type=int, nargs="*", default=[2, 8])
-    parser.add_argument(
-        "-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default=[dtypes.bf16]
-    )
-    args = parser.parse_args()
-
-    rows = []
-    for batch, sequence_q, sequence_k, heads, dtype in itertools.product(
-        args.batch, args.sequence_q, args.sequence_k, args.heads, args.dtype
-    ):
-        if dtype != dtypes.bf16:
-            aiter.logger.warning("MHA v4 BF16 benchmark skips dtype %s", dtype)
-            continue
-        rows.append(benchmark_mha_v4(batch, sequence_q, sequence_k, heads, dtype))
-    if rows:
-        frame = pd.DataFrame(rows)
-        aiter.logger.info(
-            "MHA v4 dense BF16 summary (markdown):\n%s",
-            frame.to_markdown(index=False),
+            q,
+            q,
+            q_format,
+            k_format,
+            v_format,
+            seqlens_k=torch.full((2,), 128, device="cuda", dtype=torch.int32),
         )
 
 
diff --git a/op_tests/test_mha_v4_sparse.py b/op_tests/test_mha_v4_sparse.py
new file mode 100644
index 00000000000..23990de7674
--- /dev/null
+++ b/op_tests/test_mha_v4_sparse.py
@@ -0,0 +1,1393 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Block-sparse MHA v4: LUT handling, work-table ordering, and tile-count scaling.
+
+Split out of test_mha_v4.py, which had grown past 2600 lines with the sparse cases forming
+one contiguous block of it. Absorbs the former test_mha_v4_sparse_tile_scaling.py.
+"""
+
+import argparse
+import itertools
+import os
+import subprocess
+import sys
+from typing import NamedTuple
+
+import pandas as pd
+import pytest
+import torch
+import torch._dynamo
+
+import aiter
+from aiter import dtypes
+from aiter.jit.core import AITER_ROOT_DIR
+from aiter.jit.utils.chip_info import get_gfx
+from aiter.ops.mha_v4 import (
+    AttentionFormat,
+    AttentionScaleMode,
+    mha_v4,
+    mha_v4_kv_tile,
+    mha_v4_packed,
+    mha_v4_sparse_work_table,
+    native_fp8_format,
+)
+from aiter.ops.mha_v4_quant import (
+    quantize_fp8,
+    quantize_fp8_rotated,
+)
+from aiter.ops.triton.attention.utils import block_attn_mask_to_ragged_lut
+from aiter.test_common import benchmark, checkAllclose, run_perftest
+
+
+@pytest.fixture(autouse=True)
+def isolate_dynamo_cache():
+    """Keep each test's ``torch.compile`` behaviour independent of the tests that ran before it.
+
+    Dynamo caches compiled entries per code object, and this file compiles ``mha_v4`` under many
+    format combinations. Sharing that cache across tests means the suite drifts toward the recompile
+    limit and whichever test compiles last fails under ``fullgraph=True`` -- a failure that reports
+    against a kernel while actually depending on how many earlier tests got far enough to compile.
+    """
+    torch._dynamo.reset()
+    yield
+
+
+def _mha_v4_sparse_co_available() -> bool:
+    gfx = get_gfx()
+    asm_dir = os.environ.get("AITER_ASM_DIR", os.path.join(AITER_ROOT_DIR, "hsa"))
+    if gfx == "gfx942":
+        return os.path.isfile(
+            os.path.join(
+                asm_dir, "gfx942", "fmha_v4_fwd", "MI300", "fwd_hd128_fp8_sparse.co"
+            )
+        )
+    return os.path.isfile(
+        os.path.join(asm_dir, "gfx950", "fmha_v4_fwd", "fwd_hd128_fp8_sparse.co")
+    )
+
+
+_MHA_V4_SPARSE_ARCH = get_gfx() in ("gfx942", "gfx950")
+
+
+def _mha_v4_mxfp6_sparse_co_available() -> bool:
+    """MXFP6 Q/K/V sparse is a gfx950-only row."""
+    if get_gfx() != "gfx950":
+        return False
+    asm_dir = os.environ.get("AITER_ASM_DIR", os.path.join(AITER_ROOT_DIR, "hsa"))
+    return os.path.isfile(
+        os.path.join(asm_dir, "gfx950", "fmha_v4_fwd", "fwd_hd128_mxfp6_sparse.co")
+    )
+
+
+def test_mha_v4_packed_rejects_partial_lut():
+    dummy = torch.empty(0)
+    with pytest.raises(ValueError, match="all be set or all omitted"):
+        mha_v4_packed(
+            dummy,
+            dummy,
+            dummy,
+            dummy,
+            dummy,
+            dummy,
+            AttentionFormat.INT8,
+            AttentionFormat.INT8,
+            AttentionFormat.FP8,
+            AttentionScaleMode.F32_PER_TENSOR,
+            AttentionScaleMode.F32_PER_TENSOR,
+            AttentionScaleMode.F32_PER_TENSOR,
+            kv_block_indices=dummy,
+        )
+
+
+def test_mha_v4_rejects_wrong_block_mask_shape():
+    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16)
+    mask = torch.ones((1, 2, 1, 1), dtype=torch.bool)
+    with pytest.raises(ValueError, match="block_mask must have shape"):
+        mha_v4(
+            q,
+            q,
+            q,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            block_mask=mask,
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse schema")
+def test_mha_v4_sparse_schema_mutates_only_out():
+    dense = str(torch.ops.aiter.mha_v4_fwd_launch.default._schema)
+    assert "Tensor kv_block_indices" not in dense
+    assert "Tensor(a6!) out" in dense
+
+    sparse = str(torch.ops.aiter.mha_v4_fwd_sparse_launch.default._schema)
+    assert "Tensor kv_block_indices" in sparse
+    assert "Tensor lut_start" in sparse
+    assert "Tensor lut_count" in sparse
+    assert "Tensor(a6!) out" in sparse
+    assert sparse.endswith("-> ()")
+
+
+def _work_table_counts(total, pattern):
+    """LUT lengths covering the tie structures the ordering has to get right."""
+    if pattern == "uniform":
+        return torch.full((total,), 7, device="cuda", dtype=torch.int32)
+    if pattern == "zeros":
+        return torch.zeros((total,), device="cuda", dtype=torch.int32)
+    if pattern == "random":
+        return torch.randint(0, 64, (total,), device="cuda", dtype=torch.int32)
+    if pattern == "wide_random":
+        return torch.randint(0, 8192, (total,), device="cuda", dtype=torch.int32)
+    if pattern == "two_values":
+        alternating = torch.arange(total, device="cuda") % 3 == 0
+        return torch.where(alternating, 9, 4).to(torch.int32)
+    if pattern == "descending":
+        return torch.arange(total, 0, -1, device="cuda", dtype=torch.int32)
+    return torch.arange(1, total + 1, device="cuda", dtype=torch.int32)
+
+
+def _unpack_work_table(table, nhead, q_tiles):
+    q_idx = (table & 0xFFFF).long()
+    h_idx = ((table >> 16) & 0xFF).long()
+    b_idx = ((table >> 24) & 0xFF).long()
+    return (b_idx * nhead + h_idx) * q_tiles + q_idx
+
+
+# The table has one entry per (batch, head, query tile). 8192 is the point where the builder hands
+# the sort to ATen, so straddle it, and include sizes that are not multiples of a wave or workgroup.
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.parametrize(
+    ("batch", "nhead", "q_tiles"),
+    [
+        (1, 1, 1),
+        (1, 8, 3),
+        (2, 5, 7),
+        (1, 16, 32),
+        (1, 32, 32),
+        (1, 5, 296),
+        (4, 16, 64),
+        (8, 16, 64),
+        (8, 32, 64),
+    ],
+)
+@pytest.mark.parametrize(
+    "pattern",
+    [
+        "uniform",
+        "zeros",
+        "random",
+        "wide_random",
+        "two_values",
+        "descending",
+        "ascending",
+    ],
+)
+def test_mha_v4_sparse_work_table_is_longest_lut_first(batch, nhead, q_tiles, pattern):
+    torch.manual_seed(7)
+    total = batch * nhead * q_tiles
+    counts = _work_table_counts(total, pattern)
+
+    table = mha_v4_sparse_work_table(counts, batch, nhead, q_tiles)
+    visited = _unpack_work_table(table, nhead, q_tiles)
+
+    # Every tile exactly once. This is the part a wrong table would turn into a wrong result.
+    assert torch.equal(visited.sort().values, torch.arange(total, device="cuda"))
+
+    # Longest LUT first, so no heavy tile straggles behind the rest.
+    ordered = counts[visited]
+    assert bool((ordered[:-1] >= ordered[1:]).all())
+
+    # Ties keep raster order, which is what leaves uniform counts spatially coherent. A stable
+    # reference sort pins the whole permutation, not just the two properties above.
+    expected = torch.argsort(counts, descending=True, stable=True)
+    assert torch.equal(visited, expected)
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.parametrize(
+    "batch,nhead,q_tiles", [(1, 16, 32), (1, 5, 296), (8, 16, 64), (8, 32, 64)]
+)
+def test_mha_v4_sparse_work_table_leaves_uniform_counts_in_raster_order(
+    batch, nhead, q_tiles
+):
+    """Top-k sparsity gives every tile the same LUT length, and that case must not be shuffled."""
+    total = batch * nhead * q_tiles
+    counts = torch.full((total,), 5, device="cuda", dtype=torch.int32)
+
+    table = mha_v4_sparse_work_table(counts, batch, nhead, q_tiles)
+
+    visited = _unpack_work_table(table, nhead, q_tiles)
+    assert torch.equal(visited, torch.arange(total, device="cuda"))
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize(
+    "launch",
+    [
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.BF16,
+                AttentionFormat.BF16,
+                AttentionFormat.BF16,
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(
+                get_gfx() != "gfx950", reason="gfx950 BF16 sparse"
+            ),
+            id="bf16",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.BF16,
+                AttentionFormat.BF16,
+                native_fp8_format(),
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(
+                get_gfx() != "gfx950", reason="gfx950 BF16 sparse"
+            ),
+            id="bf16fp8",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                native_fp8_format(),
+                native_fp8_format(),
+                native_fp8_format(),
+                block_mask=mask,
+            ),
+            id="fp8",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.INT8,
+                AttentionFormat.INT8,
+                native_fp8_format(),
+                block_mask=mask,
+            ),
+            id="i8fp8",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                native_fp8_format(),
+                native_fp8_format(),
+                native_fp8_format(),
+                block_mask=mask,
+                q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
+                k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
+                v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
+            ),
+            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+            id="mxfp8",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                native_fp8_format(),
+                native_fp8_format(),
+                AttentionFormat.MXFP6,
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+            id="f8f6",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.MXFP6,
+                AttentionFormat.MXFP6,
+                native_fp8_format(),
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+            id="f6f8",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.MXFP6,
+                AttentionFormat.MXFP6,
+                AttentionFormat.MXFP4,
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+            id="f6f4",
+        ),
+        pytest.param(
+            lambda q, k, v, mask: mha_v4(
+                q,
+                k,
+                v,
+                AttentionFormat.MXFP4,
+                AttentionFormat.MXFP4,
+                AttentionFormat.MXFP4,
+                block_mask=mask,
+            ),
+            marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+            id="mxfp4",
+        ),
+    ],
+)
+def test_mha_v4_sparse_all_true_mask_matches_dense(launch):
+    torch.manual_seed(41)
+    q = torch.randn((1, 511, 5, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((1, 512, 5, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    mask = torch.ones(
+        (1, 5, 2, 512 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
+    )
+    dense = launch(q, k, v, None)
+    sparse = launch(q, k, v, mask)
+    torch.cuda.synchronize()
+    _assert_sparse_matches_dense(sparse, dense)
+
+
+def _assert_sparse_matches_dense(sparse, dense, message=None):
+    """Compare code objects that use different softmax reduction schedules."""
+    cosine = torch.nn.functional.cosine_similarity(
+        sparse.float().flatten(), dense.float().flatten(), dim=0
+    )
+    assert cosine > 0.99, message
+    assert torch.isfinite(sparse).all()
+
+
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_f4f4_sparse_all_true_mask_matches_dense():
+    """Retained FP8-P sparse F4F4 remains close to dense FP6-P on an all-true mask."""
+    torch.manual_seed(41)
+    q = torch.randn((1, 511, 5, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((1, 512, 5, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    mask = torch.ones(
+        (1, 5, 2, 512 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
+    )
+    args = (AttentionFormat.MXFP4,) * 3
+    dense = mha_v4(q, k, v, *args)
+    sparse = mha_v4(q, k, v, *args, block_mask=mask)
+    torch.cuda.synchronize()
+
+    _assert_sparse_matches_dense(sparse, dense)
+
+
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MXFP6 validation")
+@pytest.mark.skipif(
+    not _mha_v4_mxfp6_sparse_co_available(),
+    reason="sorted-sparse MXFP6 code object is not deployed",
+)
+def test_mha_v4_mxfp6_accepts_block_mask():
+    """MXFP6 Q/K/V has a sorted-sparse row, so a block mask is dispatched, not rejected."""
+    q = torch.randn((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(
+        (1, 2, 1, 256 // mha_v4_kv_tile()), device="cuda", dtype=torch.bool
+    )
+    out = mha_v4(
+        q,
+        q,
+        q,
+        AttentionFormat.MXFP6,
+        AttentionFormat.MXFP6,
+        AttentionFormat.MXFP6,
+        block_mask=mask,
+    )
+    assert out.shape == q.shape
+    assert torch.isfinite(out).all()
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_sparse_block_mask_compiles_without_graph_breaks():
+    """The mask path derives its geometry from host state, which Dynamo cannot trace.
+
+    mha_v4_kv_tile() reads the manifest and get_gfx() shells out to rocminfo, so both sit behind
+    torch_compile_guard. Without that the sparse mask path costs graph breaks per trace and fails
+    under fullgraph, which no other test in this file would notice.
+    """
+    torch.manual_seed(41)
+    kv_tile = mha_v4_kv_tile()
+    q = torch.randn((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((1, 4 * kv_tile, 2, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    mask = torch.ones((1, 2, 1, 4), device="cuda", dtype=torch.bool)
+    fp8_format = native_fp8_format()
+
+    def call():
+        return mha_v4(q, k, v, fp8_format, fp8_format, fp8_format, block_mask=mask)
+
+    explained = torch._dynamo.explain(call)()
+    assert explained.break_reasons == [], [
+        str(reason.reason) for reason in explained.break_reasons
+    ]
+
+    eager = call()
+    compiled = torch.compile(call, fullgraph=True)()
+    torch.cuda.synchronize()
+
+    assert torch.equal(eager, compiled)
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize(
+    ("q_format", "v_format"),
+    [
+        pytest.param(
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            marks=pytest.mark.skipif(
+                get_gfx() != "gfx950", reason="gfx950 BF16 sparse"
+            ),
+            id="bf16",
+        ),
+        pytest.param(
+            AttentionFormat.BF16,
+            native_fp8_format(),
+            marks=pytest.mark.skipif(
+                get_gfx() != "gfx950", reason="gfx950 BF16 sparse"
+            ),
+            id="bf16fp8",
+        ),
+        pytest.param(
+            native_fp8_format(),
+            native_fp8_format(),
+            id="fp8",
+        ),
+        pytest.param(
+            AttentionFormat.MXFP4,
+            AttentionFormat.MXFP4,
+            marks=pytest.mark.skipif(
+                get_gfx() != "gfx950", reason="gfx950 MXFP4 sparse"
+            ),
+            id="mxfp4",
+        ),
+    ],
+)
+def test_mha_v4_sparse_gqa_all_true_mask_matches_repeated_kv(q_format, v_format):
+    torch.manual_seed(41)
+    query_heads = 8
+    kv_heads = 2
+    gqa_ratio = query_heads // kv_heads
+    q = torch.randn((1, 256, query_heads, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((1, 256, kv_heads, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    kv_tiles = 256 // mha_v4_kv_tile()
+    mask = torch.ones((1, query_heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
+    k_repeated = k.repeat_interleave(gqa_ratio, dim=2)
+    v_repeated = v.repeat_interleave(gqa_ratio, dim=2)
+
+    gqa_sparse = mha_v4(q, k, v, q_format, q_format, v_format, block_mask=mask)
+    mha_sparse = mha_v4(
+        q,
+        k_repeated,
+        v_repeated,
+        q_format,
+        q_format,
+        v_format,
+        block_mask=mask,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.equal(gqa_sparse, mha_sparse)
+    if q_format != AttentionFormat.MXFP4:
+        gqa_dense = mha_v4(q, k, v, q_format, q_format, v_format)
+        mha_dense = mha_v4(q, k_repeated, v_repeated, q_format, q_format, v_format)
+        assert torch.equal(gqa_dense, mha_dense)
+        _assert_sparse_matches_dense(gqa_sparse, gqa_dense)
+    assert torch.equal(gqa_sparse, mha_sparse)
+
+
+class _Operand(NamedTuple):
+    """A quantized MHA v4 operand and the descale it was produced with."""
+
+    quantized: torch.Tensor
+    descale: torch.Tensor
+
+
+def _sparse_fp8_operands(sequence_k, heads=2, sequence_q=256, batch=1, seed=0):
+    """Quantize once so sparse and reference runs share descales exactly.
+
+    Re-quantizing a KV slice would pick a different per-tensor amax, which shifts every
+    value and hides whether the kernel read the KV blocks the LUT named.
+    """
+    torch.manual_seed(seed)
+    q = torch.randn(
+        (batch, sequence_q, heads, 128), device="cuda", dtype=torch.bfloat16
+    )
+    k = torch.randn(
+        (batch, sequence_k, heads, 128), device="cuda", dtype=torch.bfloat16
+    )
+    v = torch.randn_like(k)
+    return (
+        _Operand(*quantize_fp8_rotated(q)),
+        _Operand(*quantize_fp8_rotated(k)),
+        _Operand(*quantize_fp8(v)),
+    )
+
+
+def _sparse_fp8_launch(q, k, v, block_mask=None):
+    lut = {}
+    if block_mask is not None:
+        indices, start, count = block_attn_mask_to_ragged_lut(
+            block_mask,
+            num_heads=block_mask.shape[1],
+            return_none_if_dense=False,
+        )
+        lut = {
+            "kv_block_indices": indices,
+            "lut_start": start,
+            "lut_count": count,
+        }
+    fp8_format = native_fp8_format()
+    return mha_v4_packed(
+        q.quantized,
+        k.quantized,
+        v.quantized,
+        q.descale,
+        k.descale,
+        v.descale,
+        fp8_format,
+        fp8_format,
+        fp8_format,
+        AttentionScaleMode.F32_PER_TENSOR,
+        AttentionScaleMode.F32_PER_TENSOR,
+        AttentionScaleMode.F32_PER_TENSOR,
+        **lut,
+    )
+
+
+def _gather_kv_tiles(operand, tiles):
+    """Concatenate the named KV tiles, leaving the quantized bytes and descale untouched."""
+    kv_tile = mha_v4_kv_tile()
+    gathered = torch.cat(
+        [operand.quantized[:, tile * kv_tile : (tile + 1) * kv_tile] for tile in tiles],
+        dim=1,
+    )
+    return _Operand(gathered.contiguous(), operand.descale)
+
+
+def _tile_mask(heads, kv_tiles, tiles, q_tiles=1, batch=1):
+    mask = torch.zeros(
+        (batch, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool
+    )
+    for tile in tiles:
+        mask[:, :, :, tile] = True
+    return mask
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize("tiles", [(0,), (1,), (3,), (0, 2), (1, 2, 3)])
+def test_mha_v4_sparse_reads_only_the_kv_tiles_the_lut_names(tiles):
+    """A kernel that ignored kv_block_indices would pass every all-True test."""
+    heads = 2
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
+
+    mask = _tile_mask(heads, kv_tiles, tiles)
+    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
+    # Dense over exactly the selected tiles: same quantized bytes, same descales, so the
+    # only difference is which KV blocks take part.
+    dense = _sparse_fp8_launch(
+        q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+    )
+    torch.cuda.synchronize()
+
+    _assert_sparse_matches_dense(sparse, dense)
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_sparse_distinct_kv_tiles_give_distinct_results():
+    """Guards the reference itself: selecting different tiles must change the output."""
+    heads = 2
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
+
+    outputs = [
+        _sparse_fp8_launch(q, k, v, block_mask=_tile_mask(heads, kv_tiles, (tile,)))
+        for tile in range(kv_tiles)
+    ]
+    torch.cuda.synchronize()
+
+    for tile in range(1, kv_tiles):
+        assert not torch.equal(outputs[0], outputs[tile]), (
+            f"kv tile 0 and kv tile {tile} produced identical output, so the kernel is "
+            "not reading kv_block_indices"
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_sparse_gives_each_head_its_own_kv_tiles():
+    """4-D masks may give heads different KV lists; each head must follow its own row."""
+    heads = 3
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    per_head = ((0,), (3,), (1, 2))
+    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
+
+    mask = torch.zeros((1, heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
+    for head, tiles in enumerate(per_head):
+        for tile in tiles:
+            mask[:, head, :, tile] = True
+    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
+    torch.cuda.synchronize()
+
+    for head, tiles in enumerate(per_head):
+        dense = _sparse_fp8_launch(
+            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+        )
+        torch.cuda.synchronize()
+        _assert_sparse_matches_dense(
+            sparse[:, :, head],
+            dense[:, :, head],
+            f"head {head} did not attend to tiles {tiles}",
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_sparse_follows_the_lut_across_query_tiles():
+    """Multiple query tiles exercise the work table on a real launch, not just its ordering."""
+    heads = 2
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    q_tiles = 2
+    q, k, v = _sparse_fp8_operands(
+        sequence_k=kv_tiles * kv_tile, heads=heads, sequence_q=256 * q_tiles
+    )
+
+    mask = torch.zeros((1, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool)
+    mask[:, :, 0, 0] = True
+    mask[:, :, 1, 3] = True
+    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
+    torch.cuda.synchronize()
+
+    for q_tile, tiles in ((0, (0,)), (1, (3,))):
+        dense = _sparse_fp8_launch(
+            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+        )
+        torch.cuda.synchronize()
+        rows = slice(q_tile * 256, (q_tile + 1) * 256)
+        _assert_sparse_matches_dense(
+            sparse[:, rows],
+            dense[:, rows],
+            f"query tile {q_tile} did not attend to tiles {tiles}",
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize("tail_rows", [64, 128, 200])
+def test_mha_v4_sparse_partial_query_tile_follows_the_lut(tail_rows):
+    """A trailing partial query tile has to select and zero the same way a full one does.
+
+    Production sequence lengths are not multiples of the 256-row query tile, and the empty-row
+    no-op is built on the same tail masking the partial tile uses, so the two belong in one
+    case: the short tile must still read the KV blocks its row names, and an all-False row on
+    that tile must come back zero rather than reading the masked-off remainder.
+    """
+    heads = 2
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    q_tiles = 3
+    per_tile = ((0, (0,)), (1, (1, 2)), (2, (3,)))
+    q, k, v = _sparse_fp8_operands(
+        sequence_k=kv_tiles * kv_tile,
+        heads=heads,
+        sequence_q=256 * (q_tiles - 1) + tail_rows,
+    )
+
+    mask = torch.zeros((1, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool)
+    for q_tile, tiles in per_tile:
+        for tile in tiles:
+            mask[:, :, q_tile, tile] = True
+    mask[:, 1, q_tiles - 1, :] = False  # head 1's partial-tile row selects nothing
+    sparse = _sparse_fp8_launch(q, k, v, block_mask=mask)
+    torch.cuda.synchronize()
+
+    for q_tile, tiles in per_tile:
+        dense = _sparse_fp8_launch(
+            q, _gather_kv_tiles(k, tiles), _gather_kv_tiles(v, tiles)
+        )
+        torch.cuda.synchronize()
+        rows = slice(q_tile * 256, min((q_tile + 1) * 256, sparse.shape[1]))
+        live_heads = 1 if q_tile == q_tiles - 1 else heads
+        _assert_sparse_matches_dense(
+            sparse[:, rows, :live_heads],
+            dense[:, rows, :live_heads],
+            f"query tile {q_tile} did not attend to tiles {tiles}",
+        )
+
+    tail = slice((q_tiles - 1) * 256, sparse.shape[1])
+    assert torch.equal(
+        sparse[:, tail, 1], torch.zeros_like(sparse[:, tail, 1])
+    ), "empty row on the partial query tile is not zero"
+    # Without this the case would also pass on a kernel that skipped the short tile entirely,
+    # since both sides of the comparison above would then be zero.
+    assert (
+        sparse[:, tail, 0].abs().max() > 0
+    ), "live partial-tile row came back degenerate"
+    assert torch.isfinite(sparse).all(), "partial query tile leaked NaN or infinity"
+
+
+def _gfx950_only(launch, label):
+    return pytest.param(
+        launch,
+        id=label,
+        marks=pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 MX sparse"),
+    )
+
+
+_EMPTY_ROW_LAUNCHES = [
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            block_mask=m,
+        ),
+        "bf16",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            native_fp8_format(),
+            block_mask=m,
+        ),
+        "bf16fp8",
+    ),
+    pytest.param(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            native_fp8_format(),
+            native_fp8_format(),
+            native_fp8_format(),
+            block_mask=m,
+        ),
+        id="fp8",
+    ),
+    pytest.param(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.INT8,
+            AttentionFormat.INT8,
+            native_fp8_format(),
+            block_mask=m,
+        ),
+        id="i8fp8",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            native_fp8_format(),
+            native_fp8_format(),
+            native_fp8_format(),
+            block_mask=m,
+            q_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
+            k_scale_mode=AttentionScaleMode.E8M0_PER_1X32,
+            v_scale_mode=AttentionScaleMode.F32_PER_TENSOR,
+        ),
+        "mxfp8",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            native_fp8_format(),
+            native_fp8_format(),
+            AttentionFormat.MXFP6,
+            block_mask=m,
+        ),
+        "f8f6",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.MXFP6,
+            AttentionFormat.MXFP6,
+            native_fp8_format(),
+            block_mask=m,
+        ),
+        "f6f8",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.MXFP6,
+            AttentionFormat.MXFP6,
+            AttentionFormat.MXFP4,
+            block_mask=m,
+        ),
+        "f6f4",
+    ),
+    _gfx950_only(
+        lambda q, k, v, m: mha_v4(
+            q,
+            k,
+            v,
+            AttentionFormat.MXFP4,
+            AttentionFormat.MXFP4,
+            AttentionFormat.MXFP4,
+            block_mask=m,
+        ),
+        "mxfp4",
+    ),
+]
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize("launch", _EMPTY_ROW_LAUNCHES)
+def test_mha_v4_sparse_empty_row_writes_zeros(launch):
+    """An all-False row selects no KV block, so its output tile must be zero, not garbage.
+
+    Its lut_start also sits one past the last kv_block_indices entry, which is what used to walk
+    the prologue's unguarded reads into a multi-gigabyte scalar offset and fault the kernel.
+    """
+    heads = 2
+    kv_tiles = 4
+    selected = (0, 1)
+    torch.manual_seed(0)
+    q = torch.randn((1, 256, heads, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(
+        (1, kv_tiles * mha_v4_kv_tile(), heads, 128),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    v = torch.randn_like(k)
+
+    mask = torch.zeros((1, heads, 1, kv_tiles), device="cuda", dtype=torch.bool)
+    for tile in selected:
+        mask[:, 0, :, tile] = True  # head 0 selects two tiles; head 1 stays all-False
+    out = launch(q, k, v, mask)
+    torch.cuda.synchronize()
+
+    assert torch.equal(
+        out[:, :, 1], torch.zeros_like(out[:, :, 1])
+    ), "empty row is not zero"
+    assert torch.isfinite(out).all(), "empty row leaked NaN or infinity"
+    assert out[:, :, 0].abs().max() > 0, "live row came back degenerate"
+
+    # The empty row must not perturb the row that does select tiles: give head 1 a tile and head 0's
+    # output has to stay bit-identical, since each workgroup owns one (batch, head, query tile).
+    mask[:, 1, :, 0] = True
+    populated = launch(q, k, v, mask)
+    torch.cuda.synchronize()
+    assert torch.equal(
+        out[:, :, 0], populated[:, :, 0]
+    ), "empty row disturbed the live row"
+
+
+def test_mha_v4_rejects_non_bool_block_mask():
+    """Counts come from a sum but the fill uses truthiness, so non-bool masks disagree."""
+    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16)
+    mask = torch.ones((1, 2, 1, 256 // mha_v4_kv_tile()), dtype=torch.int32)
+    with pytest.raises(ValueError, match="block_mask must be a bool tensor"):
+        mha_v4(
+            q,
+            q,
+            q,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            block_mask=mask,
+        )
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2, reason="needs two GPUs to mismatch devices"
+)
+def test_mha_v4_rejects_block_mask_on_another_device():
+    q = torch.zeros((1, 256, 2, 128), dtype=torch.bfloat16, device="cuda:0")
+    mask = torch.ones(
+        (1, 2, 1, 256 // mha_v4_kv_tile()), dtype=torch.bool, device="cuda:1"
+    )
+    with pytest.raises(ValueError, match="block_mask must be on the same device"):
+        mha_v4(
+            q,
+            q,
+            q,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            AttentionFormat.FP8,
+            block_mask=mask,
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+def test_mha_v4_sparse_rejects_empty_kv_block_indices():
+    """Rows may be empty, but the ASM still dereferences the row base, so the buffer cannot be."""
+    heads = 2
+    kv_tile = mha_v4_kv_tile()
+    kv_tiles = 4
+    q, k, v = _sparse_fp8_operands(sequence_k=kv_tiles * kv_tile, heads=heads)
+    rows = heads  # batch 1, one query tile
+    device = q.quantized.device
+    fp8_format = native_fp8_format()
+    with pytest.raises(RuntimeError, match="must be non-empty"):
+        mha_v4_packed(
+            q.quantized,
+            k.quantized,
+            v.quantized,
+            q.descale,
+            k.descale,
+            v.descale,
+            fp8_format,
+            fp8_format,
+            fp8_format,
+            AttentionScaleMode.F32_PER_TENSOR,
+            AttentionScaleMode.F32_PER_TENSOR,
+            AttentionScaleMode.F32_PER_TENSOR,
+            kv_block_indices=torch.zeros(0, dtype=torch.int32, device=device),
+            lut_start=torch.zeros(rows, dtype=torch.int32, device=device),
+            lut_count=torch.ones(rows, dtype=torch.int32, device=device),
+        )
+
+
+@pytest.mark.skipif(get_gfx() != "gfx950", reason="gfx950 BF16 sparse")
+def test_mha_v4_sparse_bf16_rejects_lut_beyond_lds_capacity():
+    max_tiles = 8192
+    kv_tile = mha_v4_kv_tile()
+    q = torch.zeros((1, 1, 1, 128), device="cuda", dtype=torch.bfloat16)
+    backing = torch.zeros(128, device="cuda", dtype=torch.bfloat16)
+    kv = backing.as_strided((1, (max_tiles + 1) * kv_tile, 1, 128), (0, 0, 0, 1))
+    indices = torch.zeros(1, device="cuda", dtype=torch.int32)
+    row = torch.zeros(1, device="cuda", dtype=torch.int32)
+
+    with pytest.raises(RuntimeError, match="supports at most 8192 KV tiles"):
+        mha_v4_packed(
+            q,
+            kv,
+            kv,
+            q,
+            kv,
+            kv,
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            AttentionFormat.BF16,
+            AttentionScaleMode.NONE,
+            AttentionScaleMode.NONE,
+            AttentionScaleMode.NONE,
+            kv_block_indices=indices,
+            lut_start=row,
+            lut_count=row,
+        )
+
+
+@pytest.mark.skipif(not _MHA_V4_SPARSE_ARCH, reason="gfx942/gfx950 sparse validation")
+@pytest.mark.skipif(
+    not _mha_v4_sparse_co_available(),
+    reason="sorted-sparse MHA v4 code object is not deployed",
+)
+@pytest.mark.parametrize(
+    "mutation,message",
+    [
+        pytest.param(
+            "indices.fill_(9999)",
+            "outside",
+            id="index_out_of_range",
+        ),
+        pytest.param(
+            "start.fill_(-1)",
+            "negative",
+            id="negative_start",
+        ),
+    ],
+)
+def test_mha_v4_sparse_validation_rejects_malformed_lut(mutation, message):
+    """Enable opt-in validation before AITER loads, without slowing the parent test process."""
+    probe = f"""
+from op_tests.test_mha_v4_sparse import (
+    AttentionFormat,
+    AttentionScaleMode,
+    _sparse_fp8_operands,
+    _tile_mask,
+    block_attn_mask_to_ragged_lut,
+    mha_v4_kv_tile,
+    mha_v4_packed,
+    native_fp8_format,
+)
+
+heads = 2
+kv_tiles = 4
+q, k, v = _sparse_fp8_operands(
+    sequence_k=kv_tiles * mha_v4_kv_tile(), heads=heads
+)
+mask = _tile_mask(heads, kv_tiles, (0, 1))
+indices, start, count = block_attn_mask_to_ragged_lut(
+    mask, num_heads=heads, return_none_if_dense=False
+)
+{mutation}
+fp8_format = native_fp8_format()
+mha_v4_packed(
+    q.quantized,
+    k.quantized,
+    v.quantized,
+    q.descale,
+    k.descale,
+    v.descale,
+    fp8_format,
+    fp8_format,
+    fp8_format,
+    AttentionScaleMode.F32_PER_TENSOR,
+    AttentionScaleMode.F32_PER_TENSOR,
+    AttentionScaleMode.F32_PER_TENSOR,
+    kv_block_indices=indices,
+    lut_start=start,
+    lut_count=count,
+)
+"""
+    env = {**os.environ, "AITER_MHA_V4_VALIDATE_LUT": "1"}
+    result = subprocess.run(
+        [sys.executable, "-c", probe],
+        cwd=AITER_ROOT_DIR,
+        env=env,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode != 0
+    assert message in result.stderr
+
+
+FP8 = native_fp8_format()
+_MX_SCALES = {
+    "q_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+    "k_scale_mode": AttentionScaleMode.E8M0_PER_1X32,
+    "v_scale_mode": AttentionScaleMode.F32_PER_TENSOR,
+}
+
+# (q/k format, v format, kwargs). K takes Q's format, as the manifest rows do.
+SPARSE_RECIPES = {
+    "bf16": (AttentionFormat.BF16, AttentionFormat.BF16, {}),
+    "bf16fp8": (AttentionFormat.BF16, FP8, {}),
+    "i8fp8": (AttentionFormat.INT8, FP8, {}),
+    "fp8": (FP8, FP8, {}),
+    "mxfp8": (FP8, FP8, _MX_SCALES),
+    "f8f6": (FP8, AttentionFormat.MXFP6, {}),
+    "f6f8": (AttentionFormat.MXFP6, FP8, {}),
+    "f6f4": (AttentionFormat.MXFP6, AttentionFormat.MXFP4, {}),
+    "mxfp4": (AttentionFormat.MXFP4, AttentionFormat.MXFP4, {}),
+}
+
+# Two tiles is the baseline the growing counts are judged against; the rest span the region the
+# suite never covered. 32 tiles is where the shipped mxfp8 object had lost a quarter of its cosine.
+BASELINE_TILES = 2
+GROWN_TILES = (4, 7, 8, 12, 16, 32)
+
+# Each row carries its own quantization error, so the bar is that accuracy does not DEGRADE with
+# tile count. Only a modest absolute floor is applied on top, to catch a row that is broken outright.
+MAX_DEGRADATION = 0.01
+ABSOLUTE_FLOOR = 0.95
+
+requires_sparse = pytest.mark.skipif(
+    get_gfx() != "gfx950" or not _mha_v4_sparse_co_available(),
+    reason="gfx950 sorted-sparse MHA v4 code object is not deployed",
+)
+
+
+def _reference(q, k, v):
+    qf, kf, vf = (t.float().transpose(1, 2) for t in (q, k, v))
+    scores = torch.matmul(qf, kf.transpose(-1, -2)) * (qf.shape[-1] ** -0.5)
+    return torch.matmul(torch.softmax(scores, dim=-1), vf).transpose(1, 2)
+
+
+def _cosine(a, b):
+    return torch.nn.functional.cosine_similarity(
+        a.float().flatten(), b.float().flatten(), dim=0
+    ).item()
+
+
+def _run(recipe_name, tiles, heads=5, sequence_q=256, all_true=True):
+    q_format, v_format, kwargs = SPARSE_RECIPES[recipe_name]
+    sequence_k = tiles * mha_v4_kv_tile()
+    torch.manual_seed(41)
+    q = torch.randn((1, sequence_q, heads, 128), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((1, sequence_k, heads, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    mask = torch.zeros(
+        (1, heads, -(-sequence_q // 256), tiles), device="cuda", dtype=torch.bool
+    )
+    mask[:] = all_true
+    if not all_true:
+        mask[:, :, :, ::2] = True
+    out = mha_v4(q, k, v, q_format, q_format, v_format, block_mask=mask, **kwargs)
+    torch.cuda.synchronize()
+    return out, _reference(q, k, v), mask
+
+
+@requires_sparse
+@pytest.mark.parametrize("recipe_name", sorted(SPARSE_RECIPES))
+def test_mha_v4_sparse_accuracy_holds_as_tile_count_grows(recipe_name):
+    """A row's own small-tile-count accuracy is its baseline; growing the count must not erode it."""
+    baseline_out, baseline_ref, _ = _run(recipe_name, BASELINE_TILES)
+    baseline = _cosine(baseline_out, baseline_ref)
+    assert torch.isfinite(
+        baseline_out
+    ).all(), f"{recipe_name}: non-finite at {BASELINE_TILES} tiles"
+    assert baseline > ABSOLUTE_FLOOR, f"{recipe_name}: baseline cosine {baseline:.5f}"
+
+    for tiles in GROWN_TILES:
+        out, ref, _ = _run(recipe_name, tiles)
+        assert torch.isfinite(out).all(), f"{recipe_name}: non-finite at {tiles} tiles"
+        cosine = _cosine(out, ref)
+        assert (
+            cosine > ABSOLUTE_FLOOR
+        ), f"{recipe_name}: cosine {cosine:.5f} at {tiles} tiles"
+        assert cosine > baseline - MAX_DEGRADATION, (
+            f"{recipe_name}: accuracy degrades with tile count -- "
+            f"{baseline:.5f} at {BASELINE_TILES} tiles, {cosine:.5f} at {tiles}"
+        )
+
+
+@requires_sparse
+@pytest.mark.parametrize("recipe_name", sorted(SPARSE_RECIPES))
+def test_mha_v4_sparse_skipping_lut_holds_as_tile_count_grows(recipe_name):
+    """Same sweep with a LUT that actually skips, so the walk has to jump rather than run affine."""
+    for tiles in (4, 8, 16, 32):
+        out, _, mask = _run(recipe_name, tiles, all_true=False)
+        assert torch.isfinite(out).all(), f"{recipe_name}: non-finite at {tiles} tiles"
+        assert mask.sum() > 0
+        assert not bool(
+            (out == 0).all()
+        ), f"{recipe_name}: all-zero output at {tiles} tiles"
+
+
+@requires_sparse
+@pytest.mark.parametrize(
+    "recipe_name", ["bf16", "bf16fp8", "i8fp8", "fp8", "mxfp8", "f6f8"]
+)
+def test_mha_v4_sparse_all_true_lut_matches_dense_bitwise(recipe_name):
+    """An all-true LUT selects every tile, so the walk must reduce exactly to the dense one.
+
+    Restricted to the rows whose V stays FP8. The MX-V rows resolve to a different V packing on
+    their dense manifest row than on their sparse one, so they cannot be compared bitwise.
+    """
+    q_format, v_format, kwargs = SPARSE_RECIPES[recipe_name]
+    for tiles in (2, 8, 16):
+        sequence_k = tiles * mha_v4_kv_tile()
+        torch.manual_seed(41)
+        q = torch.randn((1, 256, 5, 128), device="cuda", dtype=torch.bfloat16)
+        k = torch.randn((1, sequence_k, 5, 128), device="cuda", dtype=torch.bfloat16)
+        v = torch.randn_like(k)
+        mask = torch.ones((1, 5, 1, tiles), device="cuda", dtype=torch.bool)
+        sparse = mha_v4(
+            q, k, v, q_format, q_format, v_format, block_mask=mask, **kwargs
+        )
+        dense = mha_v4(q, k, v, q_format, q_format, v_format, **kwargs)
+        torch.cuda.synchronize()
+        assert torch.equal(sparse, dense), (
+            f"{recipe_name}: all-true LUT differs from dense at {tiles} tiles "
+            f"(cosine {_cosine(sparse, dense):.7f})"
+        )
+
+
+_SPARSE_BENCH_GFX = ["gfx950"]
+_SPARSE_BENCH_Q_TILE = 256
+
+
+def run_torch_mha_v4_sparse(q, k, v, block_mask, softmax_scale, kv_tile):
+    """Masked dense attention in FP32: the KV blocks a LUT row drops never reach the softmax."""
+    scores = (
+        torch.matmul(
+            q.transpose(1, 2).float(), k.transpose(1, 2).float().transpose(-1, -2)
+        )
+        * softmax_scale
+    )
+    keep = block_mask.repeat_interleave(_SPARSE_BENCH_Q_TILE, dim=2)[
+        :, :, : q.shape[1], :
+    ].repeat_interleave(kv_tile, dim=3)[..., : k.shape[1]]
+    scores = scores.masked_fill(~keep, float("-inf"))
+    return torch.matmul(
+        torch.softmax(scores, dim=-1), v.transpose(1, 2).float()
+    ).transpose(1, 2)
+
+
+@benchmark()
+def benchmark_mha_v4_sparse(batch, sequence_q, sequence_k, heads, density, dtype):
+    """Benchmark the sorted block-sparse BF16 path against masked Torch attention."""
+    head_dim = 128
+    kv_tile = mha_v4_kv_tile()
+    softmax_scale = head_dim**-0.5
+    torch.manual_seed(batch + sequence_q + sequence_k + heads)
+    q = torch.randn((batch, sequence_q, heads, head_dim), device="cuda", dtype=dtype)
+    k = torch.randn((batch, sequence_k, heads, head_dim), device="cuda", dtype=dtype)
+    v = torch.randn_like(k)
+
+    q_tiles = (sequence_q + _SPARSE_BENCH_Q_TILE - 1) // _SPARSE_BENCH_Q_TILE
+    kv_tiles = (sequence_k + kv_tile - 1) // kv_tile
+    keep = max(1, round(kv_tiles * density))
+    block_mask = torch.zeros(
+        (batch, heads, q_tiles, kv_tiles), device="cuda", dtype=torch.bool
+    )
+    # Keep the first `keep` tiles of every row: a fixed prefix makes the selected count exact, so
+    # the FLOP and byte counts below describe the work the kernel really did.
+    block_mask[..., :keep] = True
+    reference = run_torch_mha_v4_sparse(q, k, v, block_mask, softmax_scale, kv_tile)
+
+    formats = (AttentionFormat.BF16, AttentionFormat.BF16, AttentionFormat.BF16)
+    candidates = {
+        "sparse": lambda: mha_v4(
+            q, k, v, *formats, softmax_scale=softmax_scale, block_mask=block_mask
+        )
+    }
+    if keep == kv_tiles:
+        # Dense reads every tile, so it only answers the same question as the reference when the
+        # mask selects everything; at lower density it would be fast and wrong.
+        candidates["dense"] = lambda: mha_v4(
+            q, k, v, *formats, softmax_scale=softmax_scale
+        )
+
+    visited_k = keep * kv_tile
+    flops = 4 * batch * heads * sequence_q * visited_k * head_dim
+    elements = batch * heads * head_dim * (sequence_q * 2 + visited_k * 2)
+    nbytes = elements * q.element_size()
+
+    ret = {"gfx": get_gfx(), "kv_tiles": kv_tiles, "kept": keep}
+    for name, candidate in candidates.items():
+        output, us = run_perftest(candidate)
+        err = checkAllclose(
+            reference,
+            output.to(dtypes.fp32),
+            rtol=2e-2,
+            atol=2e-2,
+            msg=f"{name}: block-sparse BF16",
+        )
+        ret[f"{name} us"] = us
+        ret[f"{name} TFLOPS"] = flops / us / 1e6
+        ret[f"{name} TB/s"] = nbytes / us / 1e6
+        ret[f"{name} err"] = err
+    return ret
+
+
+def main():
+    if get_gfx() not in _SPARSE_BENCH_GFX:
+        aiter.logger.warning(
+            "MHA v4 block-sparse benchmark unsupported on %s; skipping", get_gfx()
+        )
+        return
+
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawTextHelpFormatter,
+        description="Benchmark sorted block-sparse BF16 MHA v4",
+    )
+    parser.add_argument("-b", "--batch", type=int, nargs="*", default=[1])
+    parser.add_argument("--sequence-q", type=int, nargs="*", default=[256, 512])
+    parser.add_argument("--sequence-k", type=int, nargs="*", default=[1024, 2048])
+    parser.add_argument("--heads", type=int, nargs="*", default=[8])
+    parser.add_argument(
+        "--density",
+        type=float,
+        nargs="*",
+        default=[0.25, 0.5, 1.0],
+        help="fraction of KV tiles each LUT row selects",
+    )
+    parser.add_argument(
+        "-d", "--dtype", type=dtypes.str2Dtype, nargs="*", default=[dtypes.bf16]
+    )
+    args = parser.parse_args()
+
+    rows = []
+    for batch, sequence_q, sequence_k, heads, density, dtype in itertools.product(
+        args.batch,
+        args.sequence_q,
+        args.sequence_k,
+        args.heads,
+        args.density,
+        args.dtype,
+    ):
+        if dtype != dtypes.bf16:
+            aiter.logger.warning("MHA v4 sparse benchmark skips dtype %s", dtype)
+            continue
+        rows.append(
+            benchmark_mha_v4_sparse(
+                batch, sequence_q, sequence_k, heads, density, dtype
+            )
+        )
+    if rows:
+        frame = pd.DataFrame(rows)
+        aiter.logger.info(
+            "MHA v4 block-sparse BF16 summary (markdown):\n%s",
+            frame.to_markdown(index=False),
+        )
+
+
+if __name__ == "__main__":
+    main()