Fix Triton split-K fp8 dtype and graph autotuning for gfx950 (#484)

mslk/attention/fmha/_triton/splitk_kernels.py

@@ -11,6 +11,7 @@
 import torch
 import triton
 import triton.language as tl
+from mslk.utils.device import is_gfx950, rocm_version_at_least
 
 from .vararg_kernel import unroll_varargs, VAR_ARGS_ARRAY
 
@@ -121,6 +122,7 @@ def _fwd_kernel_splitK(  # noqa: C901
     HAS_ADDITIVE_BIAS: tl.constexpr,
     NUM_PROGRAMS_DIM2_CONST: tl.constexpr,
     IS_HIP: tl.constexpr,
+    FP8_FNUZ: tl.constexpr,
     QUANTIZE_PV_TO_FP8: tl.constexpr,
     QUANTIZE_QK_TO_FP8: tl.constexpr,
     USE_FP32_SCALES: tl.constexpr,
@@ -539,6 +541,7 @@ def _fwd_kernel_splitK(  # noqa: C901
                 # pyrefly: ignore [bad-argument-type]
                 i,
                 IS_HIP,
+                FP8_FNUZ,
                 QUANTIZE_PV_TO_FP8,
                 QUANTIZE_QK_TO_FP8,
                 USE_FP32_SCALES,
@@ -780,10 +783,16 @@ def autotune_kernel(kernel: Callable):
         if block_n >= block_m
     ]
 
+    # HIP graph capture faulted (HSA_INVALID_PACKET / GPU faults) during autotuning of
+    # this kernel on gfx950 with ROCm < 7.14; fixed in ROCm 7.14. Disable graph-based
+    # autotuning only on that affected stack; gfx942, CUDA, and gfx950 on ROCm >= 7.14
+    # keep it.
+    graph_autotune_broken = is_gfx950() and not rocm_version_at_least(7, 14)
+
     kernel = triton.autotune(
         configs=TRITON_CONFIGS,
         key=AUTOTUNER_KEY,
-        use_cuda_graph=True,
+        use_cuda_graph=not graph_autotune_broken,
         prune_configs_by={
             "early_config_prune": early_config_prune,
         },
@@ -834,6 +843,7 @@ def load_dequantize_k_v_group(
     v_dtype: tl.constexpr,  # Q.dtype.element_ty
     group_id: tl.constexpr,
     IS_HIP: tl.constexpr,
+    FP8_FNUZ: tl.constexpr,
     QUANTIZE_PV_TO_FP8: tl.constexpr,
     QUANTIZE_QK_TO_FP8: tl.constexpr,
     USE_FP32_SCALES: tl.constexpr,
@@ -874,6 +884,7 @@ def load_dequantize_k_v_group(
             q_dtype,
             v_dtype,
             IS_HIP,
+            FP8_FNUZ,
             QUANTIZE_PV_TO_FP8,
             QUANTIZE_QK_TO_FP8,
             USE_FP32_SCALES,
@@ -921,6 +932,7 @@ def _process_fp8_quantization(
     q_dtype: tl.constexpr,
     v_dtype: tl.constexpr,
     IS_HIP: tl.constexpr,
+    FP8_FNUZ: tl.constexpr,
     QUANTIZE_PV_TO_FP8: tl.constexpr,
     QUANTIZE_QK_TO_FP8: tl.constexpr,
     USE_FP32_SCALES: tl.constexpr,
@@ -938,6 +950,7 @@ def _process_fp8_quantization(
             v_shift if not USE_FP32_SCALES else None,
             PACKED_PER_VAL,
             IS_HIP,
+            FP8_FNUZ,
             USE_FP32_SCALES,
         ).to(v_dtype)
     else:
@@ -951,7 +964,9 @@ def _process_fp8_quantization(
     k_scale, k_shift = _extract_scale_shift(k_scale_shift, IS_HIP, USE_FP32_SCALES)
     if IS_HIP:
         if not QUANTIZE_QK_TO_FP8:
-            k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(q_dtype)
+            k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, FP8_FNUZ).to(
+                q_dtype
+            )
         else:
             # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8
             tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
@@ -964,14 +979,15 @@ def _process_fp8_quantization(
                 tl.trans(k_shift) if not USE_FP32_SCALES else None,
                 PACKED_PER_VAL,
                 IS_HIP,
+                FP8_FNUZ,
                 USE_FP32_SCALES,
             ).to(q_dtype)
             k = tl.trans(k_t)
         else:
             # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8
             tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
             k_t = tl.trans(k)
-            k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP)
+            k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP, FP8_FNUZ)
             k = tl.trans(k_t)
 
     return k, v, v_scale, k_scale
@@ -991,7 +1007,9 @@ def _extract_scale_shift(
 
 
 @triton.jit
-def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr):
+def _unpack_fp8_tensor(
+    x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr, FP8_FNUZ: tl.constexpr
+):
     """Unpack FP8 K/V tensor from int32 packed representation."""
     tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values")
 
@@ -1006,8 +1024,9 @@ def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr):
         unpacked_values, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
     )
 
-    # Convert to FP8 through bitcast
-    fp8_type = tl.float8e4b8 if IS_HIP else tl.float8e4nv
+    # Convert to FP8 through bitcast. gfx942 uses e4m3fnuz (float8e4b8); gfx950 and
+    # CUDA use OCP e4m3fn (float8e4nv). FP8_FNUZ carries the arch decision.
+    fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv
     x_ = unpacked_values.to(tl.uint8).to(fp8_type, bitcast=True)
 
     return x_
@@ -1039,18 +1058,20 @@ def _process_int4_quantization(
     if IS_HIP:
         k_scale, k_shift = cast_uint32_to_float(k_scale_shift)
         v_scale, v_shift = cast_uint32_to_float(v_scale_shift)
-        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype)
-        k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype)
+        # int4 path never reaches the fp8 branch inside dequantize; FP8_FNUZ is unused.
+        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype)
+        k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, False).to(dtype)
     else:
         k_scale, k_shift = cast_uint32_to_half2(k_scale_shift)
         v_scale, v_shift = cast_uint32_to_half2(v_scale_shift)
-        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype)
+        v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype)
         k_t = dequantize(
             tl.trans(k),
             tl.trans(k_scale),
             tl.trans(k_shift),
             PACKED_PER_VAL,
             IS_HIP,
+            False,
         ).to(dtype)
         k = tl.trans(k_t)
 
@@ -1091,6 +1112,7 @@ def dequantize_k_hip(
     scale,
     shift,
     PACKED_PER_VAL: tl.constexpr,
+    FP8_FNUZ: tl.constexpr,
 ):
     """PACKED_PER_VAL is the number of values packed into each element x_.
     For example, for int4 quantization and x_ of type int32, PACKED_PER_VAL is 8.
@@ -1110,8 +1132,8 @@ def dequantize_k_hip(
     )
 
     if PACKED_PER_VAL == 4:
-        # FP8 quantization.
-        fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv
+        # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn.
+        fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv
         dequant = (
             quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale
             + shift
@@ -1140,6 +1162,7 @@ def dequantize(
     shift,
     PACKED_PER_VAL: tl.constexpr,
     IS_HIP: tl.constexpr,
+    FP8_FNUZ: tl.constexpr,
     # pyrefly: ignore [bad-function-definition]
     USE_FP32_SCALES: tl.constexpr = False,
 ):
@@ -1160,8 +1183,8 @@ def dequantize(
         quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
     )
     if PACKED_PER_VAL == 4:
-        # FP8 quantization.
-        fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv
+        # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn.
+        fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv
         dequant = (
             quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale
         )

mslk/attention/fmha/triton_splitk.py

@@ -21,6 +21,7 @@
 )
 
 import torch
+from mslk.utils.device import supports_float8_fnuz
 
 from ._triton.available import is_triton_available
 from .attn_bias import (
@@ -862,6 +863,12 @@ def grid(META):
         )
 
         IS_HIP = torch.version.hip is not None
+        # fp8 byte format must match how the KV cache was quantized: gfx942 e4m3fnuz,
+        # gfx950/CUDA e4m3fn. Derive from supports_float8_fnuz (NOT hardcoded per-HIP).
+        if IS_HIP:
+            FP8_FNUZ = supports_float8_fnuz(throw_on_hip_incompatibility=False)
+        else:
+            FP8_FNUZ = False
 
         if inp.quantize_pv_to_fp8:
             v = v.view(torch.int8)
@@ -937,6 +944,7 @@ def grid(META):
             HAS_ADDITIVE_BIAS=attn_bias_tensor is not None,
             NUM_PROGRAMS_DIM2_CONST=split_k,
             IS_HIP=IS_HIP,
+            FP8_FNUZ=FP8_FNUZ,
             QUANTIZE_PV_TO_FP8=inp.quantize_pv_to_fp8,
             QUANTIZE_QK_TO_FP8=inp.quantize_qk_to_fp8,
             USE_FP32_SCALES=inp.use_fp32_scales,

mslk/utils/device.py

@@ -74,6 +74,21 @@ def cuda_version_at_least(major_min: int) -> bool:
     return int(torch.version.cuda.split(".")[0]) >= major_min
 
 
+def rocm_version_at_least(major_min: int, minor_min: int = 0) -> bool:
+    """True on a ROCm build whose HIP runtime version is at least
+    ``(major_min, minor_min)``.
+
+    Parses ``torch.version.hip`` (e.g. ``"7.14.60850"``). Returns ``False`` on CUDA or
+    CPU-only builds.
+    """
+    if torch.version.hip is None:
+        return False
+    parts = torch.version.hip.split(".")
+    major = int(parts[0])
+    minor = int(parts[1]) if len(parts) > 1 else 0
+    return (major, minor) >= (major_min, minor_min)
+
+
 def get_gfx_arch_name() -> str:
     """Return the ROCm ``gcnArchName`` of the current device (e.g. ``gfx942``).
 

test/attention/fmha/test_mem_eff_attention.py

@@ -1997,13 +1997,19 @@ def test_triton_splitk_rowwise_fp8(
         inp_ref, op=fmha.triton_splitk.FwOp
     )
 
+    # ROCm gfx950 OCP e4m3fn snaps a few values differently than the fnuz grid, so
+    # loosen tolerances on ROCm only; CUDA keeps the original tight values.
+    is_hip = torch.version.hip is not None
     atol = 5e-3
-    if Hkv == 2 and torch.version.hip is not None:
-        # XXX why is this needed?
+    rtol = 5e-3
+    if Hkv == 2 and is_hip:
         atol = 1e-2
-    torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=5e-3)
+    torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=rtol)
     assert context_fp8 is not None and context_ref is not None
-    torch.testing.assert_close(context_fp8.lse, context_ref.lse, atol=5e-4, rtol=5e-4)
+    lse_tol = 5e-3 if is_hip else 5e-4
+    torch.testing.assert_close(
+        context_fp8.lse, context_ref.lse, atol=lse_tol, rtol=lse_tol
+    )
 
     # Paged K/V cache
 
@@ -2017,8 +2023,12 @@ def test_triton_splitk_rowwise_fp8(
     ) = fmha._memory_efficient_attention_forward_requires_grad(
         inp_fp8_paged, op=fmha.triton_splitk.FwOp
     )
+    # Non-paged vs paged fp8 output: a few elements snap to a different e4m3 grid point
+    # between the two layouts on ROCm; CUDA keeps the original tight value. The LSE is
+    # identical between layouts on both platforms, so it keeps the tight tolerance.
+    paged_tol = 5e-3 if is_hip else 2e-3
     torch.testing.assert_close(
-        attn_output_fp8, attn_output_fp8_paged, atol=2e-3, rtol=2e-3
+        attn_output_fp8, attn_output_fp8_paged, atol=paged_tol, rtol=paged_tol
     )
     assert context_fp8_paged is not None
     torch.testing.assert_close(

test/attention/fmha/utils.py

@@ -25,6 +25,7 @@
     ref_attention_bmhk,
 )
 from mslk.attention.fmha.triton_splitk import InputsFp8
+from mslk.utils.triton.fp8_utils import get_fp8_constants
 
 IN_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER") is not None
 
@@ -186,9 +187,9 @@ def construct_fp8_attention_inputs(
     k = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device)
     v = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device)
 
-    pt_fp8_dtype = (
-        torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn
-    )
+    # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn,
+    # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN.
+    pt_fp8_dtype = get_fp8_constants()[0]
 
     qfn = quantize_fp8_symmetric if use_symmetric else quantize_fp8_asymmetric
 
@@ -427,9 +428,9 @@ def add_q_fp8_to_inputs(
         InputsFp8 object with quantized query tensor
     """
     inp.quantize_qk_to_fp8 = True
-    pt_fp8_dtype = (
-        torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn
-    )
+    # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn,
+    # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN.
+    pt_fp8_dtype = get_fp8_constants()[0]
     # Get original query tensor
     q = inp.query