decode attn fp8

xformers/ops/fmha/_triton/splitk_kernels.py

@@ -360,13 +360,6 @@ def _fwd_kernel_splitK(
     # https://github.com/triton-lang/triton/issues/5466
     log2e = tl.full((), 1.44269504, tl.float32)
     qk_scale = sm_scale * log2e
-    # load q: it will stay in SRAM throughout
-    q: "VAR_ARGS_ARRAY"  # noqa: F821
-    for i in range(len(acc)):  # noqa: F821
-        q[i] = tl.load(  # noqa: F821
-            tl.advance(Q_block_ptr, (0, i * D_PER_GROUP)), boundary_check=(0,)
-        )
-
     if IS_CAUSAL or IS_LOCAL:
         # Why does the masking conditon below work as a causal mask?
         # Assuming num_queries <= BLOCK_M:
@@ -454,14 +447,14 @@ def _fwd_kernel_splitK(
                 V_scale_shift_block_ptr = None
             logical_block_idx += 1
 
-        k: "VAR_ARGS_ARRAY"  # noqa: F821
-        v: "VAR_ARGS_ARRAY"  # noqa: F821
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
         for i in range(len(acc)):  # noqa: F821
-            k[i], v[i] = load_dequantize_k_v_group(  # noqa: F821
+            q_i = tl.load(  # noqa: F821
+                tl.advance(Q_block_ptr, (0, i * D_PER_GROUP)), boundary_check=(0,)
+            )
+            k_i = load_dequantize_k_group(  # noqa: F821
                 K_block_ptr,
-                V_block_ptr,
                 K_scale_shift_block_ptr,
-                V_scale_shift_block_ptr,
                 BOUNDS_CHECKS_N,
                 PACKED_PER_VAL,
                 PACKED_D_PER_GROUP,
@@ -470,11 +463,7 @@ def _fwd_kernel_splitK(
                 i,
                 IS_HIP,
             )
-
-        # -- compute qk ---
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        for i in range(len(acc)):  # noqa: F821
-            qk += tl.dot(q[i], k[i])  # noqa: F821
+            qk += tl.dot(q_i, k_i)  # noqa: F821
         qk *= qk_scale
 
         if start_n == lo and ignore_in_first_block > 0:
@@ -526,10 +515,20 @@ def _fwd_kernel_splitK(
         m_i = m_i_new
         p = p.to(Q.dtype.element_ty)
 
-        # -- scale and update acc --
+        alpha_tile = alpha[:, None]
         for i in range(len(acc)):  # noqa: F821
-            acc[i] *= alpha[:, None]  # noqa: F821
-            acc[i] += tl.dot(p, v[i])  # noqa: F821
+            v_i = load_dequantize_v_group(  # noqa: F821
+                V_block_ptr,
+                V_scale_shift_block_ptr,
+                BOUNDS_CHECKS_N,
+                PACKED_PER_VAL,
+                PACKED_D_PER_GROUP,
+                FP8_QUANTIZED,
+                Q.dtype.element_ty,
+                i,
+                IS_HIP,
+            )
+            acc[i] = acc[i] * alpha_tile + tl.dot(p, v_i)  # noqa: F821
 
         if not PAGE_SIZE:
             # update pointers
@@ -646,13 +645,19 @@ def autotune_kernel(kernel: Callable):
     STAGES_VALUES = [1, 2] if torch.version.hip else [1, 2, 3]
     WARPS_VALUES = [1, 2, 4, 8]
 
+    def _is_valid_for_hip(block_m: int, block_n: int, warps: int) -> bool:
+        if not torch.version.hip:
+            return True
+        # Keep HIP configs under tight resource budgets to avoid >128 VGPR usage.
+        return block_m <= 64 and block_n <= 64 and warps <= 4
+
     TRITON_CONFIGS = [
         gen_config(block_m, block_n, stages, warps)
         for block_m in BLOCK_M_VALUES
         for block_n in BLOCK_N_VALUES
         for stages in STAGES_VALUES
         for warps in WARPS_VALUES
-        if block_n >= block_m
+        if block_n >= block_m and _is_valid_for_hip(block_m, block_n, warps)
     ]
 
     kernel = triton.autotune(
@@ -780,6 +785,96 @@ def load_dequantize_k_v_group(
     return k, v
 
 
+@triton.jit
+def load_dequantize_k_group(
+    K_block_ptr,
+    K_scale_shift_block_ptr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    PACKED_D_PER_GROUP: tl.constexpr,
+    FP8_QUANTIZED: tl.constexpr,
+    dtype: tl.constexpr,
+    group_id: tl.constexpr,
+    IS_HIP: tl.constexpr,
+):
+    K_group_ptr = tl.advance(K_block_ptr, (PACKED_D_PER_GROUP * group_id, 0))
+    k = tl.load(K_group_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ())
+
+    if FP8_QUANTIZED:
+        k_scale_shift = tl.load(
+            K_scale_shift_block_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ()
+        )
+        if IS_HIP:
+            k_scale, k_shift = cast_uint32_to_float(k_scale_shift)
+            k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype)
+        else:
+            k_scale, k_shift = cast_uint32_to_half2(k_scale_shift)
+            k_t = dequantize(
+                tl.trans(k),
+                tl.trans(k_scale),
+                tl.trans(k_shift),
+                PACKED_PER_VAL,
+                IS_HIP,
+            ).to(dtype)
+            k = tl.trans(k_t)
+    elif PACKED_PER_VAL > 1:
+        k_scale_shift_ptr = tl.advance(K_scale_shift_block_ptr, (group_id, 0))
+        k_scale_shift = tl.load(
+            k_scale_shift_ptr, boundary_check=(1,) if BOUNDS_CHECKS_N else ()
+        )
+        if IS_HIP:
+            k_scale, k_shift = cast_uint32_to_float(k_scale_shift)
+            k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype)
+        else:
+            k_scale, k_shift = cast_uint32_to_half2(k_scale_shift)
+            k_t = dequantize(
+                tl.trans(k),
+                tl.trans(k_scale),
+                tl.trans(k_shift),
+                PACKED_PER_VAL,
+                IS_HIP,
+            ).to(dtype)
+            k = tl.trans(k_t)
+    return k
+
+
+@triton.jit
+def load_dequantize_v_group(
+    V_block_ptr,
+    V_scale_shift_block_ptr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr,
+    PACKED_D_PER_GROUP: tl.constexpr,
+    FP8_QUANTIZED: tl.constexpr,
+    dtype: tl.constexpr,
+    group_id: tl.constexpr,
+    IS_HIP: tl.constexpr,
+):
+    V_group_ptr = tl.advance(V_block_ptr, (0, PACKED_D_PER_GROUP * group_id))
+    v = tl.load(V_group_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ())
+
+    if FP8_QUANTIZED:
+        v_scale_shift = tl.load(
+            V_scale_shift_block_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ()
+        )
+        if IS_HIP:
+            v_scale, v_shift = cast_uint32_to_float(v_scale_shift)
+        else:
+            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)
+    elif PACKED_PER_VAL > 1:
+        v_scale_shift_ptr = tl.advance(V_scale_shift_block_ptr, (0, group_id))
+        v_scale_shift = tl.load(
+            v_scale_shift_ptr, boundary_check=(0,) if BOUNDS_CHECKS_N else ()
+        )
+        if IS_HIP:
+            v_scale, v_shift = cast_uint32_to_float(v_scale_shift)
+        else:
+            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)
+    return v
+
+
 @triton.jit
 def cast_uint32_to_half2(scale_shift):
     """Extract two float16 packed into one int32"""
@@ -873,8 +968,21 @@ def dequantize(
         quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)
     )
     if PACKED_PER_VAL == 4:
+        if IS_HIP:
+            # Reuse the HIP-specific FP8 routine that keeps the unpacked tile transient.
+            dequant = dequantize_k_hip(
+                tl.trans(x_), tl.trans(scale), tl.trans(shift), PACKED_PER_VAL
+            )
+            return tl.trans(dequant)
         # FP8 quantization.
-        fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv
+        fp8_type = (
+            tl.float8e4b8
+            if (
+                torch.version.hip is not None
+                and triton.runtime.driver.active.get_current_target().arch == "gfx942"
+            )
+            else tl.float8e4nv
+        )
         dequant = (
             quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale
             + shift

xformers/ops/fmha/triton_splitk.py

@@ -5,6 +5,7 @@
 
 import functools
 import sys
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -247,6 +248,19 @@ class FwOp(AttentionFwOpBase):
     # On AMD or for M > 1 different NUM_STAGES and NUM_WARPS can be used.
     NUM_STAGES: int = 1
     NUM_WARPS: int = 2
+    FORCED_CONFIG: Optional[Any] = None
+
+    @classmethod
+    @contextmanager
+    def force_kernel_config(cls, config: Optional[Any]):
+        """Temporarily override the Triton launch configuration used by the forward kernel."""
+
+        previous = cls.FORCED_CONFIG
+        cls.FORCED_CONFIG = config
+        try:
+            yield
+        finally:
+            cls.FORCED_CONFIG = previous
 
     @classmethod
     def shape_not_supported_reasons(
@@ -425,6 +439,23 @@ def get_extra_args(
         attn_bias: Any,
         k_fp8_scale_shift: Any,
     ) -> Dict[str, Any]:
+        forced_config = cls.FORCED_CONFIG
+        if forced_config is not None:
+            return {
+                "BLOCK_M": forced_config.kwargs["BLOCK_M"],
+                "BLOCK_N": forced_config.kwargs["BLOCK_N"],
+                "num_warps": forced_config.num_warps,
+                "num_stages": forced_config.num_stages,
+            }
+
+        if torch.version.hip and k_fp8_scale_shift is not None:
+            return {
+                "BLOCK_M": 16,
+                "BLOCK_N": 64,
+                "num_warps": 1,
+                "num_stages": 2,
+            }
+
         BLOCK_M = cls.BLOCK_M
         BLOCK_N = cls.BLOCK_N
         if cls.AUTOTUNE: