feat(ck-tile): enable grouped AQuant/ABQuant/BQuant GEMM bridges on gfx1250 (MI400)

projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp

@@ -79,10 +79,10 @@ int dispatcher_initialize()
     }
     g_gfx_arch = props.gcnArchName;
     if(g_gfx_arch.rfind("gfx950", 0) != 0 && g_gfx_arch.rfind("gfx942", 0) != 0 &&
-       g_gfx_arch.rfind("gfx90a", 0) != 0)
+       g_gfx_arch.rfind("gfx90a", 0) != 0 && g_gfx_arch.rfind("gfx1250", 0) != 0)
     {
         std::cerr << "dispatcher_initialize: unsupported GPU architecture '" << g_gfx_arch
-                  << "' (supported: gfx90a, gfx942, gfx950)\n";
+                  << "' (supported: gfx90a, gfx942, gfx950, gfx1250)\n";
         return -1;
     }
     g_initialized = true;

projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp

@@ -77,10 +77,10 @@ int dispatcher_initialize()
     }
     g_gfx_arch = props.gcnArchName;
     if(g_gfx_arch.rfind("gfx950", 0) != 0 && g_gfx_arch.rfind("gfx942", 0) != 0 &&
-       g_gfx_arch.rfind("gfx90a", 0) != 0)
+       g_gfx_arch.rfind("gfx90a", 0) != 0 && g_gfx_arch.rfind("gfx1250", 0) != 0)
     {
         std::cerr << "dispatcher_initialize: unsupported GPU architecture '" << g_gfx_arch
-                  << "' (supported: gfx90a, gfx942, gfx950)\n";
+                  << "' (supported: gfx90a, gfx942, gfx950, gfx1250)\n";
         return -1;
     }
     g_initialized = true;

projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp

@@ -145,10 +145,10 @@ int dispatcher_run_bquant_gemm(const void* A,
         }
         const std::string arch(props.gcnArchName);
         if(arch.rfind("gfx950", 0) != 0 && arch.rfind("gfx942", 0) != 0 &&
-           arch.rfind("gfx90a", 0) != 0)
+           arch.rfind("gfx90a", 0) != 0 && arch.rfind("gfx1250", 0) != 0)
         {
             std::cerr << "dispatcher_run_bquant_gemm: unsupported GPU architecture '" << arch
-                      << "' (supported: gfx90a, gfx942, gfx950)\n";
+                      << "' (supported: gfx90a, gfx942, gfx950, gfx1250)\n";
             return -1;
         }
     }

projects/composablekernel/dispatcher/codegen/arch_specs_generated.py

@@ -109,6 +109,8 @@
     "gfx1250": {
         "fp16_fp16_fp32": [[16, 16, 32]],
         "bf16_bf16_fp32": [[16, 16, 32]],
+        "fp8_fp8_fp32": [[16, 16, 64]],
+        "bf8_bf8_fp32": [[16, 16, 64]],
     },
 }
 

projects/composablekernel/dispatcher/python/gemm_utils.py

@@ -50,7 +50,7 @@
 
 # --- Bridge shared helpers (canonical superset; byte-identical across bridges) ---
 # Supported GPU architectures for the bridge (single source of truth).
-_SUPPORTED_ARCHES = ("gfx90a", "gfx942", "gfx950")
+_SUPPORTED_ARCHES = ("gfx90a", "gfx942", "gfx950", "gfx1250")
 
 # Single source of truth for the preshuffle B-shuffle permutation used by the
 # bridge. The bridge codegen only emits the NON-permuteN preshuffle pipeline
@@ -968,6 +968,33 @@ def _bf8_u8_to_fp32(u8: np.ndarray) -> np.ndarray:
     return _fnuz_decode_table(5, 2)[u8.astype(np.intp)]
 
 
+def _fp32_to_fp8_ocp_u8(x):
+    """Encode fp32 -> fp8 E4M3 **OCP** (gfx950/gfx12xx) bit pattern (uint8)."""
+    import ml_dtypes
+    return np.ascontiguousarray(x, dtype=np.float32).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)
+
+
+def _fp32_to_bf8_ocp_u8(x):
+    """Encode fp32 -> bf8 E5M2 **OCP** (gfx950/gfx12xx) bit pattern (uint8)."""
+    import ml_dtypes
+    return np.ascontiguousarray(x, dtype=np.float32).astype(ml_dtypes.float8_e5m2).view(np.uint8)
+
+
+_OCP_FP8_CACHE = {}
+
+
+def _use_ocp_fp8():
+    """True on archs whose device fp8_t is OCP (gfx950/MI350, gfx12xx/RDNA) rather
+    than the FNUZ encoding used by gfx942/MI300. Detected once via rocminfo."""
+    if "v" not in _OCP_FP8_CACHE:
+        try:
+            a = _get_arch()
+        except Exception:
+            a = ""
+        _OCP_FP8_CACHE["v"] = a.startswith("gfx12") or a == "gfx950"
+    return _OCP_FP8_CACHE["v"]
+
+
 # Output (C) element dtype for an A/B element dtype, mirroring the codegen's
 # CommonTypeMappings.get_output_dtype: fp8/bf8 accumulate into fp16, int8 into
 # int32, everything else stores in its own dtype.
@@ -1047,11 +1074,13 @@ def run(
             A_h = _fp32_to_bf16_u16(A_lay)
             B_h = _fp32_to_bf16_u16(B_lay)
         elif dtype == "fp8":
-            A_h = _fp32_to_fp8_u8(A_lay)
-            B_h = _fp32_to_fp8_u8(B_lay)
+            _enc = _fp32_to_fp8_ocp_u8 if _use_ocp_fp8() else _fp32_to_fp8_u8
+            A_h = _enc(A_lay)
+            B_h = _enc(B_lay)
         elif dtype == "bf8":
-            A_h = _fp32_to_bf8_u8(A_lay)
-            B_h = _fp32_to_bf8_u8(B_lay)
+            _enc = _fp32_to_bf8_ocp_u8 if _use_ocp_fp8() else _fp32_to_bf8_u8
+            A_h = _enc(A_lay)
+            B_h = _enc(B_lay)
         elif dtype == "int8":
             A_h = np.ascontiguousarray(A_lay, dtype=np.int8)
             B_h = np.ascontiguousarray(B_lay, dtype=np.int8)
@@ -2130,6 +2159,7 @@ def _dim_bad(repeat: int, warp_tile: int) -> bool:
     "gfx90a": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
     "gfx942": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
     "gfx950": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
+    "gfx1250": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1], [2, 1, 1], [1, 2, 2], [4, 1, 1], [1, 4, 1], [2, 2, 1]],
     "gfx1201": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1]],
 }
 

projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py

@@ -857,3 +857,78 @@ def default_bf8_preshuffleb_config(
         k_block_per_cu=2,
         gfx_arch=gfx_arch,
     )
+
+
+# =============================================================================
+# gfx1250 (MI400) default configs
+# =============================================================================
+# gfx1250 uses the standard CompV3 pipeline (NOT the gfx950-native eightwaves /
+# preshuffleb paths, which require CK_GFX950_SUPPORT and TransposeC=True).
+#
+# Empirically (on-GPU, MI400/gfx1250, HIP_VISIBLE_DEVICES=0, fp32-dequant
+# reference), the ABQuant CompV3 fp8/bf8 kernels are correct on gfx1250 only
+# with warp_tile_k=128 (the FlatMM tile). The gfx9 stock CompV3 default uses
+# warp_tile_k=32 (plain MFMA) which returns all-zeros on gfx1250; warp_tile_k=16
+# (gfx12 WMMA) also returns all-zeros. So the only correct gfx1250 override for
+# the CompV3 pipeline is 16x16x128 with transpose_c=False.
+#
+#   fp8: warp_tile 16x16x128 -> nonzero, max_rel_err ~3e-4  (PASS)
+#   bf8: warp_tile 16x16x128 -> nonzero, max_rel_err ~3e-4  (PASS)
+#
+# The compile path injects -DCK_USE_OCP_FP8 / -DCK_TILE_USE_OCP_FP8 for gfx12
+# archs, so fp8/bf8 use the OCP encoding on gfx1250.
+
+_GFX1250_ARCH = "gfx1250"
+
+
+def default_fp8_compv3_config_gfx1250(
+    quant_group_k: int = 128,
+    bquant_group_n: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> ABQuantKernelConfig:
+    """fp8 ABQuant CompV3 config for gfx1250 (FlatMM warp_tile_k=128, TransposeC=False).
+
+    GPU-verified on MI400/gfx1250: max_rel_err ~3e-4 vs. fp32 dequant reference.
+    warp_tile_k=32 (gfx9 stock) and warp_tile_k=16 (gfx12 WMMA) both zero out.
+    """
+    return ABQuantKernelConfig(
+        variant_key="fp8",
+        layout="rcr",
+        pipeline="compv3",
+        epilogue="cshuffle",
+        scheduler="intrawave",
+        tile_m=128, tile_n=128, tile_k=128,
+        warp_m=1, warp_n=4, warp_k=1,
+        warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
+        aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k,
+        bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k,
+        preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False,
+        transpose_c=False,
+        gfx_arch=gfx_arch,
+    )
+
+
+def default_bf8_compv3_config_gfx1250(
+    quant_group_k: int = 128,
+    bquant_group_n: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> ABQuantKernelConfig:
+    """bf8 ABQuant CompV3 config for gfx1250 (FlatMM warp_tile_k=128, TransposeC=False).
+
+    GPU-verified on MI400/gfx1250: max_rel_err ~3e-4 vs. fp32 dequant reference.
+    """
+    return ABQuantKernelConfig(
+        variant_key="bf8",
+        layout="rcr",
+        pipeline="compv3",
+        epilogue="cshuffle",
+        scheduler="intrawave",
+        tile_m=128, tile_n=128, tile_k=128,
+        warp_m=1, warp_n=4, warp_k=1,
+        warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
+        aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k,
+        bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k,
+        preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False,
+        transpose_c=False,
+        gfx_arch=gfx_arch,
+    )

projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py

@@ -785,3 +785,81 @@ def default_bf8_preshuffleaq_config(
         preshuffle_aq=True,
         gfx_arch=gfx_arch,
     )
+
+
+# =============================================================================
+# gfx1250 (MI400) default configs
+# =============================================================================
+# Empirically (on-GPU, MI400/gfx1250, HIP_VISIBLE_DEVICES=0, fp32-dequant
+# reference), the AQuant fp8/bf8 kernels are correct on gfx1250 only with
+# warp_tile_k=128 (the FlatMM path, same tile the gfx9 fp8/bf8 FlatMM decode
+# uses). The gfx9 stock AQuant default uses warp_tile_k=32 (plain MFMA), which
+# produces a WRONG result on gfx1250 (max_rel_err=1.0); warp_tile_k=16 (gfx12
+# WMMA) returns all-zeros. So the only correct gfx1250 override is 16x16x128.
+#
+#   fp8: warp_tile 16x16x128 -> nonzero, max_rel_err ~3e-4  (PASS)
+#   bf8: warp_tile 16x16x128 -> nonzero, max_rel_err ~1e-4  (PASS)
+#
+# The compile path injects -DCK_USE_OCP_FP8 / -DCK_TILE_USE_OCP_FP8 for gfx12
+# archs, so fp8/bf8 use the OCP encoding on gfx1250.
+#
+# fp8i4 / bf8i4 (pk_int4 weights) do NOT compile on gfx1250 at any warp_tile_k
+# (no gfx12 WMMA/FlatMM instruction for the packed-int4 quant path). Those
+# variants are therefore not enabled on gfx1250: the C++ ctypes arch-gate is
+# opened, but on-GPU correctness is deferred pending kernel work. No gfx1250
+# config helper is provided for i4 because it would emit a non-buildable kernel.
+
+_GFX1250_ARCH = "gfx1250"
+
+
+def default_fp8_config_gfx1250(
+    quant_group_k: int = 128,
+    quant_group_m: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> AQuantKernelConfig:
+    """fp8 AQuant decode config for gfx1250 (FlatMM, warp_tile_k=128).
+
+    GPU-verified on MI400/gfx1250: max_rel_err ~3e-4 vs. fp32 dequant reference.
+    warp_tile_k=32 (gfx9 stock) is WRONG on gfx1250; warp_tile_k=16 zeros out.
+    """
+    return AQuantKernelConfig(
+        variant_key="fp8",
+        layout="rcr",
+        pipeline="mem",
+        epilogue="cshuffle",
+        scheduler="intrawave",
+        tile_m=16, tile_n=64, tile_k=256,
+        warp_m=1, warp_n=4, warp_k=1,
+        warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
+        quant_group_m=quant_group_m,
+        quant_group_n=1,
+        quant_group_k=quant_group_k,
+        preshuffle_aq=False,
+        gfx_arch=gfx_arch,
+    )
+
+
+def default_bf8_config_gfx1250(
+    quant_group_k: int = 128,
+    quant_group_m: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> AQuantKernelConfig:
+    """bf8 AQuant decode config for gfx1250 (FlatMM, warp_tile_k=128).
+
+    GPU-verified on MI400/gfx1250: max_rel_err ~1e-4 vs. fp32 dequant reference.
+    """
+    return AQuantKernelConfig(
+        variant_key="bf8",
+        layout="rcr",
+        pipeline="mem",
+        epilogue="cshuffle",
+        scheduler="intrawave",
+        tile_m=16, tile_n=64, tile_k=256,
+        warp_m=1, warp_n=4, warp_k=1,
+        warp_tile_m=16, warp_tile_n=16, warp_tile_k=128,
+        quant_group_m=quant_group_m,
+        quant_group_n=1,
+        quant_group_k=quant_group_k,
+        preshuffle_aq=False,
+        gfx_arch=gfx_arch,
+    )

projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py

@@ -1214,3 +1214,62 @@ def default_mx_bf16fp4_config(
         quant_group_k=quant_group_k,
         gfx_arch=gfx_arch,
     )
+
+
+# =============================================================================
+# gfx1250 (MI400) default configs
+# =============================================================================
+# Empirically (on-GPU, MI400/gfx1250, HIP_VISIBLE_DEVICES=0, via the
+# test_bquant_gpu_correctness fp32-dequant reference), the BQuant fp8/bf8 kernels
+# are correct on gfx1250 with the SAME config that gfx9 uses: the CompV3 FlatMM
+# tile with warp_tile_k=128. warp_tile_k=16 (gfx12 WMMA) returns all-zeros on
+# gfx1250. So the gfx1250 helpers simply return the stock fp8/bf8 config
+# (warp_tile_k=128); no distinct gfx1250 tile is required for BQuant.
+#
+#   fp8: warp_tile 16x16x128 -> nonzero, max_rel_err ~5e-4  (PASS)
+#   bf8: warp_tile 16x16x128 -> nonzero, max_rel_err ~5e-4  (PASS)
+#
+# The compile path injects -DCK_USE_OCP_FP8 / -DCK_TILE_USE_OCP_FP8 for gfx12
+# archs, so fp8/bf8 use the OCP encoding on gfx1250.
+#
+# fp8i4 / bf8i4 (pk_int4 weights) do NOT compile on gfx1250 at any warp_tile_k
+# (no gfx12 instruction for the packed-int4 quant path). Those variants are not
+# enabled on gfx1250: the C++ ctypes arch-gate is opened, but on-GPU correctness
+# is deferred pending kernel work. No gfx1250 config helper is provided for i4
+# because it would emit a non-buildable kernel.
+
+_GFX1250_ARCH = "gfx1250"
+
+
+def default_fp8_config_gfx1250(
+    quant_group_k: int = 128,
+    quant_group_n: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> BQuantKernelConfig:
+    """fp8 BQuant config for gfx1250 (compv3, FlatMM warp_tile_k=128).
+
+    Identical to the stock fp8 config; GPU-verified on MI400/gfx1250 with
+    max_rel_err ~5e-4 vs. fp32 dequant reference. warp_tile_k=16 zeros out.
+    """
+    return default_fp8_config(
+        quant_group_k=quant_group_k,
+        quant_group_n=quant_group_n,
+        gfx_arch=gfx_arch,
+    )
+
+
+def default_bf8_config_gfx1250(
+    quant_group_k: int = 128,
+    quant_group_n: int = 1,
+    gfx_arch: str = _GFX1250_ARCH,
+) -> BQuantKernelConfig:
+    """bf8 BQuant config for gfx1250 (compv3, FlatMM warp_tile_k=128).
+
+    Identical to the stock bf8 config; GPU-verified on MI400/gfx1250 with
+    max_rel_err ~5e-4 vs. fp32 dequant reference.
+    """
+    return default_bf8_config(
+        quant_group_k=quant_group_k,
+        quant_group_n=quant_group_n,
+        gfx_arch=gfx_arch,
+    )