From 630d76f1dc968ab7ac1b771cd6457a2e62a96036 Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Mon, 13 Jul 2026 04:12:49 -0400 Subject: [PATCH 01/14] feat(ck-tile): mx_gemm GEMM TE to dispatcher bridge Summary: Adds a microscaling-GEMM (mx_gemm) bridge from the Old Tile-Engine into the dispatcher's ctypes path for gfx950/MI350. fp4 (pk_fp4_t) and fp8 (e4m3), rcr layout, comp_async + cshuffle + intrawave, fixed 16x16x128 warp tile, k_batch==1. Motivation: Extend the TE->Dispatcher bridge family to the microscaling GEMM op so the dispatcher can launch byte-identical mx_gemm kernels with per-32-K e8m0 block scales. Design note: Codegen reuses Old-TE MxGemmKernelBuilder._generate_kernel_instance directly (no re-implemented header assembly) to guarantee byte-exact kernels. The ctypes lib bypasses the generic registry and calls SelectedKernel::launch() with a hand-built MxGemmHostArgs, mirroring mx_gemm_profiler.hpp for the gfx950 scale pre-shuffle. Byte accounting uses HostTensor get_element_space_size_in_bytes() so fp4 (PackedSize==2) and fp8 are both packing-correct. Changes: - dispatcher/codegen/unified_mx_gemm_codegen.py: per-config .hpp generator - dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp: direct-launch C API - dispatcher/python/mx_gemm_utils.py: config/ctypes/runner + build pipeline - dispatcher/bindings/ctypes/CMakeLists.txt: dispatcher_mx_gemm_lib target --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 46 ++ .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 304 +++++++++ .../codegen/unified_mx_gemm_codegen.py | 268 ++++++++ .../dispatcher/python/mx_gemm_utils.py | 602 ++++++++++++++++++ 4 files changed, 1220 insertions(+) create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp create mode 100644 projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py create mode 100644 projects/composablekernel/dispatcher/python/mx_gemm_utils.py diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 18314017f225..7c7b443914ec 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -9,6 +9,7 @@ # # Targets: # - dispatcher_gemm_lib : GEMM dispatcher library +# - dispatcher_mx_gemm_lib : MX-GEMM (microscaling) dispatcher library (gfx950) # - dispatcher_conv_lib : Convolution dispatcher library (forward + bwd_data) # - dispatcher_conv_bwdw_lib : Convolution backward weight library # - gpu_helper : GPU helper executable for Python @@ -71,6 +72,51 @@ else() target_link_libraries(dispatcher_gemm_lib PRIVATE hip::device) endif() +# ============================================================================= +# MX-GEMM ctypes Library (microscaling GEMM, gfx950/MI350 only) +# ============================================================================= +# +# Force-includes a generated mx_gemm kernel header (SelectedKernel/KERNEL_NAME/ +# MxGemmHostArgs/ScaleType/...). The lib also includes the Old-TE common helper +# (common/utils.hpp) for the free-function is_row_major(Layout), so add the +# tile_engine/ops include roots. + +set(MX_GEMM_TE_INCLUDE_DIRS + ${PROJECT_SOURCE_DIR}/tile_engine/ops + ${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm + ${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm/mx_gemm +) + +file(GLOB MX_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/mx_gemm_*.hpp") +if(MX_GEMM_KERNEL_HEADERS) + list(GET MX_GEMM_KERNEL_HEADERS 0 MX_GEMM_KERNEL_HEADER) + message(STATUS "Found MX-GEMM kernel for ctypes lib: ${MX_GEMM_KERNEL_HEADER}") + + add_ctypes_library(dispatcher_mx_gemm_lib + mx_gemm_ctypes_lib.cpp + KERNEL_HEADER ${MX_GEMM_KERNEL_HEADER} + ) + target_include_directories(dispatcher_mx_gemm_lib PRIVATE ${MX_GEMM_TE_INCLUDE_DIRS}) + # The generated header only exports SelectedKernel/KERNEL_NAME/ScaleType/... + # under this define, so the force-include path requires it. + target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + ) +else() + message(STATUS "No MX-GEMM kernel found for ctypes lib - building without kernel") + add_library(dispatcher_mx_gemm_lib SHARED mx_gemm_ctypes_lib.cpp) + target_include_directories(dispatcher_mx_gemm_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ${MX_GEMM_TE_INCLUDE_DIRS} + ) + target_link_libraries(dispatcher_mx_gemm_lib PRIVATE hip::device) + set_target_properties(dispatcher_mx_gemm_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 + ) +endif() + # ============================================================================= # Convolution ctypes Library (supports forward + bwd_data) # ============================================================================= diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp new file mode 100644 index 000000000000..7e795e104535 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -0,0 +1,304 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * MX-GEMM Dispatcher ctypes Library (TileEngine -> Dispatcher bridge). + * + * Provides a C API for Python ctypes integration. The kernel header is + * force-included at compile time: + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE mx_gemm_ctypes_lib.cpp + * + * Force-include defines (at global scope): SelectedKernel (with static constexpr + * TileM/TileN/TileK, WarpPerBlock_{M,N,K}, WarpTile{M,N,K} and + * static float launch(const MxGemmHostArgs&, const ck_tile::stream_config&)), + * KERNEL_NAME, ADataType/BDataType/CDataType/AccDataType/ScaleType(=e8m0_t), + * using MxGemmHostArgs = ck_tile::MxGemmHostArgs<1,1,0>, ALayout/BLayout/CLayout, + * and #include "ck_tile/ops/gemm.hpp". + * + * Registry bypass: microscaling GEMM's launch takes ck_tile::MxGemmHostArgs + * (with per-32-K e8m0 block scales that must be pre-shuffled for gfx950), which + * the generic dispatcher backend cannot express. So this lib builds the HostArgs + * from plain C arrays and calls SelectedKernel::launch() directly -- the same + * direct-launch pattern used by the batched/multi-D bridges. + * + * Memory model: host-pointer in. The lib owns device allocation/copy/free via + * ck_tile::HostTensor + ck_tile::DeviceMem (RAII), exactly like the Old-TE + * profiler. A/B/C byte sizes come from HostTensor::get_element_space_size_in_bytes(), + * which divides the logical element count by numeric_traits::PackedSize -- so + * fp8 (PackedSize==1) allocates M*K bytes while fp4 (pk_fp4_t PackedSize==2, two + * e2m1 elements per byte) allocates M*K/2 bytes. The incoming A/B host buffers + * are therefore expected PHYSICALLY packed ([M,K/PackedSize]/[N,K/PackedSize] + * bytes); strides passed to the kernel stay in LOGICAL element units. + * Layout is fixed by the compiled-in ALayout/BLayout/CLayout (rcr: A row-major + * [M,K], B col-major [K,N] == [N,K] storage, C row-major [M,N]). K % 32 == 0. + * v1 scope: k_batch == 1 (no split-K). + * + * The scale pre-shuffle mirrors mx_gemm_profiler.hpp exactly: pack params are + * derived from SelectedKernel tile dims at compile time (not hardcoded). + */ + +#include +#include +#include +#include +#include +#include +#include + +// Kernel header force-included via -include. Brings in ck_tile core + the +// ck_tile::MxGemmHostArgs type, SelectedKernel/KERNEL_NAME, and (via +// ck_tile/ops/gemm.hpp) the mx pipeline. We additionally pull the host +// helpers explicitly for HostTensor / preShuffleScaleBuffer_gfx950. +#include "ck_tile/host.hpp" +// Old-TE common helpers: provides the free-function template is_row_major(Layout) +// -> ck_tile::bool_constant<...>, matching the mx_gemm profiler usage. +#include "common/utils.hpp" + +#ifndef GFX_ARCH +#define GFX_ARCH "gfx950" +#endif + +static bool g_initialized = false; + +namespace { + +int env_int(const char* name, int fallback) +{ + const char* v = std::getenv(name); + if(!v) + return fallback; + return std::atoi(v); +} + +} // namespace + +extern "C" { + +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} +int dispatcher_init() { return dispatcher_initialize(); } + +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } +int dispatcher_get_kernel_count() { return 1; } + +void dispatcher_cleanup() { g_initialized = false; } + +/** + * Run microscaling GEMM: C[M,N] = (A[M,K] * scale_a) . (B[K,N] * scale_b). + * + * A, B, C are host pointers. scale_a/scale_b are raw e8m0 bytes, unpacked and + * unshuffled, shaped [M, K/32] and [N, K/32] row-major respectively. The lib + * builds HostTensors from them, runs the SAME preShuffleScaleBuffer_gfx950 path + * as the Old-TE profiler, uploads everything, and calls SelectedKernel::launch. + * + * Returns 0 ok, -1 HIP/bad-args, -2 kernel reports unsupported args. + */ +int dispatcher_run_mx_gemm(const void* A, + const void* B, + void* C, + const uint8_t* scale_a, + const uint8_t* scale_b, + int M, + int N, + int K, + int k_batch, + float* time_ms) +{ + if(!g_initialized) + { + std::cerr << "dispatcher_run_mx_gemm: not initialized\n"; + return -1; + } + if(!A || !B || !C || !scale_a || !scale_b) + { + std::cerr << "dispatcher_run_mx_gemm: null pointer\n"; + return -1; + } + if(M <= 0 || N <= 0 || K <= 0) + { + std::cerr << "dispatcher_run_mx_gemm: M,N,K must be > 0\n"; + return -1; + } + if(K % 32 != 0) + { + std::cerr << "dispatcher_run_mx_gemm: MX GEMM requires K to be a multiple of 32, got K=" + << K << "\n"; + return -1; + } + if(k_batch != 1) + { + // v1 scope: split-K (k_batch > 1) is not yet supported through this bridge. + std::cerr << "dispatcher_run_mx_gemm: only k_batch==1 is supported in v1, got " << k_batch + << "\n"; + return -1; + } + + const ALayout layout_a = ALayout{}; + const BLayout layout_b = BLayout{}; + const CLayout layout_c = CLayout{}; + + const ck_tile::index_t m = static_cast(M); + const ck_tile::index_t n = static_cast(N); + const ck_tile::index_t k = static_cast(K); + + // Packed default strides for the compiled-in layouts (rcr: A row-major, + // B col-major, C row-major), matching the Old-TE profiler. get_default_stride + // takes the row-major-ness as a compile-time bool_constant tag (4th arg). + const ck_tile::index_t stride_a = + static_cast(ck_tile::get_default_stride(m, k, 0, is_row_major(layout_a))); + const ck_tile::index_t stride_b = + static_cast(ck_tile::get_default_stride(k, n, 0, is_row_major(layout_b))); + const ck_tile::index_t stride_c = + static_cast(ck_tile::get_default_stride(m, n, 0, is_row_major(layout_c))); + + const ck_tile::index_t scale_k_size = k / 32; + + // ---- Scale pre-shuffle pack params, derived from SelectedKernel (compile + // time), exactly mirroring mx_gemm_profiler.hpp. ---- + constexpr ck_tile::index_t m_per_xdl = SelectedKernel::WarpTileM; + constexpr ck_tile::index_t n_per_xdl = SelectedKernel::WarpTileN; + constexpr ck_tile::index_t k_per_xdl = SelectedKernel::WarpTileK; + constexpr ck_tile::index_t m_iter_per_warp = + SelectedKernel::TileM / (SelectedKernel::WarpPerBlock_M * m_per_xdl); + constexpr ck_tile::index_t n_iter_per_warp = + SelectedKernel::TileN / (SelectedKernel::WarpPerBlock_N * n_per_xdl); + constexpr ck_tile::index_t k_iter_per_warp = SelectedKernel::TileK / k_per_xdl; + + constexpr ck_tile::index_t m_xdl_pack = + (m_iter_per_warp >= 2 && m_iter_per_warp % 2 == 0) ? 2 : 1; + constexpr ck_tile::index_t n_xdl_pack = + (n_iter_per_warp >= 2 && n_iter_per_warp % 2 == 0) ? 2 : 1; + constexpr ck_tile::index_t k_xdl_pack = + (k_iter_per_warp >= 2 && k_iter_per_warp % 2 == 0) ? 2 : 1; + + constexpr ck_tile::index_t xdl_mn_thread = SelectedKernel::WarpTileM; + constexpr ck_tile::index_t xdl_k_thread = 64 / xdl_mn_thread; + + // ---- Build unshuffled scale HostTensors from the incoming raw e8m0 bytes. + // scale_a: [M, K/32] row-major; scale_b: [N, K/32] row-major. ---- + ck_tile::HostTensor scale_a_host( + {static_cast(m), static_cast(scale_k_size)}, + {static_cast(scale_k_size), static_cast(1)}); + ck_tile::HostTensor scale_b_host( + {static_cast(n), static_cast(scale_k_size)}, + {static_cast(scale_k_size), static_cast(1)}); + + const std::size_t scale_a_count = static_cast(m) * scale_k_size; + const std::size_t scale_b_count = static_cast(n) * scale_k_size; + // e8m0_t wraps a single uint8 (raw biased exponent); construct 1:1 from the + // incoming raw byte via the explicit e8m0_t(raw_type) constructor. + for(std::size_t i = 0; i < scale_a_count; ++i) + scale_a_host.mData[i] = ScaleType(static_cast(scale_a[i])); + for(std::size_t i = 0; i < scale_b_count; ++i) + scale_b_host.mData[i] = ScaleType(static_cast(scale_b[i])); + + // ---- Shuffled scale buffers (same lengths as the profiler). ---- + ck_tile::HostTensor scale_a_shuffled( + {static_cast(m / m_xdl_pack * 2), + static_cast(scale_k_size / k_xdl_pack * 2)}, + {static_cast(scale_k_size / k_xdl_pack * 2), static_cast(1)}); + ck_tile::HostTensor scale_b_shuffled( + {static_cast(n / n_xdl_pack * 2), + static_cast(scale_k_size / k_xdl_pack * 2)}, + {static_cast(scale_k_size / k_xdl_pack * 2), static_cast(1)}); + + ck_tile::preShuffleScaleBuffer_gfx950( + scale_a_host.mData.data(), scale_a_shuffled.mData.data(), m, scale_k_size, true); + ck_tile::preShuffleScaleBuffer_gfx950( + scale_b_host.mData.data(), scale_b_shuffled.mData.data(), n, scale_k_size, true); + + // ---- Build A/B/C HostTensors with the SAME descriptors the Old-TE profiler + // uses, then allocate/copy through ck_tile::DeviceMem. This makes the byte + // accounting packing-correct for BOTH fp8 (PackedSize==1) and fp4 + // (pk_fp4_t::PackedSize==2, i.e. two logical e2m1 elements per byte): + // HostTensor::get_element_space_size_in_bytes() == + // sizeof(T) * (logical_elems / numeric_traits::PackedSize). + // The incoming A/B host buffers are already PHYSICALLY packed + // ([M, K/PackedSize] and [N, K/PackedSize] bytes for fp4; [M,K]/[N,K] for + // fp8), so we byte-copy them into mData (which is sized in physical objects). + // Strides remain in LOGICAL element units (get_default_stride above uses + // logical M,N,K), which is what the kernel expects. + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(m, k, stride_a, is_row_major(layout_a))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(k, n, stride_b, is_row_major(layout_b))); + ck_tile::HostTensor c_m_n( + ck_tile::host_tensor_descriptor(m, n, stride_c, is_row_major(layout_c))); + + const std::size_t a_bytes = a_m_k.get_element_space_size_in_bytes(); + const std::size_t b_bytes = b_k_n.get_element_space_size_in_bytes(); + const std::size_t c_bytes = c_m_n.get_element_space_size_in_bytes(); + + std::memcpy(a_m_k.mData.data(), A, a_bytes); + std::memcpy(b_k_n.mData.data(), B, b_bytes); + + float exec_time = 0.0f; + try + { + // ---- Device buffers, sized via HostTensor (packing-correct). ---- + ck_tile::DeviceMem a_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_dev_buf(c_m_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem scale_a_dev_buf(scale_a_shuffled.get_element_space_size_in_bytes()); + ck_tile::DeviceMem scale_b_dev_buf(scale_b_shuffled.get_element_space_size_in_bytes()); + + a_dev_buf.ToDevice(a_m_k.data()); + b_dev_buf.ToDevice(b_k_n.data()); + c_dev_buf.SetZero(); + scale_a_dev_buf.ToDevice(scale_a_shuffled.data()); + scale_b_dev_buf.ToDevice(scale_b_shuffled.data()); + + // ---- Build MxGemmHostArgs in the exact profiler argument order: + // a_ptr, scale_a, b_ptr, scale_b, ds{}, c_ptr, split_k, m, n, k, + // {stride_a}, {stride_b}, ds_strides{}, stride_c. ---- + MxGemmHostArgs gemm_args({a_dev_buf.GetDeviceBuffer()}, + {scale_a_dev_buf.GetDeviceBuffer()}, + {b_dev_buf.GetDeviceBuffer()}, + {scale_b_dev_buf.GetDeviceBuffer()}, + {}, + c_dev_buf.GetDeviceBuffer(), + static_cast(k_batch), + m, + n, + k, + {stride_a}, + {stride_b}, + {}, + stride_c); + + const bool do_time = (time_ms != nullptr); + const int warmup = do_time ? env_int("CK_TILE_BENCH_WARMUP", 20) : 0; + const int repeat = do_time ? env_int("CK_TILE_BENCH_REPEAT", 50) : 1; + // stream_config field order: stream_id, time_kernel, log_level, cold_niters + // (warmup), nrepeat, is_gpu_timer, flush_cache, rotating_count. + ck_tile::stream_config stream_cfg{nullptr, do_time, 0, warmup, repeat, false, false, 1}; + + exec_time = SelectedKernel::launch(gemm_args, stream_cfg); + if(exec_time < 0.0f) + { + std::cerr << "dispatcher_run_mx_gemm: kernel reports unsupported args\n"; + return -2; + } + + c_dev_buf.FromDevice(c_m_n.data()); + } + catch(const std::exception& e) + { + std::cerr << "dispatcher_run_mx_gemm: launch threw: " << e.what() << "\n"; + return -1; + } + + // C is always a non-packed output type (fp16/bf16), so c_bytes byte-copy back + // to the caller buffer is a 1:1 element copy. + std::memcpy(C, c_m_n.mData.data(), c_bytes); + + if(time_ms) + *time_ms = exec_time; + + return 0; +} + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py new file mode 100644 index 000000000000..fbcc569d3b08 --- /dev/null +++ b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py @@ -0,0 +1,268 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +MX GEMM code generator (TileEngine -> Dispatcher bridge). + +Emits ONE self-contained .hpp per concrete kernel config for the dispatcher's +ctypes path. To guarantee byte-identical kernels with Old-TE (the parity +contract), this generator does NOT re-implement the C++ header assembly -- it +directly reuses the Old-TE builder: + + tile_engine/ops/gemm/mx_gemm/mx_gemm_instance_builder.py + -> MxGemmKernelBuilder._generate_kernel_instance(tile_config, trait_combo) + +The generated instance defines, at GLOBAL scope, a ``SelectedKernel`` struct with +a static ``launch(const MxGemmHostArgs&, const ck_tile::stream_config&)``, the +``KERNEL_NAME`` string, and the ADataType/BDataType/CDataType/AccDataType/ +ScaleType/MxGemmHostArgs/ALayout/BLayout/CLayout aliases the ctypes lib expects. + +Each header is compiled per-kernel via force-include: + + hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE mx_gemm_ctypes_lib.cpp + +mx_gemm is microscaling GEMM (fp4/fp8 A*B, per-32-K e8m0 block scales), gfx950 +only. The single valid trait combo is comp_async + cshuffle + intrawave, with a +fixed 16x16x128 warp tile. +""" + +import argparse +import json +import logging +import sys +import tempfile +from pathlib import Path +from typing import Optional, Tuple + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# Robust import of the Old-TE builder (independent of CWD). +# dispatcher/codegen/ -> ../../tile_engine/ops/gemm/mx_gemm +# ============================================================================= + +_THIS_DIR = Path(__file__).resolve().parent +_GEMM_DIR = (_THIS_DIR / ".." / ".." / "tile_engine" / "ops" / "gemm").resolve() +_MX_GEMM_DIR = _GEMM_DIR / "mx_gemm" + +for _p in (str(_GEMM_DIR), str(_MX_GEMM_DIR)): + if _p not in sys.path: + sys.path.insert(0, _p) + + +def _load_mx_builder(): + """Import MxGemmKernelBuilder from the Old-TE ops tree.""" + import importlib.util + + mx_path = _MX_GEMM_DIR / "mx_gemm_instance_builder.py" + spec = importlib.util.spec_from_file_location( + "mx_gemm_instance_builder", str(mx_path) + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.MxGemmKernelBuilder + + +# ============================================================================= +# Config validation (restrict to the single valid mx_gemm combo) +# ============================================================================= + +KERNEL_NAME_PREFIX = "mx_gemm" + +VALID_DATATYPES = {"fp4", "fp8"} +VALID_LAYOUT = "rcr" +VALID_PIPELINE = "comp_async" +VALID_EPILOGUE = "cshuffle" +VALID_SCHEDULER = "intrawave" +FIXED_WARP_TILE = (16, 16, 128) + +_REQUIRED_TILE_KEYS = ( + "tile_m", + "tile_n", + "tile_k", + "warp_m", + "warp_n", + "warp_k", + "warp_tile_m", + "warp_tile_n", + "warp_tile_k", +) + + +def _validate(cfg: dict) -> None: + datatype = cfg.get("datatype") + if datatype not in VALID_DATATYPES: + raise ValueError( + f"datatype must be one of {sorted(VALID_DATATYPES)}, got {datatype!r}" + ) + + layout = cfg.get("layout") + if layout != VALID_LAYOUT: + raise ValueError(f"layout must be {VALID_LAYOUT!r}, got {layout!r}") + + pipeline = cfg.get("pipeline", VALID_PIPELINE) + if pipeline != VALID_PIPELINE: + raise ValueError( + f"pipeline must be {VALID_PIPELINE!r} for mx_gemm, got {pipeline!r}" + ) + + epilogue = cfg.get("epilogue", VALID_EPILOGUE) + if epilogue != VALID_EPILOGUE: + raise ValueError( + f"epilogue must be {VALID_EPILOGUE!r} for mx_gemm, got {epilogue!r}" + ) + + scheduler = cfg.get("scheduler", VALID_SCHEDULER) + if scheduler != VALID_SCHEDULER: + raise ValueError( + f"scheduler must be {VALID_SCHEDULER!r} for mx_gemm, got {scheduler!r}" + ) + + tc = cfg.get("tile_config") + if not isinstance(tc, dict): + raise ValueError("config must contain a 'tile_config' object") + missing = [k for k in _REQUIRED_TILE_KEYS if k not in tc] + if missing: + raise ValueError(f"tile_config missing keys: {missing}") + + warp_tile = (tc["warp_tile_m"], tc["warp_tile_n"], tc["warp_tile_k"]) + if tuple(warp_tile) != FIXED_WARP_TILE: + raise ValueError( + f"mx_gemm warp tile is fixed at {FIXED_WARP_TILE}, got {tuple(warp_tile)}" + ) + + +def _tile_config_from_cfg(cfg: dict) -> dict: + tc = cfg["tile_config"] + return {k: int(tc[k]) for k in _REQUIRED_TILE_KEYS} + + +def _trait_combo_from_cfg(cfg: dict) -> Tuple: + """7-tuple: (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, persistent).""" + return ( + cfg.get("pipeline", VALID_PIPELINE), + cfg.get("epilogue", VALID_EPILOGUE), + cfg.get("scheduler", VALID_SCHEDULER), + bool(cfg.get("pad_m", False)), + bool(cfg.get("pad_n", False)), + bool(cfg.get("pad_k", False)), + bool(cfg.get("persistent", False)), + ) + + +# ============================================================================= +# Builder driver +# ============================================================================= + +_AUTOGEN_BANNER = """\ +// Auto-generated MX GEMM kernel instance -- DO NOT EDIT. +// Regenerate via dispatcher/codegen/unified_mx_gemm_codegen.py +// This header is force-included into the ctypes translation unit: +// hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE mx_gemm_ctypes_lib.cpp +// It defines at global scope: SelectedKernel, KERNEL_NAME, ADataType/BDataType/ +// CDataType/AccDataType/ScaleType, MxGemmHostArgs, and ALayout/BLayout/CLayout. +""" + + +def _make_builder(cfg: dict): + """Instantiate MxGemmKernelBuilder against a temp working dir + temp config. + + The Old-TE __init__ loads config_json only when it is an existing FILE path, + and _generate_kernel_instance reads config.get("k_block_per_cu"). So we write + a minimal config file (carrying k_block_per_cu) and hand its path to the + builder. The builder's own side-effect .hpp lands in working_path (ignored). + """ + MxGemmKernelBuilder = _load_mx_builder() + + work_dir = Path(tempfile.mkdtemp(prefix="mx_gemm_codegen_")) + tmp_cfg = {"k_block_per_cu": int(cfg.get("k_block_per_cu", 1))} + cfg_path = work_dir / "mx_gemm_codegen_config.json" + cfg_path.write_text(json.dumps(tmp_cfg)) + + return MxGemmKernelBuilder( + KERNEL_NAME_PREFIX, + work_dir, + cfg.get("gpu_target", "gfx950"), + cfg["datatype"], + cfg["layout"], + config_json=str(cfg_path), + ) + + +def _fix_includes(code: str) -> str: + """Drop the Old-TE builder's stale ``ck_tile/ops/gemm_mx.hpp`` umbrella include. + + That umbrella does not exist on develop; the mx kernel + comp_async pipeline + are already pulled in by ``ck_tile/ops/gemm.hpp`` (also emitted in the same + header), and the host-side preshuffle helper by ``ck_tile/host.hpp`` in the + ctypes TU. Removing the stale line keeps the generated header self-contained. + """ + out_lines = [] + for line in code.splitlines(keepends=True): + if '#include "ck_tile/ops/gemm_mx.hpp"' in line: + continue + out_lines.append(line) + return "".join(out_lines) + + +def _generate(cfg: dict) -> Tuple[str, str]: + """Return (kernel_name, instance_code) for the given concrete config.""" + _validate(cfg) + builder = _make_builder(cfg) + tile_config = _tile_config_from_cfg(cfg) + trait_combo = _trait_combo_from_cfg(cfg) + name, code = builder._generate_kernel_instance(tile_config, trait_combo) + return name, _fix_includes(code) + + +def kernel_name(cfg: dict) -> str: + """Compute the kernel name without generating/writing code.""" + _validate(cfg) + builder = _make_builder(cfg) + tile_config = _tile_config_from_cfg(cfg) + trait_combo = _trait_combo_from_cfg(cfg) + return builder._format_kernel_name(trait_combo, tile_config) + + +def generate_kernel(output_dir: Path, cfg: dict) -> Optional[Path]: + output_dir.mkdir(parents=True, exist_ok=True) + name, code = _generate(cfg) + out = output_dir / f"{name}.hpp" + out.write_text(_AUTOGEN_BANNER + code) + log.info("wrote %s", out.name) + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description="MX GEMM kernel header generator") + ap.add_argument("--output-dir", type=Path) + ap.add_argument("--config", type=Path, help="path to a config JSON file") + ap.add_argument("--config-json", type=str, help="inline config JSON string") + ap.add_argument( + "--list-name", + action="store_true", + help="print the kernel name only; write no file", + ) + args = ap.parse_args() + + if args.config_json: + cfg = json.loads(args.config_json) + elif args.config: + cfg = json.loads(Path(args.config).read_text()) + else: + ap.error("one of --config-json or --config is required") + + if args.list_name: + print(kernel_name(cfg)) + return 0 + + if not args.output_dir: + ap.error("--output-dir is required unless --list-name is given") + + return 0 if generate_kernel(args.output_dir, cfg) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py new file mode 100644 index 000000000000..518c4d4af341 --- /dev/null +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -0,0 +1,602 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +MX-GEMM dispatcher utilities (TileEngine -> Dispatcher bridge). + +Three-layer Python bridge for the dispatcher's microscaling-GEMM path +(fp4/fp8 A.B with per-32-K e8m0 block scales, gfx950/MI350 only): + + MxGemmKernelConfig -- describes one kernel; .name is byte-exact with the + codegen KERNEL_NAME (obtained by shelling the codegen + CLI with --list-name so utils and codegen never drift) + MxGemmDispatcherLib -- thin ctypes wrapper around a compiled .so + GpuMxGemmRunner -- high-level runner: generates quantized inputs, calls the + kernel, and provides a numpy microscaled reference + +Build helper (self-contained): + setup_multiple_mx_gemm_dispatchers(configs, ...) : codegen -> hipcc -> .so paths + +Data types (verified against ck_tile headers / example/ck_tile/42_mx_gemm): + fp8 : ck_tile::fp8_t == float8_e4m3_t (OCP e4m3, bias 7, on gfx950 device) + fp4 : ck_tile::pk_fp4_t == pk_float4_e2m1_t (two e2m1 values packed per byte; + low nibble = even-K element, high nibble = odd-K element) + C : ck_tile::fp16_t + scale: ck_tile::e8m0_t (biased exponent, byte e decodes to 2^(e-127); 127 == 1.0) +""" + +import concurrent.futures +import ctypes +import json +import logging +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import numpy as np + +log = logging.getLogger(__name__) + +_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_mx_gemm_codegen.py" +_CTYPES_LIB_SRC = ( + Path(__file__).parent.parent / "bindings" / "ctypes" / "mx_gemm_ctypes_lib.cpp" +) +# ck root == the composablekernel dir (three levels up from dispatcher/python). +_CK_ROOT = Path(__file__).parent.parent.parent +_HIPCC = os.environ.get("CK_TILE_HIPCC", "/opt/rocm/bin/hipcc") +_DEFAULT_ARCH = "gfx950" + +# MX GEMM scales every 32 K-elements with one e8m0 byte. +SCALE_BLOCK = 32 + +# e8m0 byte for scale == 1.0 (2^(127-127)). +E8M0_ONE = 127 + +# ============================================================================= +# fp4 (e2m1) exact value grid, byte codes match pk_fp4 e2m1_to_fp32_table. +# index -> value ; index is the 4-bit e2m1 code. +# ============================================================================= +_FP4_E2M1_VALUES = np.array( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=np.float32, +) +# value -> 4-bit code (skip -0.0; +0.0 code 0 covers zero). +_FP4_VALUE_TO_CODE = { + 0.0: 0, 0.5: 1, 1.0: 2, 1.5: 3, 2.0: 4, 3.0: 5, 4.0: 6, 6.0: 7, + -0.5: 9, -1.0: 10, -1.5: 11, -2.0: 12, -3.0: 13, -4.0: 14, -6.0: 15, +} + +# ============================================================================= +# fp8 (e4m3, OCP) raw-byte codec for the exact grid we test on. +# These values are exactly representable in BOTH OCP and FNUZ e4m3, and the +# device (gfx950) uses OCP e4m3 (bias 7). We emit OCP bytes to match device. +# Bytes verified against pk_fp4 e2m1_to_fp8_table (OCP branch) + hand-derived +# grid entries (sign|exp4|mant3, bias 7). +# ============================================================================= +_FP8_OCP_VALUE_TO_BYTE = { + 0.0: 0x00, 0.5: 0x30, 1.0: 0x38, 1.5: 0x3C, 2.0: 0x40, 2.5: 0x42, + 3.0: 0x44, 4.0: 0x48, 6.0: 0x4C, + -0.5: 0xB0, -1.0: 0xB8, -1.5: 0xBC, -2.0: 0xC0, -2.5: 0xC2, + -3.0: 0xC4, -4.0: 0xC8, -6.0: 0xCC, +} +_FP8_OCP_BYTE_TO_VALUE = {b: v for v, b in _FP8_OCP_VALUE_TO_BYTE.items()} + + +def e8m0_to_float(byte) -> np.ndarray: + """Decode e8m0 bytes to float: 2^(e-127); e==255 is NaN.""" + b = np.asarray(byte, dtype=np.int32) + out = np.exp2((b - 127).astype(np.float32)) + out = np.where(b == 255, np.nan, out) + return out.astype(np.float32) + + +def float_to_e8m0(scale) -> np.ndarray: + """Encode a power-of-two float scale to an e8m0 byte (round exponent).""" + s = np.asarray(scale, dtype=np.float32) + e = np.rint(np.log2(np.where(s > 0, s, 1.0))).astype(np.int32) + 127 + e = np.clip(e, 0, 254).astype(np.uint8) + return e + + +# ============================================================================= +# Config +# ============================================================================= + + +@dataclass +class MxGemmKernelConfig: + datatype: str = "fp8" # fp8 | fp4 + layout: str = "rcr" # a/b/c ; only rcr supported by mx_gemm + gpu_target: str = _DEFAULT_ARCH + pipeline: str = "comp_async" + epilogue: str = "cshuffle" + scheduler: str = "intrawave" + pad_m: bool = False + pad_n: bool = False + pad_k: bool = False + persistent: bool = False + + tile_m: int = 128 + tile_n: int = 128 + tile_k: int = 128 + warp_m: int = 2 + warp_n: int = 2 + warp_k: int = 1 + warp_tile_m: int = 16 + warp_tile_n: int = 16 + warp_tile_k: int = 128 + + k_block_per_cu: int = 1 + + # cached name so we do not re-shell the codegen repeatedly. + _name_cache: Optional[str] = field(default=None, repr=False, compare=False) + + def to_codegen_config(self) -> dict: + return { + "datatype": self.datatype, + "layout": self.layout, + "gpu_target": self.gpu_target, + "pipeline": self.pipeline, + "epilogue": self.epilogue, + "scheduler": self.scheduler, + "pad_m": self.pad_m, + "pad_n": self.pad_n, + "pad_k": self.pad_k, + "persistent": self.persistent, + "tile_config": { + "tile_m": self.tile_m, "tile_n": self.tile_n, "tile_k": self.tile_k, + "warp_m": self.warp_m, "warp_n": self.warp_n, "warp_k": self.warp_k, + "warp_tile_m": self.warp_tile_m, "warp_tile_n": self.warp_tile_n, + "warp_tile_k": self.warp_tile_k, + }, + "k_block_per_cu": self.k_block_per_cu, + } + + @property + def name(self) -> str: + """Byte-exact kernel name from the codegen CLI (--list-name). + + The codegen owns the canonical name format; we never reconstruct it here + to avoid drift. Falls back to a locally computed name only if the codegen + script is not yet present (so utils stays import/unit testable before the + codegen component lands). + """ + if self._name_cache is not None: + return self._name_cache + if _CODEGEN_SCRIPT.exists(): + try: + r = subprocess.run( + [sys.executable, str(_CODEGEN_SCRIPT), + "--output-dir", tempfile.gettempdir(), + "--config-json", json.dumps(self.to_codegen_config()), + "--list-name"], + capture_output=True, text=True, timeout=120, + ) + if r.returncode == 0 and r.stdout.strip(): + self._name_cache = r.stdout.strip().splitlines()[-1].strip() + return self._name_cache + log.warning("codegen --list-name failed for %s:\n%s", + self._fallback_name(), r.stderr[-800:]) + except Exception as exc: # noqa: BLE001 + log.warning("codegen --list-name error: %s", exc) + self._name_cache = self._fallback_name() + return self._name_cache + + def _fallback_name(self) -> str: + """Local reconstruction mirroring the documented codegen name format. + + Format (per contract): + mx_gemm_{dtype}_{layout}_comp_async_cshuffle_intrawave_ + {PadM}_{PadN}_{PadK}[_{Persistent}]_ + {tileM}x{tileN}x{tileK}_{warpM}x{warpN}x{warpK}_{wtM}x{wtN}x{wtK} + Only used when the codegen script is absent (never at real build time). + """ + parts = [ + "mx_gemm", self.datatype, self.layout, + self.pipeline, self.epilogue, self.scheduler, + str(self.pad_m), str(self.pad_n), str(self.pad_k), + ] + if self.persistent: + parts.append(str(self.persistent)) + parts.append(f"{self.tile_m}x{self.tile_n}x{self.tile_k}") + parts.append(f"{self.warp_m}x{self.warp_n}x{self.warp_k}") + parts.append(f"{self.warp_tile_m}x{self.warp_tile_n}x{self.warp_tile_k}") + return "_".join(parts) + + def is_valid(self) -> bool: + if self.layout != "rcr": + return False + if self.datatype not in ("fp8", "fp4"): + return False + if not ( + self.tile_m % (self.warp_m * self.warp_tile_m) == 0 + and self.tile_n % (self.warp_n * self.warp_tile_n) == 0 + and self.tile_k % (self.warp_k * self.warp_tile_k) == 0 + ): + return False + # MX XDL uses 16x16x128 warp tiles on gfx950. + return (self.warp_tile_m, self.warp_tile_n, self.warp_tile_k) == (16, 16, 128) + + +# ============================================================================= +# Problem / result +# ============================================================================= + + +@dataclass +class MxGemmProblem: + M: int + N: int + K: int + k_batch: int = 1 + + def __post_init__(self): + if self.K % SCALE_BLOCK != 0: + raise ValueError(f"MX GEMM requires K % {SCALE_BLOCK} == 0, got K={self.K}") + + @property + def scale_k(self) -> int: + return self.K // SCALE_BLOCK + + @property + def flops(self) -> int: + return 2 * self.M * self.N * self.K + + +@dataclass +class MxGemmResult: + C: object + time_ms: float + kernel_name: str + + +# ============================================================================= +# ctypes wrapper +# ============================================================================= + + +class MxGemmDispatcherLib: + def __init__(self, so_path: Path): + self.so_path = Path(so_path) + if not self.so_path.exists(): + raise FileNotFoundError(f"mx_gemm .so not found: {self.so_path}") + self._lib = ctypes.CDLL(str(self.so_path)) + self._setup() + # Contract exposes both dispatcher_initialize() and dispatcher_init(). + init = getattr(self._lib, "dispatcher_initialize", None) + if init is None: + init = self._lib.dispatcher_init + if init() != 0: + raise RuntimeError("dispatcher_initialize failed") + + def _setup(self): + lib = self._lib + for fn in ("dispatcher_initialize", "dispatcher_init"): + f = getattr(lib, fn, None) + if f is not None: + f.restype = ctypes.c_int + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + if hasattr(lib, "dispatcher_get_kernel_count"): + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + lib.dispatcher_cleanup.restype = None + lib.dispatcher_run_mx_gemm.restype = ctypes.c_int + lib.dispatcher_run_mx_gemm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_void_p, # C + ctypes.POINTER(ctypes.c_uint8), # scale_a + ctypes.POINTER(ctypes.c_uint8), # scale_b + ctypes.c_int, # M + ctypes.c_int, # N + ctypes.c_int, # K + ctypes.c_int, # k_batch + ctypes.POINTER(ctypes.c_float), # time_ms + ] + + def kernel_name(self) -> str: + raw = self._lib.dispatcher_get_kernel_name() + return raw.decode() if raw else "" + + def kernel_count(self) -> int: + if hasattr(self._lib, "dispatcher_get_kernel_count"): + return int(self._lib.dispatcher_get_kernel_count()) + return 1 + + def run(self, A, B, C, scale_a, scale_b, prob: MxGemmProblem) -> float: + A = np.ascontiguousarray(A) + B = np.ascontiguousarray(B) + C = np.ascontiguousarray(C) + sa = np.ascontiguousarray(scale_a, dtype=np.uint8) + sb = np.ascontiguousarray(scale_b, dtype=np.uint8) + tms = ctypes.c_float(0.0) + rc = self._lib.dispatcher_run_mx_gemm( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + sa.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8)), + sb.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8)), + int(prob.M), int(prob.N), int(prob.K), int(prob.k_batch), + ctypes.byref(tms), + ) + if rc != 0: + raise RuntimeError(f"dispatcher_run_mx_gemm rc={rc} " + f"({'unsupported' if rc == -2 else 'error'})") + return tms.value + + def cleanup(self): + try: + self._lib.dispatcher_cleanup() + except Exception: # noqa: BLE001 + pass + + def __del__(self): + self.cleanup() + + +# ============================================================================= +# Quantization helpers (host-side; produce the exact bytes the kernel expects) +# ============================================================================= + + +def quantize_fp8(vals: np.ndarray) -> np.ndarray: + """Grid float values -> raw OCP e4m3 bytes (uint8), one per element. + + Values MUST come from the exact e4m3 grid (see _FP8_OCP_VALUE_TO_BYTE) so the + mapping is lossless; that keeps the numpy reference unambiguous. + """ + flat = vals.reshape(-1) + out = np.empty(flat.shape, dtype=np.uint8) + for i, v in enumerate(flat): + fv = float(v) + if fv == 0.0: + fv = 0.0 # collapse -0.0 + out[i] = _FP8_OCP_VALUE_TO_BYTE[round(fv, 3)] + return out.reshape(vals.shape) + + +def dequantize_fp8(bytes_arr: np.ndarray) -> np.ndarray: + flat = np.asarray(bytes_arr, dtype=np.uint8).reshape(-1) + out = np.empty(flat.shape, dtype=np.float32) + for i, b in enumerate(flat): + out[i] = _FP8_OCP_BYTE_TO_VALUE[int(b)] + return out.reshape(bytes_arr.shape) + + +def quantize_fp4_packed(vals: np.ndarray) -> np.ndarray: + """[M,K] grid floats -> packed pk_fp4 bytes [M,K//2]. + + Low nibble holds the even-K element, high nibble the odd-K element, matching + pk_fp4 _pack(x0,x1) = (x1<<4)|(x0&0xF) and the reference unpack<0>=lo. + + Returns the physically-packed pk_fp4 buffer ([M, K//2] bytes, one byte == two + fp4). mx_gemm_ctypes_lib sizes the A/B device buffer via + HostTensor::get_element_space_size_in_bytes(), which divides the logical + element count by numeric_traits::PackedSize==2, so this packed layout + is exactly what it expects (fp8 with PackedSize==1 is the 1 byte/element case). + """ + M, K = vals.shape + assert K % 2 == 0, "fp4 packs two K elements per byte" + codes = np.empty((M, K), dtype=np.uint8) + flat = vals.reshape(-1) + for i, v in enumerate(flat): + fv = float(v) + if fv == 0.0: + fv = 0.0 + codes.reshape(-1)[i] = _FP4_VALUE_TO_CODE[round(fv, 3)] + lo = codes[:, 0::2] + hi = codes[:, 1::2] + return ((hi << 4) | (lo & 0x0F)).astype(np.uint8) + + +def dequantize_fp4_packed(packed: np.ndarray, K: int) -> np.ndarray: + """Packed pk_fp4 bytes [M,K//2] -> [M,K] float grid values.""" + p = np.asarray(packed, dtype=np.uint8) + M = p.shape[0] + lo = (p & 0x0F).astype(np.int32) + hi = ((p >> 4) & 0x0F).astype(np.int32) + out = np.empty((M, K), dtype=np.float32) + out[:, 0::2] = _FP4_E2M1_VALUES[lo] + out[:, 1::2] = _FP4_E2M1_VALUES[hi] + return out + + +# ============================================================================= +# Numpy microscaled reference +# ============================================================================= + + +def mx_gemm_reference(A_deq, B_deq, scale_a_byte, scale_b_byte, prob: MxGemmProblem): + """C[m,n] = sum_kb sa[m,kb]*sb[n,kb] * sum_{j in block} A[m,kb*32+j]*B[kb*32+j,n]. + + A_deq: [M,K] dequantized floats (logical, unscaled A values). + B_deq: [K,N] dequantized floats (logical, unscaled B values). + scale_a_byte: e8m0 bytes [M, K/32] ; scale_b_byte: e8m0 bytes [N, K/32]. + Accumulated in fp32, returned as fp16 (kernel output dtype). + """ + M, N, K = prob.M, prob.N, prob.K + nkb = prob.scale_k + A = np.asarray(A_deq, dtype=np.float32).reshape(M, K) + B = np.asarray(B_deq, dtype=np.float32).reshape(K, N) + sa = e8m0_to_float(scale_a_byte).reshape(M, nkb) # [M, K/32] + sb = e8m0_to_float(scale_b_byte).reshape(N, nkb) # [N, K/32] + + # Scale A per (m, K-block) and B per (n, K-block). Expand block scales to K. + sa_k = np.repeat(sa, SCALE_BLOCK, axis=1) # [M, K] + sb_k = np.repeat(sb, SCALE_BLOCK, axis=1) # [N, K] + A_scaled = A * sa_k # [M, K] + B_scaled = (B.T * sb_k).T # [K, N] + C = A_scaled.astype(np.float32) @ B_scaled.astype(np.float32) # [M, N] + return C.astype(np.float16) + + +# ============================================================================= +# Runner +# ============================================================================= + + +class GpuMxGemmRunner: + def __init__(self, so_path: Path, dtype: str = "fp8"): + self._lib = MxGemmDispatcherLib(so_path) + self.dtype = dtype + + @property + def kernel_name(self) -> str: + return self._lib.kernel_name() + + def make_inputs(self, prob: MxGemmProblem, scale: float = 1.0, seed: int = 0): + """Generate grid-exact A/B and e8m0 scales. + + Returns (A_deq[M,K], B_deq[K,N], A_bytes, B_bytes, sa_bytes[M,K/32], + sb_bytes[N,K/32]). A_bytes/B_bytes are the raw device buffers. + scale is a single power-of-two applied uniformly (1.0 == e8m0 byte 127). + """ + M, N, K = prob.M, prob.N, prob.K + rng = np.random.default_rng(seed) + # Draw from a small grid that is exact in both fp8 e4m3 and fp4 e2m1. + grid = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], + dtype=np.float32) + A_deq = rng.choice(grid, size=(M, K)).astype(np.float32) + B_deq = rng.choice(grid, size=(K, N)).astype(np.float32) + + if self.dtype == "fp8": + A_bytes = quantize_fp8(A_deq) # [M,K] uint8 + # B stored col-major (rcr) => [N,K] row-major bytes. + B_bytes = quantize_fp8(B_deq.T) # [N,K] uint8 + elif self.dtype == "fp4": + A_bytes = quantize_fp4_packed(A_deq) # [M,K//2] uint8 + B_bytes = quantize_fp4_packed(B_deq.T) # [N,K//2] uint8 + else: + raise ValueError(f"unsupported dtype {self.dtype}") + + sbyte = int(float_to_e8m0(np.float32(scale))) + sa_bytes = np.full((M, prob.scale_k), sbyte, dtype=np.uint8) + sb_bytes = np.full((N, prob.scale_k), sbyte, dtype=np.uint8) + return A_deq, B_deq, A_bytes, B_bytes, sa_bytes, sb_bytes + + def run(self, prob: MxGemmProblem, A_bytes, B_bytes, sa_bytes, sb_bytes): + C = np.zeros((prob.M, prob.N), dtype=np.float16) + t = self._lib.run(A_bytes, B_bytes, C, sa_bytes, sb_bytes, prob) + return MxGemmResult(C=C, time_ms=t, kernel_name=self.kernel_name) + + @staticmethod + def reference(A_deq, B_deq, sa_bytes, sb_bytes, prob: MxGemmProblem): + return mx_gemm_reference(A_deq, B_deq, sa_bytes, sb_bytes, prob) + + +# ============================================================================= +# Build pipeline +# ============================================================================= + + +def _generate_kernel(cfg: MxGemmKernelConfig, headers_dir: Path) -> Optional[Path]: + cmd = [ + sys.executable, str(_CODEGEN_SCRIPT), "--output-dir", str(headers_dir), + "--config-json", json.dumps(cfg.to_codegen_config()), + ] + r = subprocess.run(cmd, capture_output=True, text=True, timeout=180) + if r.returncode != 0: + log.error("codegen failed for %s:\n%s", cfg.name, r.stderr[-1500:]) + return None + hpp = headers_dir / f"{cfg.name}.hpp" + return hpp if hpp.exists() else None + + +def _compile_kernel(hpp: Path, so: Path, arch: str) -> bool: + inc = [ + f"-I{_CK_ROOT}/include", f"-I{_CK_ROOT}", + f"-I{_CK_ROOT}/tile_engine/ops", + f"-I{_CK_ROOT}/tile_engine/ops/gemm", + f"-I{_CK_ROOT}/tile_engine/ops/gemm/mx_gemm", + ] + cmd = [ + _HIPCC, "-shared", "-fPIC", "-O3", "-std=c++17", + *inc, + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f"-include{hpp}", + "-D__HIP_PLATFORM_AMD__", f"--offload-arch={arch}", f'-DGFX_ARCH="{arch}"', + "-Wno-undefined-func-template", "-Wno-float-equal", + str(_CTYPES_LIB_SRC), "-o", str(so), + ] + r = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) + if r.returncode != 0: + log.error("compile failed for %s:\n%s", so.name, r.stderr[-3000:]) + return False + return True + + +def setup_multiple_mx_gemm_dispatchers( + configs: List[MxGemmKernelConfig], + output_dir: Optional[Path] = None, + gfx_arch: str = _DEFAULT_ARCH, + parallel: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """codegen -> hipcc -> .so for each config. Returns paths aligned with + `configs` (None on failure). Dedups by .name; per-arch .so cache.""" + if not configs: + return [] + arch = gfx_arch or _DEFAULT_ARCH + base = Path(output_dir) if output_dir else Path(tempfile.mkdtemp(prefix="mx_gemm_bridge_")) + headers = base / "generated_kernels" + libs = base / "libs" + headers.mkdir(parents=True, exist_ok=True) + libs.mkdir(parents=True, exist_ok=True) + + seen: Dict[str, int] = {} + deduped: List[Tuple[int, MxGemmKernelConfig]] = [] + for i, c in enumerate(configs): + if c.name not in seen: + seen[c.name] = i + deduped.append((i, c)) + results: List[Optional[Path]] = [None] * len(configs) + + def build_one(idx: int, cfg: MxGemmKernelConfig): + so = libs / f"lib_{cfg.name}_{arch}.so" + if so.exists(): + return idx, so + hpp = _generate_kernel(cfg, headers) + if hpp is None: + return idx, None + return idx, (so if _compile_kernel(hpp, so, arch) else None) + + if parallel and len(deduped) > 1: + workers = max_workers or min(len(deduped), os.cpu_count() or 4) + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex: + futs = {ex.submit(build_one, i, c): i for i, c in deduped} + for f in concurrent.futures.as_completed(futs): + idx, so = f.result() + results[idx] = so + else: + for i, c in deduped: + _, so = build_one(i, c) + results[i] = so + + # Propagate deduped results to duplicate slots. + for i, c in enumerate(configs): + if results[i] is None and seen.get(c.name, i) != i: + results[i] = results[seen[c.name]] + built = sum(1 for r in results if r) + log.info("built %d/%d mx_gemm kernels for %s", built, len(configs), arch) + return results + + +def default_fp8_config(gfx_arch: str = _DEFAULT_ARCH) -> MxGemmKernelConfig: + return MxGemmKernelConfig( + datatype="fp8", layout="rcr", gpu_target=gfx_arch, + pipeline="comp_async", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + ) + + +def default_fp4_config(gfx_arch: str = _DEFAULT_ARCH) -> MxGemmKernelConfig: + cfg = default_fp8_config(gfx_arch) + cfg.datatype = "fp4" + cfg._name_cache = None + return cfg From 3119eeee25c06f7327a1af298e3381e4fd8bb1dd Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Mon, 13 Jul 2026 05:03:55 -0400 Subject: [PATCH 02/14] test(ck-tile): add CPU-only unit tests for mx_gemm bridge Adds dispatcher/tests/test_mx_gemm_bridge.py covering the config name format, codegen-JSON projection, dtype/layout/warp-tile validity gate, e8m0 scale codec, fp8/fp4 quantization round-trips, and the numpy microscaled reference. Satisfies the repo unit-test policy for the mx_gemm bridge source files. No GPU required. --- .../dispatcher/tests/test_mx_gemm_bridge.py | 210 ++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py diff --git a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py new file mode 100644 index 000000000000..77d23cfbcc3e --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the microscaling (mx) GEMM TileEngine -> Dispatcher bridge. + +Locks the config name format, the codegen-JSON projection, the dtype/layout/warp-tile +validity gate, the e8m0 scale codec, the fp8/fp4 quantization round-trips, and the numpy +microscaled reference. No GPU, no hipcc, no Old-TE builder import required. +""" + +import sys +import unittest +from pathlib import Path + +import numpy as np + +_DISP = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_DISP / "python")) +sys.path.insert(0, str(_DISP / "codegen")) + +from mx_gemm_utils import ( # noqa: E402 + SCALE_BLOCK, + E8M0_ONE, + MxGemmKernelConfig, + MxGemmProblem, + default_fp8_config, + default_fp4_config, + e8m0_to_float, + float_to_e8m0, + quantize_fp8, + dequantize_fp8, + quantize_fp4_packed, + dequantize_fp4_packed, + mx_gemm_reference, +) + + +class TestConfigName(unittest.TestCase): + def test_fallback_name_prefix(self): + cfg = default_fp8_config() + self.assertTrue(cfg._fallback_name().startswith("mx_gemm_fp8_rcr_")) + + def test_fallback_name_encodes_tiles(self): + cfg = MxGemmKernelConfig( + datatype="fp4", tile_m=64, tile_n=128, tile_k=256, + warp_m=1, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + ) + name = cfg._fallback_name() + self.assertIn("_fp4_rcr_", name) + self.assertIn("64x128x256", name) + self.assertIn("1x2x1", name) + self.assertIn("16x16x128", name) + self.assertNotIn(" ", name) + + def test_persistent_suffix_only_when_set(self): + self.assertNotIn("True", default_fp8_config()._fallback_name()) + cfg = default_fp8_config() + cfg.persistent = True + self.assertIn("True", cfg._fallback_name()) + + +class TestCodegenJson(unittest.TestCase): + def test_projection_roundtrip(self): + cfg = MxGemmKernelConfig( + datatype="fp8", tile_m=128, tile_n=128, tile_k=128, + warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + k_block_per_cu=3, + ) + j = cfg.to_codegen_config() + self.assertEqual(j["datatype"], "fp8") + self.assertEqual(j["layout"], "rcr") + self.assertEqual(j["tile_config"]["tile_k"], 128) + self.assertEqual(j["tile_config"]["warp_tile_k"], 128) + self.assertEqual(j["k_block_per_cu"], 3) + + +class TestValidity(unittest.TestCase): + def test_default_configs_valid(self): + self.assertTrue(default_fp8_config().is_valid()) + self.assertTrue(default_fp4_config().is_valid()) + + def test_non_rcr_rejected(self): + cfg = default_fp8_config() + cfg.layout = "rrr" + self.assertFalse(cfg.is_valid()) + + def test_bad_dtype_rejected(self): + cfg = default_fp8_config() + cfg.datatype = "bf16" + self.assertFalse(cfg.is_valid()) + + def test_wrong_warp_tile_rejected(self): + cfg = default_fp8_config() + cfg.warp_tile_k = 64 + self.assertFalse(cfg.is_valid()) + + def test_indivisible_tile_rejected(self): + cfg = default_fp8_config() + cfg.tile_m = 100 # not a multiple of warp_m * warp_tile_m (2*16=32) + self.assertFalse(cfg.is_valid()) + + +class TestProblem(unittest.TestCase): + def test_scale_k_and_flops(self): + p = MxGemmProblem(M=64, N=128, K=256) + self.assertEqual(p.scale_k, 256 // SCALE_BLOCK) + self.assertEqual(p.flops, 2 * 64 * 128 * 256) + + def test_k_not_multiple_of_32_raises(self): + with self.assertRaises(ValueError): + MxGemmProblem(M=32, N=32, K=48) + + +class TestE8m0Codec(unittest.TestCase): + def test_one_is_byte_127(self): + self.assertEqual(int(float_to_e8m0(np.float32(1.0))), E8M0_ONE) + self.assertEqual(float(e8m0_to_float(E8M0_ONE)), 1.0) + + def test_power_of_two_roundtrip(self): + for s in (0.25, 0.5, 1.0, 2.0, 4.0, 8.0): + b = float_to_e8m0(np.float32(s)) + self.assertAlmostEqual(float(e8m0_to_float(b)), s, places=6) + + def test_255_is_nan(self): + with np.errstate(over="ignore"): + self.assertTrue(np.isnan(float(e8m0_to_float(255)))) + + +class TestFp8Codec(unittest.TestCase): + def test_known_bytes(self): + self.assertEqual(int(quantize_fp8(np.array([1.0], np.float32))[0]), 0x38) + self.assertEqual(int(quantize_fp8(np.array([-2.0], np.float32))[0]), 0xC0) + + def test_roundtrip_grid(self): + grid = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], np.float32) + vals = np.tile(grid, (4, 1)) + self.assertTrue(np.array_equal(dequantize_fp8(quantize_fp8(vals)), vals)) + + +class TestFp4Codec(unittest.TestCase): + def test_pack_two_per_byte(self): + # one row, K=2 -> a single packed byte; low nibble even-K, high nibble odd-K. + vals = np.array([[1.0, 2.0]], np.float32) # codes: 1.0->2, 2.0->4 + packed = quantize_fp4_packed(vals) + self.assertEqual(packed.shape, (1, 1)) + self.assertEqual(int(packed[0, 0]), (4 << 4) | 2) + + def test_roundtrip_grid(self): + grid = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], np.float32) + rng = np.random.default_rng(0) + vals = rng.choice(grid, size=(4, 8)).astype(np.float32) + packed = quantize_fp4_packed(vals) + self.assertEqual(packed.shape, (4, 4)) + self.assertTrue(np.array_equal(dequantize_fp4_packed(packed, 8), vals)) + + +class TestReference(unittest.TestCase): + def _inputs(self, M, N, K, seed=0): + rng = np.random.default_rng(seed) + grid = np.array([-2.0, -1.0, -0.5, 0.5, 1.0, 2.0], np.float32) + A = rng.choice(grid, size=(M, K)).astype(np.float32) + B = rng.choice(grid, size=(K, N)).astype(np.float32) + return A, B + + def test_unit_scales_equals_plain_matmul(self): + M, N, K = 4, 3, 32 + A, B = self._inputs(M, N, K) + prob = MxGemmProblem(M=M, N=N, K=K) + one = int(float_to_e8m0(np.float32(1.0))) + sa = np.full((M, prob.scale_k), one, np.uint8) + sb = np.full((N, prob.scale_k), one, np.uint8) + ref = mx_gemm_reference(A, B, sa, sb, prob).astype(np.float32) + plain = (A @ B).astype(np.float16).astype(np.float32) + self.assertTrue(np.allclose(ref, plain, atol=1e-2)) + + def test_power_of_two_scale_multiplies(self): + M, N, K = 4, 3, 32 + A, B = self._inputs(M, N, K, seed=1) + prob = MxGemmProblem(M=M, N=N, K=K) + two = int(float_to_e8m0(np.float32(2.0))) + sa = np.full((M, prob.scale_k), two, np.uint8) + sb = np.full((N, prob.scale_k), two, np.uint8) + ref = mx_gemm_reference(A, B, sa, sb, prob).astype(np.float32) + # scale_a=2 and scale_b=2 -> product scaled by 4. + expected = (4.0 * (A @ B)).astype(np.float16).astype(np.float32) + self.assertTrue(np.allclose(ref, expected, atol=1e-1)) + + +class TestCodegenNameContract(unittest.TestCase): + """Optional: byte-exact name parity when the Old-TE builder is importable.""" + + def test_codegen_name_matches_config(self): + try: + from unified_mx_gemm_codegen import kernel_name + except Exception as exc: # noqa: BLE001 + self.skipTest(f"codegen import unavailable: {exc}") + cfg = default_fp8_config() + try: + name = kernel_name(cfg.to_codegen_config()) + except Exception as exc: # noqa: BLE001 + self.skipTest(f"Old-TE builder unavailable: {exc}") + self.assertTrue(name.startswith("mx_gemm_fp8_rcr_")) + + +if __name__ == "__main__": + unittest.main() From 4a5b4639c2c7bed7185caa2b0ccc902a0e0f9c60 Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Tue, 14 Jul 2026 23:20:15 -0400 Subject: [PATCH 03/14] harden(ck-tile): mx_gemm bridge divisibility guard, fp8 OCP precondition, bf8-rejected test Round-2 review hardening for the mx_gemm TE->Dispatcher bridge (no functional gap remained; these are LOW-severity robustness fixes): - mx_gemm_ctypes_lib.cpp: guard M/N/scale_k divisibility by the xdl pack factors before sizing the shuffled scale buffers; return -1 (matching the existing error convention) instead of silently truncating the buffers and feeding the kernel a corrupt scale layout. - mx_gemm_utils.py: make the fp8 OCP e4m3 assumption explicit. Add fp8_ocp_is_default_for_arch()/assert_fp8_ocp_supported() mirroring config.hpp (OCP is the gfx950/gfx12 device default; the lib compiles without an explicit -DCK_TILE_USE_OCP_FP8), document the codec precondition, and have GpuMxGemmRunner assert it for fp8 so an FNUZ arch cannot silently diverge from the numpy reference. - test_mx_gemm_bridge.py: add test_bf8_rejected asserting is_valid() and the codegen _validate() both refuse dtype="bf8" (Old-TE does not compile bf8). --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 17 ++++++ .../dispatcher/python/mx_gemm_utils.py | 61 +++++++++++++++++-- .../dispatcher/tests/test_mx_gemm_bridge.py | 16 +++++ 3 files changed, 89 insertions(+), 5 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 7e795e104535..02903d538e2b 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -177,6 +177,23 @@ int dispatcher_run_mx_gemm(const void* A, constexpr ck_tile::index_t xdl_mn_thread = SelectedKernel::WarpTileM; constexpr ck_tile::index_t xdl_k_thread = 64 / xdl_mn_thread; + // ---- Divisibility guard. The shuffled scale-buffer sizes below use integer + // division by the xdl pack factors (m/m_xdl_pack, n/n_xdl_pack, + // scale_k_size/k_xdl_pack). If M/N/scale_k are not exact multiples of their + // pack factors, that division silently truncates the shuffled buffers, so the + // pre-shuffle would read/write past valid data and the kernel would consume a + // corrupt scale layout. Fail loudly instead (matches the -1 error convention + // used above). ---- + if(m % m_xdl_pack != 0 || n % n_xdl_pack != 0 || scale_k_size % k_xdl_pack != 0) + { + std::cerr << "dispatcher_run_mx_gemm: M, N, and scale_k (=K/32) must be divisible by the " + "xdl pack factors (m_xdl_pack=" + << m_xdl_pack << ", n_xdl_pack=" << n_xdl_pack << ", k_xdl_pack=" << k_xdl_pack + << ") for the selected warp tile; got M=" << m << ", N=" << n + << ", scale_k=" << scale_k_size << "\n"; + return -1; + } + // ---- Build unshuffled scale HostTensors from the incoming raw e8m0 bytes. // scale_a: [M, K/32] row-major; scale_b: [N, K/32] row-major. ---- ck_tile::HostTensor scale_a_host( diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index 518c4d4af341..72f91cc5b1f2 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -75,10 +75,20 @@ # ============================================================================= # fp8 (e4m3, OCP) raw-byte codec for the exact grid we test on. -# These values are exactly representable in BOTH OCP and FNUZ e4m3, and the -# device (gfx950) uses OCP e4m3 (bias 7). We emit OCP bytes to match device. -# Bytes verified against pk_fp4 e2m1_to_fp8_table (OCP branch) + hand-derived -# grid entries (sign|exp4|mant3, bias 7). +# +# PRECONDITION (hard): the byte codes below are OCP e4m3 (bias 7). They are ONLY +# a valid reference for a kernel compiled with CK_TILE_USE_OCP_FP8 == 1. This is +# the gfx950 default (see include/ck_tile/core/config.hpp: OCP is selected for +# __gfx950__ / __gfx12__ and the bridge compiles the lib with +# --offload-arch=gfx950 without overriding the flag). On any arch where CK +# defaults to (or is built for) FNUZ e4m3 (bias 8), the DEVICE bytes for the same +# float value differ, so this numpy reference would SILENTLY disagree with the +# kernel. Callers that target a non-OCP arch must not use this codec. +# +# The grid values are exactly representable in BOTH OCP and FNUZ, but their +# RAW BYTES are not the same between the two formats, which is why the codec is +# OCP-specific. Bytes verified against pk_fp4 e2m1_to_fp8_table (OCP branch) + +# hand-derived grid entries (sign|exp4|mant3, bias 7). # ============================================================================= _FP8_OCP_VALUE_TO_BYTE = { 0.0: 0x00, 0.5: 0x30, 1.0: 0x38, 1.5: 0x3C, 2.0: 0x40, 2.5: 0x42, @@ -89,6 +99,36 @@ _FP8_OCP_BYTE_TO_VALUE = {b: v for v, b in _FP8_OCP_VALUE_TO_BYTE.items()} +def fp8_ocp_is_default_for_arch(arch: str) -> bool: + """True iff CK_TILE_USE_OCP_FP8 defaults to 1 for `arch`. + + Mirrors include/ck_tile/core/config.hpp: OCP e4m3 is the device default only + for gfx950 and gfx12; every other arch defaults to FNUZ e4m3. The bridge + compiles the mx_gemm lib without an explicit -DCK_TILE_USE_OCP_FP8, so the + effective fp8 format is exactly this arch default. Use this to guard the fp8 + numpy reference (see assert_fp8_ocp_supported / quantize_fp8). + """ + a = (arch or "").lower() + return a.startswith("gfx950") or a.startswith("gfx12") + + +def assert_fp8_ocp_supported(arch: str = _DEFAULT_ARCH) -> None: + """Fail loudly if the fp8 OCP codec would silently disagree with the device. + + The fp8 quantize/dequantize helpers emit/decode OCP e4m3 bytes; that is only + correct when the kernel is compiled with CK_TILE_USE_OCP_FP8 == 1 (see the + codec precondition above). Raise instead of producing bytes that mismatch a + FNUZ-built kernel. + """ + if not fp8_ocp_is_default_for_arch(arch): + raise ValueError( + f"fp8 mx_gemm reference is OCP e4m3 only, but arch '{arch}' defaults to " + "FNUZ e4m3 (CK_TILE_USE_OCP_FP8 == 0). The bridge lib is compiled without " + "an explicit -DCK_TILE_USE_OCP_FP8, so its device fp8 bytes would not match " + "this numpy reference. OCP fp8 is supported for gfx950/gfx12 only." + ) + + def e8m0_to_float(byte) -> np.ndarray: """Decode e8m0 bytes to float: 2^(e-127); e==255 is NaN.""" b = np.asarray(byte, dtype=np.int32) @@ -350,6 +390,10 @@ def quantize_fp8(vals: np.ndarray) -> np.ndarray: Values MUST come from the exact e4m3 grid (see _FP8_OCP_VALUE_TO_BYTE) so the mapping is lossless; that keeps the numpy reference unambiguous. + + PRECONDITION: OCP e4m3 only. These bytes match the device only for a kernel + compiled with CK_TILE_USE_OCP_FP8 == 1 (gfx950/gfx12 default). See the codec + comment above and assert_fp8_ocp_supported(); GpuMxGemmRunner enforces this. """ flat = vals.reshape(-1) out = np.empty(flat.shape, dtype=np.uint8) @@ -442,9 +486,16 @@ def mx_gemm_reference(A_deq, B_deq, scale_a_byte, scale_b_byte, prob: MxGemmProb class GpuMxGemmRunner: - def __init__(self, so_path: Path, dtype: str = "fp8"): + def __init__(self, so_path: Path, dtype: str = "fp8", arch: str = _DEFAULT_ARCH): + # The fp8 numpy reference (quantize_fp8/dequantize_fp8) emits OCP e4m3 + # bytes; that only matches the device when the lib was compiled with + # CK_TILE_USE_OCP_FP8 == 1 (the gfx950 default). Guard loudly so an + # FNUZ arch can't silently diverge from the reference. + if dtype == "fp8": + assert_fp8_ocp_supported(arch) self._lib = MxGemmDispatcherLib(so_path) self.dtype = dtype + self.arch = arch @property def kernel_name(self) -> str: diff --git a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py index 77d23cfbcc3e..f6e0b4eec6b3 100644 --- a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py @@ -93,6 +93,22 @@ def test_bad_dtype_rejected(self): cfg.datatype = "bf16" self.assertFalse(cfg.is_valid()) + def test_bf8_rejected(self): + # Old-TE argparse (choices=["fp4","fp8"]) + validate_gemm_mx never + # compile bf8/e5m2 for mx_gemm, so the bridge must reject it too. Assert + # both the config-level gate (is_valid) and, when importable, the + # codegen-level gate (_validate) refuse dtype="bf8". + cfg = default_fp8_config() + cfg.datatype = "bf8" + self.assertFalse(cfg.is_valid()) + + try: + from unified_mx_gemm_codegen import _validate # noqa: E402 + except Exception as exc: # noqa: BLE001 + self.skipTest(f"codegen import unavailable: {exc}") + with self.assertRaises(Exception): + _validate(cfg.to_codegen_config()) + def test_wrong_warp_tile_rejected(self): cfg = default_fp8_config() cfg.warp_tile_k = 64 From ec36f78a365273ef550ca764c21bd233d1aa60e9 Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Thu, 16 Jul 2026 14:53:40 -0400 Subject: [PATCH 04/14] fix(ck-tile): address Copilot round-3 findings on mx_gemm bridge - mx_gemm_ctypes_lib.cpp: make the TU compile in the CMake "no-kernel" fallback branch. Guard the kernel-dependent code (KERNEL_NAME, dispatcher_run_mx_gemm body) behind CK_TILE_SINGLE_KERNEL_INCLUDE; without a force-included kernel the lib now returns "" for the name and -2 (unsupported) from run, instead of failing to compile on SelectedKernel/MxGemmHostArgs/ScaleType. - unified_mx_gemm_codegen.py: stop leaking mx_gemm_codegen_* temp dirs. Turn _make_builder into a context manager backed by TemporaryDirectory so the work dir is cleaned up after each name/code generation (kernel_name/generate are shelled frequently via --list-name). - mx_gemm_utils.py: * float_to_e8m0: validate strictly-positive, finite input and raise instead of silently mapping non-positive scales to 1.0 (wrong scale byte). * quantize_fp8/quantize_fp4_packed: replace per-element Python loops with a vectorized sorted-grid LUT + searchsorted (off-grid values now raise). Removes the O(M*K) Python overhead that made the bridge look slow in sweeps. * MxGemmDispatcherLib.run: C is an in-place output buffer; validate it is a contiguous, writeable ndarray and raise rather than silently ascontiguousarray- copying a non-contiguous view (which would drop the caller's results). - test_mx_gemm_bridge.py: cover the new behaviors (e8m0 non-positive/non-finite raise, fp8/fp4 off-grid raise, -0.0 collapse, shape preservation). Divisibility guard (scale-buffer OOB) was already handled in the prior hardening commit; it returns -1 to match this file's existing error convention. --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 29 +++++++- .../codegen/unified_mx_gemm_codegen.py | 50 ++++++++------ .../dispatcher/python/mx_gemm_utils.py | 69 ++++++++++++++----- .../dispatcher/tests/test_mx_gemm_bridge.py | 33 +++++++++ 4 files changed, 141 insertions(+), 40 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 02903d538e2b..25eb4c07ebda 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -81,7 +81,15 @@ int dispatcher_initialize() } int dispatcher_init() { return dispatcher_initialize(); } -const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } +const char* dispatcher_get_kernel_name() +{ +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE + return KERNEL_NAME; +#else + // No kernel force-included (CMake no-kernel fallback): report an empty name. + return ""; +#endif +} int dispatcher_get_kernel_count() { return 1; } void dispatcher_cleanup() { g_initialized = false; } @@ -107,6 +115,24 @@ int dispatcher_run_mx_gemm(const void* A, int k_batch, float* time_ms) { +#ifndef CK_TILE_SINGLE_KERNEL_INCLUDE + // No kernel was force-included (see the CMake no-kernel fallback). The body + // below needs SelectedKernel/MxGemmHostArgs/ScaleType/... which only exist + // under -DCK_TILE_SINGLE_KERNEL_INCLUDE, so this build cannot run anything: + // report "unsupported" (-2) instead of failing to compile/link. + (void)A; + (void)B; + (void)C; + (void)scale_a; + (void)scale_b; + (void)M; + (void)N; + (void)K; + (void)k_batch; + (void)time_ms; + std::cerr << "dispatcher_run_mx_gemm: library built without a kernel; unsupported\n"; + return -2; +#else if(!g_initialized) { std::cerr << "dispatcher_run_mx_gemm: not initialized\n"; @@ -316,6 +342,7 @@ int dispatcher_run_mx_gemm(const void* A, *time_ms = exec_time; return 0; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE } } // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py index fbcc569d3b08..5ccc34a0311b 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py @@ -27,12 +27,13 @@ """ import argparse +import contextlib import json import logging import sys import tempfile from pathlib import Path -from typing import Optional, Tuple +from typing import Iterator, Optional, Tuple logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") log = logging.getLogger(__name__) @@ -166,29 +167,36 @@ def _trait_combo_from_cfg(cfg: dict) -> Tuple: """ -def _make_builder(cfg: dict): - """Instantiate MxGemmKernelBuilder against a temp working dir + temp config. +@contextlib.contextmanager +def _make_builder(cfg: dict) -> Iterator["object"]: + """Yield an MxGemmKernelBuilder bound to a self-cleaning temp working dir. The Old-TE __init__ loads config_json only when it is an existing FILE path, and _generate_kernel_instance reads config.get("k_block_per_cu"). So we write a minimal config file (carrying k_block_per_cu) and hand its path to the - builder. The builder's own side-effect .hpp lands in working_path (ignored). + builder. The builder's own side-effect .hpp lands in the working dir (ignored). + + This is a context manager (not a plain factory) because it owns a temp dir: + kernel_name()/generate_kernel() are shelled frequently (e.g. --list-name), so + a leaked mkdtemp() per call would pile up mx_gemm_codegen_* dirs under /tmp. + The builder is only used inside the `with` block, so cleanup on exit is safe. """ MxGemmKernelBuilder = _load_mx_builder() - work_dir = Path(tempfile.mkdtemp(prefix="mx_gemm_codegen_")) - tmp_cfg = {"k_block_per_cu": int(cfg.get("k_block_per_cu", 1))} - cfg_path = work_dir / "mx_gemm_codegen_config.json" - cfg_path.write_text(json.dumps(tmp_cfg)) - - return MxGemmKernelBuilder( - KERNEL_NAME_PREFIX, - work_dir, - cfg.get("gpu_target", "gfx950"), - cfg["datatype"], - cfg["layout"], - config_json=str(cfg_path), - ) + with tempfile.TemporaryDirectory(prefix="mx_gemm_codegen_") as work_dir_str: + work_dir = Path(work_dir_str) + tmp_cfg = {"k_block_per_cu": int(cfg.get("k_block_per_cu", 1))} + cfg_path = work_dir / "mx_gemm_codegen_config.json" + cfg_path.write_text(json.dumps(tmp_cfg)) + + yield MxGemmKernelBuilder( + KERNEL_NAME_PREFIX, + work_dir, + cfg.get("gpu_target", "gfx950"), + cfg["datatype"], + cfg["layout"], + config_json=str(cfg_path), + ) def _fix_includes(code: str) -> str: @@ -210,20 +218,20 @@ def _fix_includes(code: str) -> str: def _generate(cfg: dict) -> Tuple[str, str]: """Return (kernel_name, instance_code) for the given concrete config.""" _validate(cfg) - builder = _make_builder(cfg) tile_config = _tile_config_from_cfg(cfg) trait_combo = _trait_combo_from_cfg(cfg) - name, code = builder._generate_kernel_instance(tile_config, trait_combo) + with _make_builder(cfg) as builder: + name, code = builder._generate_kernel_instance(tile_config, trait_combo) return name, _fix_includes(code) def kernel_name(cfg: dict) -> str: """Compute the kernel name without generating/writing code.""" _validate(cfg) - builder = _make_builder(cfg) tile_config = _tile_config_from_cfg(cfg) trait_combo = _trait_combo_from_cfg(cfg) - return builder._format_kernel_name(trait_combo, tile_config) + with _make_builder(cfg) as builder: + return builder._format_kernel_name(trait_combo, tile_config) def generate_kernel(output_dir: Path, cfg: dict) -> Optional[Path]: diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index 72f91cc5b1f2..0434394f740e 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -138,9 +138,20 @@ def e8m0_to_float(byte) -> np.ndarray: def float_to_e8m0(scale) -> np.ndarray: - """Encode a power-of-two float scale to an e8m0 byte (round exponent).""" + """Encode a strictly-positive power-of-two float scale to an e8m0 byte. + + The e8m0 format only encodes 2^(e-127) for e in [0, 254] (255 == NaN), so + the contract is a positive, finite input. A non-positive or non-finite value + is a caller bug (e.g. a zeroed/garbage scale); fail loudly rather than + silently substituting 1.0, which would emit a wrong scale byte. + """ s = np.asarray(scale, dtype=np.float32) - e = np.rint(np.log2(np.where(s > 0, s, 1.0))).astype(np.int32) + 127 + if not np.all(np.isfinite(s)) or np.any(s <= 0.0): + raise ValueError( + "float_to_e8m0 requires strictly positive, finite scales " + "(power-of-two); got a non-positive or non-finite value" + ) + e = np.rint(np.log2(s)).astype(np.int32) + 127 e = np.clip(e, 0, 254).astype(np.uint8) return e @@ -352,7 +363,16 @@ def kernel_count(self) -> int: def run(self, A, B, C, scale_a, scale_b, prob: MxGemmProblem) -> float: A = np.ascontiguousarray(A) B = np.ascontiguousarray(B) - C = np.ascontiguousarray(C) + # C is the OUTPUT buffer: the kernel writes results into it in place via a + # raw pointer. ascontiguousarray() would silently redirect the write into a + # throwaway copy and leave the caller's array unchanged -- a surprising bug + # for an output API. Require a contiguous, writeable ndarray instead. + if not isinstance(C, np.ndarray): + raise TypeError(f"C output must be a numpy.ndarray, got {type(C).__name__}") + if not C.flags["C_CONTIGUOUS"]: + raise ValueError("C output must be C-contiguous (results are written in place)") + if not C.flags["WRITEABLE"]: + raise ValueError("C output must be writeable (results are written in place)") sa = np.ascontiguousarray(scale_a, dtype=np.uint8) sb = np.ascontiguousarray(scale_b, dtype=np.uint8) tms = ctypes.c_float(0.0) @@ -385,6 +405,32 @@ def __del__(self): # ============================================================================= +def _map_grid_to_bytes(vals: np.ndarray, value_to_byte: dict, grid_name: str) -> np.ndarray: + """Vectorized exact-grid float -> byte code lookup (uint8, flattened). + + The inputs are drawn from a small fixed grid (see make_inputs), so instead of + a per-element Python loop -- which dominates runtime for realistic M*K -- we + round to the grid resolution and index a sorted LUT with a single + searchsorted. Off-grid values are a caller contract violation, so raise + (mirrors the old dict-KeyError) rather than silently snapping to a neighbour. + """ + keys = np.array(sorted(value_to_byte), dtype=np.float32) + byts = np.array([value_to_byte[float(k)] for k in keys], dtype=np.uint8) + + r = np.round(np.asarray(vals, dtype=np.float32), 3) + r = r + np.float32(0.0) # collapse -0.0 -> +0.0 so it matches the 0.0 key + flat = r.reshape(-1) + + idx = np.clip(np.searchsorted(keys, flat), 0, keys.size - 1) + if not np.all(keys[idx] == flat): + bad = flat[keys[idx] != flat] + raise KeyError( + f"{grid_name}: value(s) not on the exact quantization grid: " + f"{np.unique(bad)[:8].tolist()}" + ) + return byts[idx] + + def quantize_fp8(vals: np.ndarray) -> np.ndarray: """Grid float values -> raw OCP e4m3 bytes (uint8), one per element. @@ -395,14 +441,7 @@ def quantize_fp8(vals: np.ndarray) -> np.ndarray: compiled with CK_TILE_USE_OCP_FP8 == 1 (gfx950/gfx12 default). See the codec comment above and assert_fp8_ocp_supported(); GpuMxGemmRunner enforces this. """ - flat = vals.reshape(-1) - out = np.empty(flat.shape, dtype=np.uint8) - for i, v in enumerate(flat): - fv = float(v) - if fv == 0.0: - fv = 0.0 # collapse -0.0 - out[i] = _FP8_OCP_VALUE_TO_BYTE[round(fv, 3)] - return out.reshape(vals.shape) + return _map_grid_to_bytes(vals, _FP8_OCP_VALUE_TO_BYTE, "fp8 e4m3").reshape(vals.shape) def dequantize_fp8(bytes_arr: np.ndarray) -> np.ndarray: @@ -427,13 +466,7 @@ def quantize_fp4_packed(vals: np.ndarray) -> np.ndarray: """ M, K = vals.shape assert K % 2 == 0, "fp4 packs two K elements per byte" - codes = np.empty((M, K), dtype=np.uint8) - flat = vals.reshape(-1) - for i, v in enumerate(flat): - fv = float(v) - if fv == 0.0: - fv = 0.0 - codes.reshape(-1)[i] = _FP4_VALUE_TO_CODE[round(fv, 3)] + codes = _map_grid_to_bytes(vals, _FP4_VALUE_TO_CODE, "fp4 e2m1").reshape(M, K) lo = codes[:, 0::2] hi = codes[:, 1::2] return ((hi << 4) | (lo & 0x0F)).astype(np.uint8) diff --git a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py index f6e0b4eec6b3..69680b9ab648 100644 --- a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py @@ -145,6 +145,21 @@ def test_255_is_nan(self): with np.errstate(over="ignore"): self.assertTrue(np.isnan(float(e8m0_to_float(255)))) + def test_non_positive_raises(self): + # Contract is a strictly-positive power-of-two; a 0.0/negative scale is a + # caller bug and must fail loudly rather than silently encode 1.0. + for bad in (0.0, -1.0, -0.0): + with self.assertRaises(ValueError): + float_to_e8m0(np.float32(bad)) + with self.assertRaises(ValueError): + float_to_e8m0(np.array([1.0, 0.0, 2.0], np.float32)) + + def test_non_finite_raises(self): + with np.errstate(invalid="ignore"): + for bad in (np.inf, np.nan): + with self.assertRaises(ValueError): + float_to_e8m0(np.float32(bad)) + class TestFp8Codec(unittest.TestCase): def test_known_bytes(self): @@ -156,6 +171,19 @@ def test_roundtrip_grid(self): vals = np.tile(grid, (4, 1)) self.assertTrue(np.array_equal(dequantize_fp8(quantize_fp8(vals)), vals)) + def test_off_grid_raises(self): + # The vectorized codec must reject values not on the exact e4m3 grid + # instead of snapping them to a neighbour byte. + with self.assertRaises(KeyError): + quantize_fp8(np.array([[0.3]], np.float32)) + + def test_neg_zero_collapses_to_zero_byte(self): + self.assertEqual(int(quantize_fp8(np.array([-0.0], np.float32))[0]), 0x00) + + def test_shape_preserved_2d(self): + vals = np.full((3, 5), 1.0, np.float32) + self.assertEqual(quantize_fp8(vals).shape, (3, 5)) + class TestFp4Codec(unittest.TestCase): def test_pack_two_per_byte(self): @@ -173,6 +201,11 @@ def test_roundtrip_grid(self): self.assertEqual(packed.shape, (4, 4)) self.assertTrue(np.array_equal(dequantize_fp4_packed(packed, 8), vals)) + def test_off_grid_raises(self): + # 2.5 exists in fp8 e4m3 but NOT in the fp4 e2m1 grid -> must reject. + with self.assertRaises(KeyError): + quantize_fp4_packed(np.array([[2.5, 1.0]], np.float32)) + class TestReference(unittest.TestCase): def _inputs(self, M, N, K, seed=0): From 0488a9d33baad6375b058ca0d0ad8fe8885457f4 Mon Sep 17 00:00:00 2001 From: Mert Ozturk Date: Thu, 16 Jul 2026 23:44:16 -0400 Subject: [PATCH 05/14] polish(ck-tile): mx_gemm bridge LOW nits from round-1 review Two optional polish items from the round-1 reviewer (both LOW severity): 1. Divisibility guard returns -2 instead of -1 The xdl-pack divisibility guard in dispatcher_run_mx_gemm() previously returned -1 (bad-args convention). -2 is semantically more accurate for "shape unsupported by the selected warp tile", consistent with the IsSupportedArguments-style rejects used elsewhere in the ctypes layer. Update the comment and the function docstring to match. 2. fp8 OCP/FNUZ guard added at MxGemmDispatcherLib.run() level assert_fp8_ocp_supported() was only called in GpuMxGemmRunner.__init__. A direct caller of MxGemmDispatcherLib.run() with fp8 on an FNUZ arch would bypass it and silently diverge from the numpy reference. Add optional dtype/arch params to MxGemmDispatcherLib.__init__(); if dtype=="fp8", call assert_fp8_ocp_supported at run() entry. Backward compatible (both params default to None; existing call sites unchanged). All 28 CPU tests pass; no test asserted on the specific -1 return code. --- .../dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp | 9 +++++---- .../composablekernel/dispatcher/python/mx_gemm_utils.py | 9 ++++++++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 25eb4c07ebda..7b8899e64e9d 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -102,7 +102,7 @@ void dispatcher_cleanup() { g_initialized = false; } * builds HostTensors from them, runs the SAME preShuffleScaleBuffer_gfx950 path * as the Old-TE profiler, uploads everything, and calls SelectedKernel::launch. * - * Returns 0 ok, -1 HIP/bad-args, -2 kernel reports unsupported args. + * Returns 0 ok, -1 HIP/bad-args, -2 unsupported shape (divisibility) or kernel rejects. */ int dispatcher_run_mx_gemm(const void* A, const void* B, @@ -208,8 +208,9 @@ int dispatcher_run_mx_gemm(const void* A, // scale_k_size/k_xdl_pack). If M/N/scale_k are not exact multiples of their // pack factors, that division silently truncates the shuffled buffers, so the // pre-shuffle would read/write past valid data and the kernel would consume a - // corrupt scale layout. Fail loudly instead (matches the -1 error convention - // used above). ---- + // corrupt scale layout. Return -2 (unsupported shape for the selected warp + // tile) rather than -1 (bad-args/HIP error), consistent with the IsSupportedArguments- + // style rejects used elsewhere in the dispatcher ctypes layer. ---- if(m % m_xdl_pack != 0 || n % n_xdl_pack != 0 || scale_k_size % k_xdl_pack != 0) { std::cerr << "dispatcher_run_mx_gemm: M, N, and scale_k (=K/32) must be divisible by the " @@ -217,7 +218,7 @@ int dispatcher_run_mx_gemm(const void* A, << m_xdl_pack << ", n_xdl_pack=" << n_xdl_pack << ", k_xdl_pack=" << k_xdl_pack << ") for the selected warp tile; got M=" << m << ", N=" << n << ", scale_k=" << scale_k_size << "\n"; - return -1; + return -2; } // ---- Build unshuffled scale HostTensors from the incoming raw e8m0 bytes. diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index 0434394f740e..d635f00ebee7 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -314,11 +314,13 @@ class MxGemmResult: class MxGemmDispatcherLib: - def __init__(self, so_path: Path): + def __init__(self, so_path: Path, dtype: Optional[str] = None, arch: Optional[str] = None): self.so_path = Path(so_path) if not self.so_path.exists(): raise FileNotFoundError(f"mx_gemm .so not found: {self.so_path}") self._lib = ctypes.CDLL(str(self.so_path)) + self._dtype = dtype + self._arch = arch self._setup() # Contract exposes both dispatcher_initialize() and dispatcher_init(). init = getattr(self._lib, "dispatcher_initialize", None) @@ -361,6 +363,11 @@ def kernel_count(self) -> int: return 1 def run(self, A, B, C, scale_a, scale_b, prob: MxGemmProblem) -> float: + # Guard: fp8 OCP bytes only match the device on gfx950/gfx12. If the lib + # was constructed with dtype="fp8", fail loudly now rather than silently + # diverging from the numpy reference on an FNUZ arch. + if self._dtype == "fp8": + assert_fp8_ocp_supported(self._arch or _DEFAULT_ARCH) A = np.ascontiguousarray(A) B = np.ascontiguousarray(B) # C is the OUTPUT buffer: the kernel writes results into it in place via a From 5ecc48e65c2eb52127f59d86b93ad12671b75ff8 Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Fri, 17 Jul 2026 00:06:22 -0400 Subject: [PATCH 06/14] fix(ck-tile): match Old-TE compile flags in mx_gemm bridge _compile_kernel The bridge runtime kernel compile path _compile_kernel built the byte-identical mx_gemm device kernel with only "-O3 -std=c++17", while Old-TE compiles the same kernel with the CK global AMDGPU codegen flags (C++20 + -mllvm tuning + -fno-offload-uniform-block + --offload-compress). The missing flags changed inlining/register allocation and occupancy, making bridge-built fp8 kernels run ~30% slower than Old-TE in production (700 TF vs 1076 TF at 4096^3 on a byte-identical 128x128 fp8 kernel). Add the exact flag set from projects/composablekernel/CMakeLists.txt (add_compile_options: -std=c++20, -mllvm -amdgpu-early-inline-all=true, -mllvm -amdgpu-function-calls=false, -mllvm --lsr-drop-solution=1, -mllvm -enable-post-misched=0, -mllvm -amdgpu-coerce-illegal-types=1, -fno-offload-uniform-block) plus mx_gemm/CMakeLists.txt's --offload-compress, matching the pattern already used by the gemm bridge's _tile_engine_codegen_flags. This brings fp8 to parity (measured +4.8% at 4096^3 on the same kernel). --- .../dispatcher/python/mx_gemm_utils.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index d635f00ebee7..d645abbdba63 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -607,7 +607,19 @@ def _compile_kernel(hpp: Path, so: Path, arch: str) -> bool: f"-I{_CK_ROOT}/tile_engine/ops/gemm/mx_gemm", ] cmd = [ - _HIPCC, "-shared", "-fPIC", "-O3", "-std=c++17", + # C++ standard and optimization flags must match Old-TE's mx_gemm build + # exactly, otherwise the byte-identical device kernel runs slower under + # the bridge (~30% on fp8). The flags below mirror the global + # add_compile_options() set in projects/composablekernel/CMakeLists.txt + # (CK_CXX_STANDARD=20, the -mllvm tuning flags, -fno-offload-uniform-block) + # plus mx_gemm/CMakeLists.txt's --offload-compress. + _HIPCC, "-shared", "-fPIC", "-O3", "-std=c++20", + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-mllvm", "-amdgpu-coerce-illegal-types=1", + "-fno-offload-uniform-block", "--offload-compress", *inc, "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f"-include{hpp}", "-D__HIP_PLATFORM_AMD__", f"--offload-arch={arch}", f'-DGFX_ARCH="{arch}"', From 09327f25a5aa49f205450034cb347a154281c1c3 Mon Sep 17 00:00:00 2001 From: ozturkosu Date: Fri, 17 Jul 2026 00:09:47 -0400 Subject: [PATCH 07/14] fix(ck-tile): probe-gate mx_gemm coerce-illegal-types flag for toolchain safety The previous commit added -mllvm -amdgpu-coerce-illegal-types=1 unconditionally, but that LLVM option is rejected by clang on ROCm 7.2 ("Unknown command line argument"), hard-failing the kernel compile. CK's CMake only adds that flag when check_cxx_compiler_flag passes; the shipping gemm bridge mirrors this via a probe. Restructure _compile_kernel to use the same pattern as gemm_utils._tile_engine_codegen_flags: an unconditional core set (-amdgpu-early-inline-all, -amdgpu-function-calls, --lsr-drop-solution, -enable-post-misched, -fno-offload-uniform-block, --offload-compress) plus a probe-gated -amdgpu-coerce-illegal-types=1 added only when hipcc accepts it. Keeps the bridge matched to Old-TE on toolchains that support the flag and buildable on those that do not. Verified: builds on ROCm 7.2/gfx950 and brings fp8 128x128 to parity with Old-TE at 4096^3. --- .../dispatcher/python/mx_gemm_utils.py | 74 ++++++++++++++++--- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index d645abbdba63..ea4627dfe1c9 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -29,6 +29,7 @@ import concurrent.futures import ctypes +import functools import json import logging import os @@ -599,6 +600,63 @@ def _generate_kernel(cfg: MxGemmKernelConfig, headers_dir: Path) -> Optional[Pat return hpp if hpp.exists() else None +# AMDGPU codegen / optimization flags that Old-TE compiles the mx_gemm device +# kernel with. They MUST match Old-TE flag-for-flag: the generated device-kernel +# header is byte-identical, so any flag difference (inlining, register +# allocation, occupancy) shows up as a large perf gap. Compiling with only +# "-O3 -std=c++17" made bridge fp8 kernels run ~30% slower than Old-TE. +# +# Source of truth: projects/composablekernel/CMakeLists.txt add_compile_options() +# (inherited by tile_engine/ops/gemm/mx_gemm) + mx_gemm/CMakeLists.txt's +# --offload-compress. This mirrors gemm_utils._tile_engine_codegen_flags so the +# two bridges stay consistent. +# +# Unconditional set (CK CMake adds these on every supported toolchain): +_MX_CODEGEN_FLAGS = ( + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-fno-offload-uniform-block", + "--offload-compress", +) +# Probe-gated set: CK's CMake only adds these when check_cxx_compiler_flag +# passes (newer -mllvm options some clang builds reject, e.g. ROCm 7.2 does not +# know -amdgpu-coerce-illegal-types). Mirror that probe so the bridge matches +# Old-TE wherever the compiler accepts it and stays buildable where it does not. +_MX_PROBED_CODEGEN_FLAGS = ( + ("-mllvm", "-amdgpu-coerce-illegal-types=1"), +) + + +@functools.lru_cache(maxsize=None) +def _hipcc_accepts(flag_tuple: Tuple[str, ...]) -> bool: + """Mirror CMake check_cxx_compiler_flag: does hipcc compile a trivial TU with + these flags? Cached so the probe runs at most once per distinct flag set.""" + try: + with tempfile.TemporaryDirectory() as d: + src = Path(d) / "probe.cpp" + src.write_text("int main(){}\n") + r = subprocess.run( + [_HIPCC, *flag_tuple, "-c", str(src), "-o", str(Path(d) / "probe.o")], + capture_output=True, timeout=120, + ) + return r.returncode == 0 + except Exception: + return False + + +@functools.lru_cache(maxsize=1) +def _mx_codegen_flags() -> Tuple[str, ...]: + """Old-TE's mx_gemm codegen flags plus any probe-gated flags the compiler + accepts -- the exact backend flag set the TE benchmark TU is built with.""" + flags = list(_MX_CODEGEN_FLAGS) + for pair in _MX_PROBED_CODEGEN_FLAGS: + if _hipcc_accepts(pair): + flags = list(pair) + flags + return tuple(flags) + + def _compile_kernel(hpp: Path, so: Path, arch: str) -> bool: inc = [ f"-I{_CK_ROOT}/include", f"-I{_CK_ROOT}", @@ -607,19 +665,11 @@ def _compile_kernel(hpp: Path, so: Path, arch: str) -> bool: f"-I{_CK_ROOT}/tile_engine/ops/gemm/mx_gemm", ] cmd = [ - # C++ standard and optimization flags must match Old-TE's mx_gemm build - # exactly, otherwise the byte-identical device kernel runs slower under - # the bridge (~30% on fp8). The flags below mirror the global - # add_compile_options() set in projects/composablekernel/CMakeLists.txt - # (CK_CXX_STANDARD=20, the -mllvm tuning flags, -fno-offload-uniform-block) - # plus mx_gemm/CMakeLists.txt's --offload-compress. + # -std=c++20 matches CK_CXX_STANDARD; codegen flags match Old-TE (see + # _mx_codegen_flags above). Without them the byte-identical fp8 device + # kernel ran ~30% slower than Old-TE. _HIPCC, "-shared", "-fPIC", "-O3", "-std=c++20", - "-mllvm", "-amdgpu-early-inline-all=true", - "-mllvm", "-amdgpu-function-calls=false", - "-mllvm", "--lsr-drop-solution=1", - "-mllvm", "-enable-post-misched=0", - "-mllvm", "-amdgpu-coerce-illegal-types=1", - "-fno-offload-uniform-block", "--offload-compress", + *_mx_codegen_flags(), *inc, "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f"-include{hpp}", "-D__HIP_PLATFORM_AMD__", f"--offload-arch={arch}", f'-DGFX_ARCH="{arch}"', From c582f48ff9a712d02616c186945b6b2b19140779 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:13:52 -0400 Subject: [PATCH 08/14] =?UTF-8?q?fix(ck-tile):=20mx=5Fgemm=20bridge=20?= =?UTF-8?q?=E2=80=94=20require=20GPU=20arch=20via=20get=5Farch,=20do=20not?= =?UTF-8?q?=20default=20to=20a=20specific=20arch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 3 +- .../codegen/unified_mx_gemm_codegen.py | 8 +++- .../dispatcher/python/mx_gemm_utils.py | 40 ++++++++++++++----- 3 files changed, 39 insertions(+), 12 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 7b8899e64e9d..56447a9b026f 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -55,7 +55,8 @@ #include "common/utils.hpp" #ifndef GFX_ARCH -#define GFX_ARCH "gfx950" +#error \ + "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=); do not default to a specific GPU architecture." #endif static bool g_initialized = false; diff --git a/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py index 5ccc34a0311b..c05882d5153f 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_mx_gemm_codegen.py @@ -189,10 +189,16 @@ def _make_builder(cfg: dict) -> Iterator["object"]: cfg_path = work_dir / "mx_gemm_codegen_config.json" cfg_path.write_text(json.dumps(tmp_cfg)) + gpu_target = cfg.get("gpu_target") + if not gpu_target: + raise ValueError( + "mx_gemm codegen requires an explicit 'gpu_target' in the config; " + "do not default to a specific GPU architecture." + ) yield MxGemmKernelBuilder( KERNEL_NAME_PREFIX, work_dir, - cfg.get("gpu_target", "gfx950"), + gpu_target, cfg["datatype"], cfg["layout"], config_json=str(cfg_path), diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index ea4627dfe1c9..a56682da8311 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -51,7 +51,25 @@ # ck root == the composablekernel dir (three levels up from dispatcher/python). _CK_ROOT = Path(__file__).parent.parent.parent _HIPCC = os.environ.get("CK_TILE_HIPCC", "/opt/rocm/bin/hipcc") -_DEFAULT_ARCH = "gfx950" + + +def _get_arch() -> str: + """Detect GPU arch via rocminfo and validate; raise on failure. Never defaults.""" + import subprocess + arch = "" + try: + out = subprocess.check_output(["rocminfo"], text=True, stderr=subprocess.DEVNULL) + for line in out.splitlines(): + if "Name:" in line and "gfx" in line: + arch = line.split()[-1].strip(); break + except Exception: + arch = "" + if not arch: + raise RuntimeError("Could not detect GPU architecture from rocminfo; refusing to default. Pass gfx_arch explicitly.") + _supported = ("gfx90a", "gfx942", "gfx950") + if arch not in _supported: + raise ValueError(f"Unsupported GPU architecture {arch!r}; supported: {list(_supported)}") + return arch # MX GEMM scales every 32 K-elements with one e8m0 byte. SCALE_BLOCK = 32 @@ -113,7 +131,7 @@ def fp8_ocp_is_default_for_arch(arch: str) -> bool: return a.startswith("gfx950") or a.startswith("gfx12") -def assert_fp8_ocp_supported(arch: str = _DEFAULT_ARCH) -> None: +def assert_fp8_ocp_supported(arch: Optional[str] = None) -> None: """Fail loudly if the fp8 OCP codec would silently disagree with the device. The fp8 quantize/dequantize helpers emit/decode OCP e4m3 bytes; that is only @@ -121,6 +139,7 @@ def assert_fp8_ocp_supported(arch: str = _DEFAULT_ARCH) -> None: codec precondition above). Raise instead of producing bytes that mismatch a FNUZ-built kernel. """ + arch = arch or _get_arch() if not fp8_ocp_is_default_for_arch(arch): raise ValueError( f"fp8 mx_gemm reference is OCP e4m3 only, but arch '{arch}' defaults to " @@ -166,7 +185,7 @@ def float_to_e8m0(scale) -> np.ndarray: class MxGemmKernelConfig: datatype: str = "fp8" # fp8 | fp4 layout: str = "rcr" # a/b/c ; only rcr supported by mx_gemm - gpu_target: str = _DEFAULT_ARCH + gpu_target: Optional[str] = None pipeline: str = "comp_async" epilogue: str = "cshuffle" scheduler: str = "intrawave" @@ -194,7 +213,7 @@ def to_codegen_config(self) -> dict: return { "datatype": self.datatype, "layout": self.layout, - "gpu_target": self.gpu_target, + "gpu_target": self.gpu_target or _get_arch(), "pipeline": self.pipeline, "epilogue": self.epilogue, "scheduler": self.scheduler, @@ -368,7 +387,7 @@ def run(self, A, B, C, scale_a, scale_b, prob: MxGemmProblem) -> float: # was constructed with dtype="fp8", fail loudly now rather than silently # diverging from the numpy reference on an FNUZ arch. if self._dtype == "fp8": - assert_fp8_ocp_supported(self._arch or _DEFAULT_ARCH) + assert_fp8_ocp_supported(self._arch or _get_arch()) A = np.ascontiguousarray(A) B = np.ascontiguousarray(B) # C is the OUTPUT buffer: the kernel writes results into it in place via a @@ -527,11 +546,12 @@ def mx_gemm_reference(A_deq, B_deq, scale_a_byte, scale_b_byte, prob: MxGemmProb class GpuMxGemmRunner: - def __init__(self, so_path: Path, dtype: str = "fp8", arch: str = _DEFAULT_ARCH): + def __init__(self, so_path: Path, dtype: str = "fp8", arch: Optional[str] = None): # The fp8 numpy reference (quantize_fp8/dequantize_fp8) emits OCP e4m3 # bytes; that only matches the device when the lib was compiled with # CK_TILE_USE_OCP_FP8 == 1 (the gfx950 default). Guard loudly so an # FNUZ arch can't silently diverge from the reference. + arch = arch or _get_arch() if dtype == "fp8": assert_fp8_ocp_supported(arch) self._lib = MxGemmDispatcherLib(so_path) @@ -686,7 +706,7 @@ def _compile_kernel(hpp: Path, so: Path, arch: str) -> bool: def setup_multiple_mx_gemm_dispatchers( configs: List[MxGemmKernelConfig], output_dir: Optional[Path] = None, - gfx_arch: str = _DEFAULT_ARCH, + gfx_arch: Optional[str] = None, parallel: bool = True, max_workers: Optional[int] = None, ) -> List[Optional[Path]]: @@ -694,7 +714,7 @@ def setup_multiple_mx_gemm_dispatchers( `configs` (None on failure). Dedups by .name; per-arch .so cache.""" if not configs: return [] - arch = gfx_arch or _DEFAULT_ARCH + arch = gfx_arch or _get_arch() base = Path(output_dir) if output_dir else Path(tempfile.mkdtemp(prefix="mx_gemm_bridge_")) headers = base / "generated_kernels" libs = base / "libs" @@ -739,7 +759,7 @@ def build_one(idx: int, cfg: MxGemmKernelConfig): return results -def default_fp8_config(gfx_arch: str = _DEFAULT_ARCH) -> MxGemmKernelConfig: +def default_fp8_config(gfx_arch: Optional[str] = None) -> MxGemmKernelConfig: return MxGemmKernelConfig( datatype="fp8", layout="rcr", gpu_target=gfx_arch, pipeline="comp_async", epilogue="cshuffle", scheduler="intrawave", @@ -748,7 +768,7 @@ def default_fp8_config(gfx_arch: str = _DEFAULT_ARCH) -> MxGemmKernelConfig: ) -def default_fp4_config(gfx_arch: str = _DEFAULT_ARCH) -> MxGemmKernelConfig: +def default_fp4_config(gfx_arch: Optional[str] = None) -> MxGemmKernelConfig: cfg = default_fp8_config(gfx_arch) cfg.datatype = "fp4" cfg._name_cache = None From f1944add260edede0e572594062bf5361c1a8d14 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:56:13 -0400 Subject: [PATCH 09/14] =?UTF-8?q?fix(ck-tile):=20mx=5Fgemm=20bridge=20?= =?UTF-8?q?=E2=80=94=20assert=20gfx950-only=20at=20compile=20time=20for=20?= =?UTF-8?q?preShuffleScaleBuffer=5Fgfx950?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 56447a9b026f..ed847c124b30 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -43,6 +43,7 @@ #include #include #include +#include #include // Kernel header force-included via -include. Brings in ck_tile core + the @@ -59,6 +60,14 @@ "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=); do not default to a specific GPU architecture." #endif +// The MX (microscaling) block-scale pre-shuffle below uses ck_tile's gfx950-only +// preShuffleScaleBuffer_gfx950 host helper, so this bridge is inherently gfx950 +// only. Make that scope explicit: fail the build clearly on any other arch +// instead of silently mis-calling the gfx950 helper. +static_assert(std::string_view(GFX_ARCH) == "gfx950", + "mx_gemm dispatcher bridge is gfx950-only (uses preShuffleScaleBuffer_gfx950); " + "build with -DGFX_ARCH=gfx950."); + static bool g_initialized = false; namespace { From d59b62692bd5d5908feb83ff9f6ce01edc973f8f Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:59:11 -0400 Subject: [PATCH 10/14] =?UTF-8?q?fix(ck-tile):=20mx=5Fgemm=20bridge=20?= =?UTF-8?q?=E2=80=94=20add=20runtime=20gfx950=20device=20guard=20for=20pre?= =?UTF-8?q?ShuffleScaleBuffer=5Fgfx950?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index ed847c124b30..f0ed8be59f11 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -172,6 +172,26 @@ int dispatcher_run_mx_gemm(const void* A, return -1; } + // MX block-scale pre-shuffle is gfx950-only (preShuffleScaleBuffer_gfx950). + // The compile-time static_assert already pins the build to gfx950; also guard + // at runtime so a gfx950-built .so run on a non-gfx950 device fails clearly + // instead of launching an arch-mismatched kernel. + { + int dev = 0; + hipDeviceProp_t props{}; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) + { + std::cerr << "dispatcher_run_mx_gemm: could not query device architecture\n"; + return -1; + } + if(std::string_view(props.gcnArchName).substr(0, 6) != "gfx950") + { + std::cerr << "dispatcher_run_mx_gemm: MX GEMM is gfx950-only; running device is " + << props.gcnArchName << "\n"; + return -1; + } + } + const ALayout layout_a = ALayout{}; const BLayout layout_b = BLayout{}; const CLayout layout_c = CLayout{}; From 53bfd2a33c623854182c290960b48a2e69a4a6cf Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:32:48 -0400 Subject: [PATCH 11/14] =?UTF-8?q?fix(ck-tile):=20mx=5Fgemm=20bridge=20?= =?UTF-8?q?=E2=80=94=20address=20review=20(unsupported=20-2=20code,=20arch?= =?UTF-8?q?=20propagation=20in=20setup=5Fmultiple,=20grid-map=20contract)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 5 +- .../dispatcher/python/mx_gemm_utils.py | 57 +++++++++++++++---- 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index f0ed8be59f11..8cc3f8452a4a 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -167,9 +167,12 @@ int dispatcher_run_mx_gemm(const void* A, if(k_batch != 1) { // v1 scope: split-K (k_batch > 1) is not yet supported through this bridge. + // This is a v1 feature limitation, not a bad-argument/HIP error, so report + // -2 (unsupported) per the documented convention above -- the Python + // wrapper then surfaces it as "unsupported" rather than a generic error. std::cerr << "dispatcher_run_mx_gemm: only k_batch==1 is supported in v1, got " << k_batch << "\n"; - return -1; + return -2; } // MX block-scale pre-shuffle is gfx950-only (preShuffleScaleBuffer_gfx950). diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index a56682da8311..bf83364c7209 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -432,29 +432,48 @@ def __del__(self): # ============================================================================= +# Tolerance for on-grid membership in _map_grid_to_bytes. The quantization grid +# entries are all exactly representable in float32 (0, +/-0.5, +/-1.0, ...), so +# grid-exact inputs match to the bit; this epsilon only absorbs the last few ULPs +# of float-repr noise (e.g. a value stored via float64 -> float32). Any input +# farther than this from the nearest grid point is a genuine off-grid value and +# is rejected. Values like 0.5004 (0.0004 off the 0.5 grid point) exceed this and +# now raise, matching the documented "rejects off-grid values" contract. +_GRID_MATCH_EPS = np.float32(1e-4) + + def _map_grid_to_bytes(vals: np.ndarray, value_to_byte: dict, grid_name: str) -> np.ndarray: """Vectorized exact-grid float -> byte code lookup (uint8, flattened). The inputs are drawn from a small fixed grid (see make_inputs), so instead of a per-element Python loop -- which dominates runtime for realistic M*K -- we - round to the grid resolution and index a sorted LUT with a single - searchsorted. Off-grid values are a caller contract violation, so raise - (mirrors the old dict-KeyError) rather than silently snapping to a neighbour. + snap each input to the nearest grid point via a single searchsorted and index + a sorted LUT. An input is accepted only if it lies within _GRID_MATCH_EPS + (1e-4) of that nearest grid point; anything farther off (e.g. 0.5004) is a + caller contract violation and raises (mirrors the old dict-KeyError) rather + than silently snapping to a neighbour. """ keys = np.array(sorted(value_to_byte), dtype=np.float32) byts = np.array([value_to_byte[float(k)] for k in keys], dtype=np.uint8) - r = np.round(np.asarray(vals, dtype=np.float32), 3) - r = r + np.float32(0.0) # collapse -0.0 -> +0.0 so it matches the 0.0 key - flat = r.reshape(-1) + v = np.asarray(vals, dtype=np.float32) + np.float32(0.0) # collapse -0.0 -> +0.0 + flat = v.reshape(-1) - idx = np.clip(np.searchsorted(keys, flat), 0, keys.size - 1) - if not np.all(keys[idx] == flat): - bad = flat[keys[idx] != flat] + # Nearest grid key for each value: searchsorted gives the insertion point; + # the closer of the two straddling keys is the nearest grid point. + pos = np.clip(np.searchsorted(keys, flat), 1, keys.size - 1) + left = keys[pos - 1] + right = keys[pos] + nearest = np.where(np.abs(flat - left) <= np.abs(flat - right), left, right) + + if not np.all(np.abs(flat - nearest) <= _GRID_MATCH_EPS): + off = np.abs(flat - nearest) > _GRID_MATCH_EPS + bad = flat[off] raise KeyError( - f"{grid_name}: value(s) not on the exact quantization grid: " - f"{np.unique(bad)[:8].tolist()}" + f"{grid_name}: value(s) not on the exact quantization grid " + f"(tolerance {float(_GRID_MATCH_EPS):g}): {np.unique(bad)[:8].tolist()}" ) + idx = np.searchsorted(keys, nearest) return byts[idx] @@ -715,6 +734,22 @@ def setup_multiple_mx_gemm_dispatchers( if not configs: return [] arch = gfx_arch or _get_arch() + # When an explicit gfx_arch is passed, pin every config to it BEFORE computing + # cfg.name / running codegen. Otherwise cfg.to_codegen_config() falls back to + # `self.gpu_target or _get_arch()`, so codegen (and the cached name) would use + # the host-detected arch (or fail if rocminfo is missing) while the .so is + # compiled for `arch` -- an arch mismatch between the header and the binary. + # Resetting _name_cache forces the name to be recomputed for the chosen arch. + if gfx_arch is not None: + _supported = ("gfx90a", "gfx942", "gfx950") + if gfx_arch not in _supported: + raise ValueError( + f"Unsupported GPU architecture {gfx_arch!r}; supported: {list(_supported)}" + ) + for c in configs: + if c.gpu_target != gfx_arch: + c.gpu_target = gfx_arch + c._name_cache = None base = Path(output_dir) if output_dir else Path(tempfile.mkdtemp(prefix="mx_gemm_bridge_")) headers = base / "generated_kernels" libs = base / "libs" From 0c40bdaa7325e3e5ecda8040c34c616c229d38b5 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:47:37 -0400 Subject: [PATCH 12/14] =?UTF-8?q?perf(ck-tile):=20mx=5Fgemm=20bridge=20?= =?UTF-8?q?=E2=80=94=20sweep=20buffer=20reuse,=20vectorized=20fp8=20dequan?= =?UTF-8?q?t,=20parity=20bench=20defaults?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses tough-review findings #2, #3, #5: - #2: dispatcher_run_mx_gemm now reuses persistent device allocations across calls (opt-out via CK_TILE_MX_GEMM_REUSE_BUFFERS=0), eliminating per-call hipMalloc/hipFree churn that slowed throughput sweeps and perturbed thermal steady state. A/B/scales are still uploaded and C still zeroed every call, so per-shape data is never stale; cache is freed in dispatcher_cleanup(). - #3: dequantize_fp8 replaced its interpreted per-element loop with a 256-entry LUT gather (consistent with dequantize_fp4_packed); foreign bytes map to NaN. - #5: bench warmup/repeat defaults raised to 50/100 to match the bridge parity-sweep convention for apples-to-apples bridge-vs-Old-TE timing. --- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 162 +++++++++++++++--- .../dispatcher/python/mx_gemm_utils.py | 19 +- 2 files changed, 155 insertions(+), 26 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index 8cc3f8452a4a..a8dc3a6952b2 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -43,6 +43,8 @@ #include #include #include +#include +#include #include #include @@ -80,6 +82,58 @@ int env_int(const char* name, int fallback) return std::atoi(v); } +bool env_flag(const char* name, bool fallback) +{ + const char* v = std::getenv(name); + if(!v) + return fallback; + // Accept 1/true/on/yes (case-insensitive first char) as true; 0/false/off as + // false. + const char c = v[0]; + return !(c == '0' || c == 'f' || c == 'F' || c == 'n' || c == 'N'); +} + +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE +// ---- Optional persistent device-buffer cache for sweep mode. ---- +// By default (opt-out) each dispatcher_run_mx_gemm() allocates and frees its +// device buffers, which is the simplest lifetime model but churns hipMalloc/ +// hipFree on every call. For throughput sweeps that call run() thousands of times +// this alloc/free churn both slows the sweep and perturbs the thermal steady +// state relative to a persistent-buffer profiler. When +// CK_TILE_MX_GEMM_REUSE_BUFFERS is enabled (default on), the buffers below are +// kept alive between calls and only (re)allocated when a larger capacity is +// needed. Correctness is unchanged: A/B/scales are still uploaded and C is still +// zeroed every call, so per-shape data is never stale; only the allocations are +// reused. Set CK_TILE_MX_GEMM_REUSE_BUFFERS=0 to restore per-call alloc/free. +struct ReusableDeviceBuffer +{ + std::unique_ptr mem; + std::size_t capacity = 0; + + // Ensure the buffer holds at least `bytes`, reallocating (growing) only when + // needed. Returns the device pointer. + void* ensure(std::size_t bytes) + { + if(!mem || capacity < bytes) + { + mem = std::make_unique(bytes); + capacity = bytes; + } + return mem->GetDeviceBuffer(); + } +}; + +struct MxGemmBufferCache +{ + ReusableDeviceBuffer a, b, c, scale_a, scale_b; +}; + +// One cache for the single compiled-in kernel. dispatcher is used single-kernel +// per .so and callers drive it from one thread at a time (the Python ctypes +// wrapper holds the GIL across the call), so a plain static is sufficient. +MxGemmBufferCache g_buf_cache; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE + } // namespace extern "C" { @@ -102,7 +156,15 @@ const char* dispatcher_get_kernel_name() } int dispatcher_get_kernel_count() { return 1; } -void dispatcher_cleanup() { g_initialized = false; } +void dispatcher_cleanup() +{ + g_initialized = false; +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE + // Release any persistent sweep-mode device buffers so cleanup() frees all + // device memory this lib allocated. + g_buf_cache = MxGemmBufferCache{}; +#endif +} /** * Run microscaling GEMM: C[M,N] = (A[M,K] * scale_a) . (B[K,N] * scale_b). @@ -312,31 +374,82 @@ int dispatcher_run_mx_gemm(const void* A, std::memcpy(a_m_k.mData.data(), A, a_bytes); std::memcpy(b_k_n.mData.data(), B, b_bytes); + const std::size_t scale_a_bytes = scale_a_shuffled.get_element_space_size_in_bytes(); + const std::size_t scale_b_bytes = scale_b_shuffled.get_element_space_size_in_bytes(); + + // Sweep mode reuses persistent device allocations across calls (default on) to + // avoid per-call hipMalloc/hipFree churn; opt out via + // CK_TILE_MX_GEMM_REUSE_BUFFERS=0 to get simple per-call alloc/free. + const bool reuse_buffers = env_flag("CK_TILE_MX_GEMM_REUSE_BUFFERS", true); + float exec_time = 0.0f; try { - // ---- Device buffers, sized via HostTensor (packing-correct). ---- - ck_tile::DeviceMem a_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_dev_buf(c_m_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem scale_a_dev_buf(scale_a_shuffled.get_element_space_size_in_bytes()); - ck_tile::DeviceMem scale_b_dev_buf(scale_b_shuffled.get_element_space_size_in_bytes()); - - a_dev_buf.ToDevice(a_m_k.data()); - b_dev_buf.ToDevice(b_k_n.data()); - c_dev_buf.SetZero(); - scale_a_dev_buf.ToDevice(scale_a_shuffled.data()); - scale_b_dev_buf.ToDevice(scale_b_shuffled.data()); + // Device pointers, resolved below either from the persistent cache (reuse + // path) or from call-local DeviceMem (per-call path). Both paths upload + // A/B/scales and zero C every call, so per-shape data is never stale. + void* a_ptr = nullptr; + void* b_ptr = nullptr; + void* c_ptr = nullptr; + void* scale_a_ptr = nullptr; + void* scale_b_ptr = nullptr; + + // Call-local buffers kept in scope for the non-reuse path; empty (never + // allocated) when reusing the cache. + std::unique_ptr a_local, b_local, c_local, sa_local, sb_local; + + auto upload = [](void* dst, const void* src, std::size_t bytes) { + if(hipMemcpy(dst, src, bytes, hipMemcpyHostToDevice) != hipSuccess) + throw std::runtime_error("hipMemcpy H2D failed"); + }; + + if(reuse_buffers) + { + a_ptr = g_buf_cache.a.ensure(a_bytes); + b_ptr = g_buf_cache.b.ensure(b_bytes); + c_ptr = g_buf_cache.c.ensure(c_bytes); + scale_a_ptr = g_buf_cache.scale_a.ensure(scale_a_bytes); + scale_b_ptr = g_buf_cache.scale_b.ensure(scale_b_bytes); + + // Upload exactly the shape's byte count (cached buffers may be larger). + upload(a_ptr, a_m_k.data(), a_bytes); + upload(b_ptr, b_k_n.data(), b_bytes); + upload(scale_a_ptr, scale_a_shuffled.data(), scale_a_bytes); + upload(scale_b_ptr, scale_b_shuffled.data(), scale_b_bytes); + if(hipMemset(c_ptr, 0, c_bytes) != hipSuccess) + throw std::runtime_error("hipMemset C failed"); + } + else + { + // ---- Call-local device buffers, sized via HostTensor (packing-correct). ---- + a_local = std::make_unique(a_bytes); + b_local = std::make_unique(b_bytes); + c_local = std::make_unique(c_bytes); + sa_local = std::make_unique(scale_a_bytes); + sb_local = std::make_unique(scale_b_bytes); + + a_local->ToDevice(a_m_k.data()); + b_local->ToDevice(b_k_n.data()); + c_local->SetZero(); + sa_local->ToDevice(scale_a_shuffled.data()); + sb_local->ToDevice(scale_b_shuffled.data()); + + a_ptr = a_local->GetDeviceBuffer(); + b_ptr = b_local->GetDeviceBuffer(); + c_ptr = c_local->GetDeviceBuffer(); + scale_a_ptr = sa_local->GetDeviceBuffer(); + scale_b_ptr = sb_local->GetDeviceBuffer(); + } // ---- Build MxGemmHostArgs in the exact profiler argument order: // a_ptr, scale_a, b_ptr, scale_b, ds{}, c_ptr, split_k, m, n, k, // {stride_a}, {stride_b}, ds_strides{}, stride_c. ---- - MxGemmHostArgs gemm_args({a_dev_buf.GetDeviceBuffer()}, - {scale_a_dev_buf.GetDeviceBuffer()}, - {b_dev_buf.GetDeviceBuffer()}, - {scale_b_dev_buf.GetDeviceBuffer()}, + MxGemmHostArgs gemm_args({a_ptr}, + {scale_a_ptr}, + {b_ptr}, + {scale_b_ptr}, {}, - c_dev_buf.GetDeviceBuffer(), + c_ptr, static_cast(k_batch), m, n, @@ -347,8 +460,12 @@ int dispatcher_run_mx_gemm(const void* A, stride_c); const bool do_time = (time_ms != nullptr); - const int warmup = do_time ? env_int("CK_TILE_BENCH_WARMUP", 20) : 0; - const int repeat = do_time ? env_int("CK_TILE_BENCH_REPEAT", 50) : 1; + // Defaults match the bridge parity-sweep convention (warmup 50 / repeat + // 100) so bridge-vs-Old-TE timings are apples-to-apples out of the box; + // both remain env-overridable for custom sweeps (set identical values on + // both sides). + const int warmup = do_time ? env_int("CK_TILE_BENCH_WARMUP", 50) : 0; + const int repeat = do_time ? env_int("CK_TILE_BENCH_REPEAT", 100) : 1; // stream_config field order: stream_id, time_kernel, log_level, cold_niters // (warmup), nrepeat, is_gpu_timer, flush_cache, rotating_count. ck_tile::stream_config stream_cfg{nullptr, do_time, 0, warmup, repeat, false, false, 1}; @@ -360,7 +477,10 @@ int dispatcher_run_mx_gemm(const void* A, return -2; } - c_dev_buf.FromDevice(c_m_n.data()); + // D2H from the resolved c_ptr (works for both the cached and per-call + // paths); copy exactly c_bytes since a cached buffer may be larger. + if(hipMemcpy(c_m_n.data(), c_ptr, c_bytes, hipMemcpyDeviceToHost) != hipSuccess) + throw std::runtime_error("hipMemcpy D2H failed"); } catch(const std::exception& e) { diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index bf83364c7209..dce3a76b4869 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -116,6 +116,11 @@ def _get_arch() -> str: -3.0: 0xC4, -4.0: 0xC8, -6.0: 0xCC, } _FP8_OCP_BYTE_TO_VALUE = {b: v for v, b in _FP8_OCP_VALUE_TO_BYTE.items()} +# 256-entry byte -> value LUT for vectorized dequantize_fp8. Untested bytes stay +# NaN (they cannot come from quantize_fp8, so a NaN result flags a foreign byte). +_FP8_OCP_BYTE_LUT = np.full(256, np.nan, dtype=np.float32) +for _b, _v in _FP8_OCP_BYTE_TO_VALUE.items(): + _FP8_OCP_BYTE_LUT[_b] = np.float32(_v) def fp8_ocp_is_default_for_arch(arch: str) -> bool: @@ -491,11 +496,15 @@ def quantize_fp8(vals: np.ndarray) -> np.ndarray: def dequantize_fp8(bytes_arr: np.ndarray) -> np.ndarray: - flat = np.asarray(bytes_arr, dtype=np.uint8).reshape(-1) - out = np.empty(flat.shape, dtype=np.float32) - for i, b in enumerate(flat): - out[i] = _FP8_OCP_BYTE_TO_VALUE[int(b)] - return out.reshape(bytes_arr.shape) + """Raw OCP e4m3 bytes (uint8) -> float grid values, preserving shape. + + Vectorized LUT-gather (consistent with dequantize_fp4_packed) instead of a + per-element Python loop, so this stays fast if ever applied to large buffers. + Bytes outside the tested grid map to NaN (they cannot appear from quantize_fp8, + which only emits grid codes, so a NaN here flags a corrupt/foreign buffer). + """ + b = np.asarray(bytes_arr, dtype=np.uint8) + return _FP8_OCP_BYTE_LUT[b] def quantize_fp4_packed(vals: np.ndarray) -> np.ndarray: From 3cb9b927cf1b6f4ecc92dca0295e2b5a49b0bac4 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:36:57 -0400 Subject: [PATCH 13/14] fix(ck-tile): address Copilot review on mx_gemm bridge (arch/CMake/tests/cache) - CMake: define GFX_ARCH="gfx950" on both dispatcher_mx_gemm_lib branches. mx_gemm_ctypes_lib.cpp #errors when GFX_ARCH is unset, so the CMake build path (kernel and no-kernel fallback) failed to compile; gfx950 is the only arch this op supports (static_assert on preShuffleScaleBuffer_gfx950). - Python arch validation: restrict _get_arch()/setup_multiple_mx_gemm_dispatchers to gfx950 only, matching the C++ static_assert. A gfx942/gfx90a caller no longer passes detection only to fail later with a less actionable error. - Tests: pin gpu_target="gfx950" at the three to_codegen_config() call sites so the CPU-only suite never shells out to rocminfo on non-ROCm runners. - Thread-safety: remove the g_buf_cache device-buffer reuse path entirely and go back to per-call DeviceMem alloc/free (as the batched/multi_abd bridges do). The cache was a shared global with no lock; ctypes.CDLL releases the GIL so concurrent callers could race across upload->launch->D2H. Removing it (vs adding a mutex) also drops an unrequested, untested optimization. CPU tests: 27 passed, 1 skipped (codegen import unavailable off-node). --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 10 +- .../bindings/ctypes/mx_gemm_ctypes_lib.cpp | 147 +++--------------- .../dispatcher/python/mx_gemm_utils.py | 20 ++- .../dispatcher/tests/test_mx_gemm_bridge.py | 7 + 4 files changed, 55 insertions(+), 129 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 7c7b443914ec..a3743613275d 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -98,9 +98,13 @@ if(MX_GEMM_KERNEL_HEADERS) ) target_include_directories(dispatcher_mx_gemm_lib PRIVATE ${MX_GEMM_TE_INCLUDE_DIRS}) # The generated header only exports SelectedKernel/KERNEL_NAME/ScaleType/... - # under this define, so the force-include path requires it. + # under this define, so the force-include path requires it. GFX_ARCH is + # required by mx_gemm_ctypes_lib.cpp (#error if unset) and the bridge is + # gfx950-only (static_assert on preShuffleScaleBuffer_gfx950), so pin it to + # gfx950 here -- the only value this op can compile for. target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="gfx950" ) else() message(STATUS "No MX-GEMM kernel found for ctypes lib - building without kernel") @@ -111,6 +115,10 @@ else() ${MX_GEMM_TE_INCLUDE_DIRS} ) target_link_libraries(dispatcher_mx_gemm_lib PRIVATE hip::device) + # Even the no-kernel fallback compiles mx_gemm_ctypes_lib.cpp, which #errors + # if GFX_ARCH is unset; define it (gfx950, the only arch this op supports) so + # the fallback target builds instead of failing the whole configure/build. + target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE GFX_ARCH="gfx950") set_target_properties(dispatcher_mx_gemm_lib PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 17 diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp index a8dc3a6952b2..018ada94f625 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/mx_gemm_ctypes_lib.cpp @@ -43,7 +43,6 @@ #include #include #include -#include #include #include #include @@ -82,58 +81,6 @@ int env_int(const char* name, int fallback) return std::atoi(v); } -bool env_flag(const char* name, bool fallback) -{ - const char* v = std::getenv(name); - if(!v) - return fallback; - // Accept 1/true/on/yes (case-insensitive first char) as true; 0/false/off as - // false. - const char c = v[0]; - return !(c == '0' || c == 'f' || c == 'F' || c == 'n' || c == 'N'); -} - -#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE -// ---- Optional persistent device-buffer cache for sweep mode. ---- -// By default (opt-out) each dispatcher_run_mx_gemm() allocates and frees its -// device buffers, which is the simplest lifetime model but churns hipMalloc/ -// hipFree on every call. For throughput sweeps that call run() thousands of times -// this alloc/free churn both slows the sweep and perturbs the thermal steady -// state relative to a persistent-buffer profiler. When -// CK_TILE_MX_GEMM_REUSE_BUFFERS is enabled (default on), the buffers below are -// kept alive between calls and only (re)allocated when a larger capacity is -// needed. Correctness is unchanged: A/B/scales are still uploaded and C is still -// zeroed every call, so per-shape data is never stale; only the allocations are -// reused. Set CK_TILE_MX_GEMM_REUSE_BUFFERS=0 to restore per-call alloc/free. -struct ReusableDeviceBuffer -{ - std::unique_ptr mem; - std::size_t capacity = 0; - - // Ensure the buffer holds at least `bytes`, reallocating (growing) only when - // needed. Returns the device pointer. - void* ensure(std::size_t bytes) - { - if(!mem || capacity < bytes) - { - mem = std::make_unique(bytes); - capacity = bytes; - } - return mem->GetDeviceBuffer(); - } -}; - -struct MxGemmBufferCache -{ - ReusableDeviceBuffer a, b, c, scale_a, scale_b; -}; - -// One cache for the single compiled-in kernel. dispatcher is used single-kernel -// per .so and callers drive it from one thread at a time (the Python ctypes -// wrapper holds the GIL across the call), so a plain static is sufficient. -MxGemmBufferCache g_buf_cache; -#endif // CK_TILE_SINGLE_KERNEL_INCLUDE - } // namespace extern "C" { @@ -156,15 +103,7 @@ const char* dispatcher_get_kernel_name() } int dispatcher_get_kernel_count() { return 1; } -void dispatcher_cleanup() -{ - g_initialized = false; -#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE - // Release any persistent sweep-mode device buffers so cleanup() frees all - // device memory this lib allocated. - g_buf_cache = MxGemmBufferCache{}; -#endif -} +void dispatcher_cleanup() { g_initialized = false; } /** * Run microscaling GEMM: C[M,N] = (A[M,K] * scale_a) . (B[K,N] * scale_b). @@ -377,69 +316,30 @@ int dispatcher_run_mx_gemm(const void* A, const std::size_t scale_a_bytes = scale_a_shuffled.get_element_space_size_in_bytes(); const std::size_t scale_b_bytes = scale_b_shuffled.get_element_space_size_in_bytes(); - // Sweep mode reuses persistent device allocations across calls (default on) to - // avoid per-call hipMalloc/hipFree churn; opt out via - // CK_TILE_MX_GEMM_REUSE_BUFFERS=0 to get simple per-call alloc/free. - const bool reuse_buffers = env_flag("CK_TILE_MX_GEMM_REUSE_BUFFERS", true); - float exec_time = 0.0f; try { - // Device pointers, resolved below either from the persistent cache (reuse - // path) or from call-local DeviceMem (per-call path). Both paths upload - // A/B/scales and zero C every call, so per-shape data is never stale. - void* a_ptr = nullptr; - void* b_ptr = nullptr; - void* c_ptr = nullptr; - void* scale_a_ptr = nullptr; - void* scale_b_ptr = nullptr; - - // Call-local buffers kept in scope for the non-reuse path; empty (never - // allocated) when reusing the cache. - std::unique_ptr a_local, b_local, c_local, sa_local, sb_local; - - auto upload = [](void* dst, const void* src, std::size_t bytes) { - if(hipMemcpy(dst, src, bytes, hipMemcpyHostToDevice) != hipSuccess) - throw std::runtime_error("hipMemcpy H2D failed"); - }; - - if(reuse_buffers) - { - a_ptr = g_buf_cache.a.ensure(a_bytes); - b_ptr = g_buf_cache.b.ensure(b_bytes); - c_ptr = g_buf_cache.c.ensure(c_bytes); - scale_a_ptr = g_buf_cache.scale_a.ensure(scale_a_bytes); - scale_b_ptr = g_buf_cache.scale_b.ensure(scale_b_bytes); - - // Upload exactly the shape's byte count (cached buffers may be larger). - upload(a_ptr, a_m_k.data(), a_bytes); - upload(b_ptr, b_k_n.data(), b_bytes); - upload(scale_a_ptr, scale_a_shuffled.data(), scale_a_bytes); - upload(scale_b_ptr, scale_b_shuffled.data(), scale_b_bytes); - if(hipMemset(c_ptr, 0, c_bytes) != hipSuccess) - throw std::runtime_error("hipMemset C failed"); - } - else - { - // ---- Call-local device buffers, sized via HostTensor (packing-correct). ---- - a_local = std::make_unique(a_bytes); - b_local = std::make_unique(b_bytes); - c_local = std::make_unique(c_bytes); - sa_local = std::make_unique(scale_a_bytes); - sb_local = std::make_unique(scale_b_bytes); - - a_local->ToDevice(a_m_k.data()); - b_local->ToDevice(b_k_n.data()); - c_local->SetZero(); - sa_local->ToDevice(scale_a_shuffled.data()); - sb_local->ToDevice(scale_b_shuffled.data()); - - a_ptr = a_local->GetDeviceBuffer(); - b_ptr = b_local->GetDeviceBuffer(); - c_ptr = c_local->GetDeviceBuffer(); - scale_a_ptr = sa_local->GetDeviceBuffer(); - scale_b_ptr = sb_local->GetDeviceBuffer(); - } + // Call-local device buffers, sized via HostTensor (packing-correct). + // Each call allocates and (via RAII) frees its own DeviceMem -- the same + // per-call lifetime model the batched/multi_abd bridges use. A/B/scales + // are uploaded and C is zeroed every call, so per-shape data is fresh. + ck_tile::DeviceMem a_dev(a_bytes); + ck_tile::DeviceMem b_dev(b_bytes); + ck_tile::DeviceMem c_dev(c_bytes); + ck_tile::DeviceMem sa_dev(scale_a_bytes); + ck_tile::DeviceMem sb_dev(scale_b_bytes); + + a_dev.ToDevice(a_m_k.data()); + b_dev.ToDevice(b_k_n.data()); + c_dev.SetZero(); + sa_dev.ToDevice(scale_a_shuffled.data()); + sb_dev.ToDevice(scale_b_shuffled.data()); + + void* a_ptr = a_dev.GetDeviceBuffer(); + void* b_ptr = b_dev.GetDeviceBuffer(); + void* c_ptr = c_dev.GetDeviceBuffer(); + void* scale_a_ptr = sa_dev.GetDeviceBuffer(); + void* scale_b_ptr = sb_dev.GetDeviceBuffer(); // ---- Build MxGemmHostArgs in the exact profiler argument order: // a_ptr, scale_a, b_ptr, scale_b, ds{}, c_ptr, split_k, m, n, k, @@ -477,8 +377,7 @@ int dispatcher_run_mx_gemm(const void* A, return -2; } - // D2H from the resolved c_ptr (works for both the cached and per-call - // paths); copy exactly c_bytes since a cached buffer may be larger. + // D2H from c_ptr into the host output tensor. if(hipMemcpy(c_m_n.data(), c_ptr, c_bytes, hipMemcpyDeviceToHost) != hipSuccess) throw std::runtime_error("hipMemcpy D2H failed"); } diff --git a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py index dce3a76b4869..2090acf2a154 100644 --- a/projects/composablekernel/dispatcher/python/mx_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/mx_gemm_utils.py @@ -66,9 +66,17 @@ def _get_arch() -> str: arch = "" if not arch: raise RuntimeError("Could not detect GPU architecture from rocminfo; refusing to default. Pass gfx_arch explicitly.") - _supported = ("gfx90a", "gfx942", "gfx950") + # mx_gemm is gfx950-ONLY: the C++ bridge (mx_gemm_ctypes_lib.cpp) has a + # static_assert(GFX_ARCH == "gfx950") because it uses the gfx950-only + # preShuffleScaleBuffer_gfx950 host helper. Validating a broader set here + # would let a gfx942/gfx90a caller past detection only to fail later at + # build/runtime with a less actionable error, so restrict it to gfx950. + _supported = ("gfx950",) if arch not in _supported: - raise ValueError(f"Unsupported GPU architecture {arch!r}; supported: {list(_supported)}") + raise ValueError( + f"mx_gemm is gfx950-only; detected {arch!r} is not supported " + f"(supported: {list(_supported)})" + ) return arch # MX GEMM scales every 32 K-elements with one e8m0 byte. @@ -750,10 +758,14 @@ def setup_multiple_mx_gemm_dispatchers( # compiled for `arch` -- an arch mismatch between the header and the binary. # Resetting _name_cache forces the name to be recomputed for the chosen arch. if gfx_arch is not None: - _supported = ("gfx90a", "gfx942", "gfx950") + # mx_gemm is gfx950-only (the C++ bridge static_asserts GFX_ARCH==gfx950); + # reject any other explicit arch here so a build that can never succeed + # fails early with a clear message instead of at compile/runtime. + _supported = ("gfx950",) if gfx_arch not in _supported: raise ValueError( - f"Unsupported GPU architecture {gfx_arch!r}; supported: {list(_supported)}" + f"mx_gemm is gfx950-only; requested {gfx_arch!r} is not supported " + f"(supported: {list(_supported)})" ) for c in configs: if c.gpu_target != gfx_arch: diff --git a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py index 69680b9ab648..4eb5da021e16 100644 --- a/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_mx_gemm_bridge.py @@ -69,6 +69,9 @@ def test_projection_roundtrip(self): warp_m=2, warp_n=2, warp_k=1, warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, k_block_per_cu=3, + # Pin the arch so to_codegen_config() does not shell out to rocminfo; + # keeps this a CPU-only test on non-ROCm runners. + gpu_target="gfx950", ) j = cfg.to_codegen_config() self.assertEqual(j["datatype"], "fp8") @@ -100,6 +103,8 @@ def test_bf8_rejected(self): # codegen-level gate (_validate) refuse dtype="bf8". cfg = default_fp8_config() cfg.datatype = "bf8" + # Pin the arch so to_codegen_config() below stays CPU-only (no rocminfo). + cfg.gpu_target = "gfx950" self.assertFalse(cfg.is_valid()) try: @@ -248,6 +253,8 @@ def test_codegen_name_matches_config(self): except Exception as exc: # noqa: BLE001 self.skipTest(f"codegen import unavailable: {exc}") cfg = default_fp8_config() + # Pin the arch so to_codegen_config() stays CPU-only (no rocminfo). + cfg.gpu_target = "gfx950" try: name = kernel_name(cfg.to_codegen_config()) except Exception as exc: # noqa: BLE001 From 1a6765ea146ea545c5642c28e7e44f4b268494a0 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:41:27 -0400 Subject: [PATCH 14/14] fix(ck-tile): mx_gemm CMake resolves GFX_ARCH generically (no hardcoded gfx950) Per reviewer feedback ("no hardcoding to gfx950 arch"; "make it generic and throw if the arch is not gfx950"): the mx_gemm ctypes CMake no longer passes a literal GFX_ARCH="gfx950". It now resolves the configured GPU target from GPU_TARGETS (falling back to CK_TILE_GEMM_GPU_TARGET) and passes it through as -DGFX_ARCH="${MX_GEMM_GPU_TARGET}". The gfx950-only requirement stays enforced in the source (mx_gemm_ctypes_lib.cpp #errors if GFX_ARCH is unset and static_asserts GFX_ARCH == "gfx950", plus a runtime device guard) and in Python (_get_arch detects generically and raises if != gfx950). The target is only wired up when the build targets gfx950, so a gfx942/other dispatcher build is neither broken by the static_assert nor silently built for an unsupported arch. --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 92 ++++++++++++------- 1 file changed, 57 insertions(+), 35 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index a3743613275d..7c36e26b9cbf 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -87,42 +87,64 @@ set(MX_GEMM_TE_INCLUDE_DIRS ${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm/mx_gemm ) -file(GLOB MX_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/mx_gemm_*.hpp") -if(MX_GEMM_KERNEL_HEADERS) - list(GET MX_GEMM_KERNEL_HEADERS 0 MX_GEMM_KERNEL_HEADER) - message(STATUS "Found MX-GEMM kernel for ctypes lib: ${MX_GEMM_KERNEL_HEADER}") - - add_ctypes_library(dispatcher_mx_gemm_lib - mx_gemm_ctypes_lib.cpp - KERNEL_HEADER ${MX_GEMM_KERNEL_HEADER} - ) - target_include_directories(dispatcher_mx_gemm_lib PRIVATE ${MX_GEMM_TE_INCLUDE_DIRS}) - # The generated header only exports SelectedKernel/KERNEL_NAME/ScaleType/... - # under this define, so the force-include path requires it. GFX_ARCH is - # required by mx_gemm_ctypes_lib.cpp (#error if unset) and the bridge is - # gfx950-only (static_assert on preShuffleScaleBuffer_gfx950), so pin it to - # gfx950 here -- the only value this op can compile for. - target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE - CK_TILE_SINGLE_KERNEL_INCLUDE - GFX_ARCH="gfx950" - ) +# Resolve the configured GPU target GENERICALLY -- do NOT hardcode/default the +# arch to gfx950. mx_gemm is gfx950-only (it calls the gfx950-only +# preShuffleScaleBuffer_gfx950 helper), and that requirement is enforced IN THE +# SOURCE: mx_gemm_ctypes_lib.cpp #errors if GFX_ARCH is unset and static_asserts +# GFX_ARCH == "gfx950" (plus a runtime device guard). Here we simply pass the +# real build target through as -DGFX_ARCH and only wire up the target when the +# build actually targets gfx950, so a gfx942/other dispatcher build is neither +# broken by the static_assert nor silently built for an unsupported arch. +if(DEFINED GPU_TARGETS AND NOT GPU_TARGETS STREQUAL "") + string(REPLACE ";" " " _mx_gpu_targets_space "${GPU_TARGETS}") + string(REPLACE " " ";" _mx_gpu_targets_list "${_mx_gpu_targets_space}") + list(GET _mx_gpu_targets_list 0 MX_GEMM_GPU_TARGET) +elseif(DEFINED CK_TILE_GEMM_GPU_TARGET AND NOT CK_TILE_GEMM_GPU_TARGET STREQUAL "") + set(MX_GEMM_GPU_TARGET "${CK_TILE_GEMM_GPU_TARGET}") else() - message(STATUS "No MX-GEMM kernel found for ctypes lib - building without kernel") - add_library(dispatcher_mx_gemm_lib SHARED mx_gemm_ctypes_lib.cpp) - target_include_directories(dispatcher_mx_gemm_lib PRIVATE - ${PROJECT_SOURCE_DIR}/include - ${PROJECT_SOURCE_DIR}/dispatcher/include - ${MX_GEMM_TE_INCLUDE_DIRS} - ) - target_link_libraries(dispatcher_mx_gemm_lib PRIVATE hip::device) - # Even the no-kernel fallback compiles mx_gemm_ctypes_lib.cpp, which #errors - # if GFX_ARCH is unset; define it (gfx950, the only arch this op supports) so - # the fallback target builds instead of failing the whole configure/build. - target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE GFX_ARCH="gfx950") - set_target_properties(dispatcher_mx_gemm_lib PROPERTIES - POSITION_INDEPENDENT_CODE ON - CXX_STANDARD 17 - ) + set(MX_GEMM_GPU_TARGET "") +endif() + +if(NOT MX_GEMM_GPU_TARGET STREQUAL "gfx950") + message(STATUS + "MX-GEMM ctypes lib is gfx950-only; skipping for GPU target " + "'${MX_GEMM_GPU_TARGET}' (configure with GPU_TARGETS=gfx950 to enable).") +else() + file(GLOB MX_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/mx_gemm_*.hpp") + if(MX_GEMM_KERNEL_HEADERS) + list(GET MX_GEMM_KERNEL_HEADERS 0 MX_GEMM_KERNEL_HEADER) + message(STATUS "Found MX-GEMM kernel for ctypes lib: ${MX_GEMM_KERNEL_HEADER}") + + add_ctypes_library(dispatcher_mx_gemm_lib + mx_gemm_ctypes_lib.cpp + KERNEL_HEADER ${MX_GEMM_KERNEL_HEADER} + ) + target_include_directories(dispatcher_mx_gemm_lib PRIVATE ${MX_GEMM_TE_INCLUDE_DIRS}) + # The generated header only exports SelectedKernel/KERNEL_NAME/ScaleType/... + # under this define. Pass the resolved build target through as -DGFX_ARCH + # (the source #errors if unset and static_asserts it is gfx950). + target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="${MX_GEMM_GPU_TARGET}" + ) + else() + message(STATUS "No MX-GEMM kernel found for ctypes lib - building without kernel") + add_library(dispatcher_mx_gemm_lib SHARED mx_gemm_ctypes_lib.cpp) + target_include_directories(dispatcher_mx_gemm_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ${MX_GEMM_TE_INCLUDE_DIRS} + ) + target_link_libraries(dispatcher_mx_gemm_lib PRIVATE hip::device) + # Pass the resolved build target through as -DGFX_ARCH (the source #errors + # if unset and static_asserts it is gfx950). + target_compile_definitions(dispatcher_mx_gemm_lib PRIVATE + GFX_ARCH="${MX_GEMM_GPU_TARGET}") + set_target_properties(dispatcher_mx_gemm_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 + ) + endif() endif() # =============================================================================