feat(ck-tile): add dispatcher correctness and performance CI lanes for gfx942/gfx950

projects/composablekernel/Dockerfile

@@ -100,7 +100,15 @@ RUN set -x && \
     wget https://github.com/Yelp/dumb-init/releases/download/v1.2.0/dumb-init_1.2.0_amd64.deb && \
     dpkg -i dumb-init_*.deb && rm dumb-init_*.deb && \
 # Install packages for processing the performance results
-    pip3 install --break-system-packages --upgrade pytest pymysql pandas==2.2.3 sqlalchemy==2.0.3 setuptools-rust setuptools sshtunnel==0.4.0 && \
+# ml_dtypes is pinned like its neighbours, and for a sharper reason than most:
+# it supplies both float8_e4m3fn/float8_e5m2 (OCP) and float8_e4m3fnuz/
+# float8_e5m2fnuz (FNUZ), and the dispatcher tests pick between those pairs per
+# arch. An unpinned upgrade would silently move fp8 encoding semantics -- and
+# the reference values every parity/correctness test compares against -- from
+# under CI on an image rebuild. 0.6.0 is what this image already resolved on
+# Ubuntu 24.04 / Python 3.12, so this records current behaviour rather than
+# bumping it.
+    pip3 install --break-system-packages --upgrade pytest pymysql pandas==2.2.3 sqlalchemy==2.0.3 setuptools-rust setuptools sshtunnel==0.4.0 ml_dtypes==0.6.0 && \
 # Add render group
     groupadd -f render && \
 # Install the new rocm-cmake version

projects/composablekernel/dispatcher/python/batched_contraction_utils.py

@@ -693,11 +693,40 @@ def expand_sweep(config: dict, dtype: str = "fp16", layout: str = "rcr") -> List
     return out
 
 
-def default_fp16_config(gfx_arch: Optional[str] = None) -> BatchedContractionKernelConfig:
+# Warp tile per dtype, chosen from the XDL allow-list in is_valid(). fp16 and
+# bf16 share the same MFMA shape; fp32 has no 32x32x16 instruction and needs
+# 32x32x8 instead. Picking a shape outside the allow-list is not an error you
+# would notice -- is_valid() just returns False and the config is silently
+# dropped from the sweep (test_batched_contraction_bridge.py asserts exactly
+# that), so a "default" config with the wrong tile would look like a dtype that
+# simply has no instances.
+_DEFAULT_WARP_TILE = {
+    "fp16": (32, 32, 16),
+    "bf16": (32, 32, 16),
+    "fp32": (32, 32, 8),
+}
+
+
+def default_config(
+    dtype: str = "fp16", gfx_arch: Optional[str] = None
+) -> BatchedContractionKernelConfig:
+    """A known-valid single config for ``dtype``, for smoke tests and defaults."""
+    try:
+        wt_m, wt_n, wt_k = _DEFAULT_WARP_TILE[dtype]
+    except KeyError:
+        raise ValueError(
+            f"no default warp tile for dtype {dtype!r}; "
+            f"known: {', '.join(sorted(_DEFAULT_WARP_TILE))}"
+        ) from None
     gfx_arch = _validate_arch(gfx_arch) if gfx_arch else _get_arch()
     return BatchedContractionKernelConfig(
-        dtype="fp16", layout="rcr", pipeline="compv3", epilogue="cshuffle", scheduler="intrawave",
+        dtype=dtype, layout="rcr", pipeline="compv3", epilogue="cshuffle", scheduler="intrawave",
         tile_m=128, tile_n=128, tile_k=64, warp_m=2, warp_n=2, warp_k=1,
-        warp_tile_m=32, warp_tile_n=32, warp_tile_k=16,
+        warp_tile_m=wt_m, warp_tile_n=wt_n, warp_tile_k=wt_k,
         num_dim_g=1, num_dim_m=1, num_dim_n=1, num_dim_k=1, gfx_arch=gfx_arch,
     )
+
+
+def default_fp16_config(gfx_arch: Optional[str] = None) -> BatchedContractionKernelConfig:
+    """Thin alias kept for existing callers; prefer :func:`default_config`."""
+    return default_config("fp16", gfx_arch=gfx_arch)

projects/composablekernel/dispatcher/python/dispatcher_common.py

@@ -101,6 +101,120 @@ def detect_gpu_arch(fallback: str = "gfx942") -> str:
     return fallback
 
 
+# ============================================================================
+# fp8 / bf8 encoding format per architecture
+# ============================================================================
+#
+# CK has two incompatible 8-bit float encodings and picks between them per arch:
+#
+#   OCP  (gfx950, gfx12): e4m3fn   / e5m2    -- exponent bias 7 / 15
+#   FNUZ (everything else, notably gfx942): e4m3fnuz / e5m2fnuz -- bias 8 / 16
+#
+# include/ck_tile/core/config.hpp:361-371 resolves this at compile time, but only
+# the *device* pass sees __gfx950__; the host pass of the very same header falls
+# back to FNUZ. So a host-side reference encoder that guesses from the header ends
+# up disagreeing with the kernel it is validating. Every caller must therefore
+# decide from the target arch string, and pass -DCK_TILE_USE_OCP_FP8 explicitly so
+# both compiler passes agree. These helpers are the single source of truth for
+# both halves of that contract.
+
+
+def normalize_arch(arch: Optional[str]) -> str:
+    """Lowercase `arch` and strip the target-feature suffix, if any.
+
+    rocminfo, hipcc and the HSA runtime all hand back full target triples --
+    ``gfx950:sramecc+:xnack-`` -- while configs, CI parameters and these helpers
+    are written in terms of the bare name.  Every arch predicate below matches on
+    a prefix, so a triple happens to survive unnormalized, but anything carrying
+    a leading component (an ``amdgcn-amd-amdhsa--gfx950`` offload target, say)
+    silently falls through to the FNUZ default.  That miss drops
+    ``-DCK_USE_OCP_FP8``, and host and device then disagree on the fp8 encoding
+    with no diagnostic at all -- the kernel builds and returns wrong numbers.
+
+    Normalizing in one place, up front, is what keeps that from depending on
+    which spelling of the arch a given caller happened to be handed.
+    """
+    a = (arch or "").lower().strip()
+    # Feature suffix: gfx950:sramecc+:xnack- -> gfx950
+    a = a.split(":", 1)[0]
+    # Offload-target prefix: amdgcn-amd-amdhsa--gfx950 -> gfx950
+    idx = a.rfind("gfx")
+    return a[idx:] if idx > 0 else a
+
+
+def fp8_uses_ocp(arch: Optional[str]) -> bool:
+    """True iff `arch` uses the OCP fp8/bf8 encoding rather than FNUZ.
+
+    Mirrors the __gfx950__ / __gfx12__ test in include/ck_tile/core/config.hpp.
+    Use this to select the host-side codec (ml_dtypes.float8_e4m3fn vs
+    float8_e4m3fnuz) so it matches the bytes the kernel actually produces.
+
+    Accepts bare names and full target triples alike; see `normalize_arch`.
+    """
+    a = normalize_arch(arch)
+    return a.startswith("gfx950") or a.startswith("gfx12")
+
+
+def ocp_arch_defines(arch: Optional[str]) -> List[str]:
+    """hipcc defines that pin the fp8/bf8 encoding for `arch`.
+
+    Returned for OCP archs only; FNUZ is the default for both compiler passes and
+    needs no define. Passing these makes the host pass agree with the device pass
+    instead of silently falling back to FNUZ.
+    """
+    if not fp8_uses_ocp(arch):
+        return []
+    return ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"]
+
+
+def arch_feature_defines(arch: Optional[str]) -> List[str]:
+    """`ocp_arch_defines` plus the per-arch feature-enablement defines.
+
+    The top-level CMakeLists.txt (:456-512) sets these for a normal build; they
+    are absent in the standalone hipcc JIT path, so any bridge that compiles a
+    kernel out-of-tree has to re-supply them.  This mirrors that block for the
+    arches the dispatcher targets:
+
+        gfx950  -> CK_USE_NATIVE_MX_SUPPORT, CK_GFX950_SUPPORT
+        gfx1250 -> CK_USE_GFX1250, CK_USE_NATIVE_MX_SUPPORT, CK_GFX1250_SUPPORT
+        gfx11/gfx12 -> CK_TILE_USE_WMMA=1  (gfx12 also CK_GFX12_SUPPORT)
+
+    CK_TILE_USE_WMMA is the one that must be passed even when it is 0: CMake
+    always defines it (:480), and the JIT path leaving it undefined only happens
+    to work because the preprocessor reads an undefined identifier as 0, which is
+    the right answer on gfx942/gfx950 and the wrong one on every WMMA part.
+
+    CK_USE_GFX950 is deliberately *not* emitted -- CMake sets it, but it is read
+    only by the legacy ck/library conv instances, never by ck_tile, so it is
+    dead weight on this path.
+
+    On the warp tile: these defines do select the branch of
+    tile_gemm_shape.hpp get_k_warp_tile() (CK_TILE_USE_WMMA outermost, then
+    CK_USE_GFX1250 / CK_GFX950_SUPPORT), but the emitted kernels never *call*
+    that function -- codegen writes a literal warp_tile_k, mirroring it in Python
+    via codegen_common.fp8_warp_tile_k_for_arch.  The hazard is therefore a
+    mismatch between that literal and the warp-gemm the target arch actually has:
+    there is no 16x16x128 fp8 warp-gemm on gfx942, and asking for one compiles
+    cleanly and returns all zeros.  Keep warp_tile_k arch-aware alongside these
+    defines, and derive it from the one helper rather than a local copy.
+    """
+    a = normalize_arch(arch)
+    defines = ocp_arch_defines(arch)
+    if a.startswith("gfx950"):
+        defines = defines + ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"]
+    elif a.startswith("gfx11") or a.startswith("gfx12"):
+        defines = defines + ["-DCK_TILE_USE_WMMA=1"]
+        if a.startswith("gfx12"):
+            defines = defines + ["-DCK_GFX12_SUPPORT"]
+        if a.startswith("gfx1250"):
+            defines = defines + [
+                "-DCK_USE_GFX1250",
+                "-DCK_USE_NATIVE_MX_SUPPORT",
+                "-DCK_GFX1250_SUPPORT",
+            ]
+    return defines
+
+
 # ============================================================================
 # Architecture Filter Data
 # ============================================================================