From 1c455b1bf5dc3954d4c90be37ad13f339539fd80 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Thu, 9 Jul 2026 23:58:00 +0000 Subject: [PATCH] [rocm-libraries] ROCm/rocm-libraries#8998 (commit 5501ef1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts) (#8998) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Extends the Tile Engine ↔ Dispatcher GEMM **bridge** to the remaining data types TE's plain GEMM has MFMA warp tiles for, beyond the fp16/bf16 surface of #8479: - **fp8** (E4M3) and **bf8** (E5M2) → fp16 output, fp32 accumulate - **int8** → int32 output and accumulate (gfx942) All four A/B layout combinations per dtype (row-major C only, matching #8479). `fp32`/`fp64` are intentionally **excluded** — they appear in TE's dtype-string map but have no MFMA warp tiles in `GEMM_WARP_TILE_SUPPORTED_COMBINATIONS`, so no kernel can be generated/run. **Depends on the fp16/bf16 bridge in #8997** (`users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16`), which carries the bridge infrastructure and is not yet merged. This PR targets `develop`, so until #8997 merges its diff also includes the base bridge changes; please merge #8997 first. ## Changes - **Codegen** (`codegen_common.py`, `unified_gemm_codegen.py`): add `int32` to the dtype maps; `get_output_dtype` int8→int32; new `get_acc_dtype` (int8→int32, else fp32); derive `AccDataType`/`CDataType`, the `GEMM_KEY_DTYPE_{C,ACC}` macros, and the registry `dtype_c`/`dtype_acc` from the dtype instead of hard-coding `float`/`fp32`. - **Host harness** (`gemm_utils.py`): fp8/bf8 **FNUZ** (gfx942) uint8 codecs — exact decode (matches device `fp8_t`/`bf8_t`), nearest-representable saturating encode (same pattern as the existing bf16 helper); `GpuGemmRunner.run` encodes A/B and sizes the C buffer per dtype; `expand_sweep` sets `dtype_c`/`dtype_acc`. - **Tests**: `test_gemm_utils.py` adds CPU-only fp8/bf8 codec + output-dtype tests (all green); `test_gemm_parity.py` adds fp8/bf8/int8 cases with dtype-aware inputs/references/tolerances (int8 is bit-exact), GPU-gated like the existing cases. ## Verification done - `test_gemm_utils.py` + `test_codegen_common.py`: **54 passed** (CPU). - Codegen smoke: fp8/int8/fp16 each generate 1 kernel + 1 wrapper, 0 failed; emitted `ADataType/CDataType/AccDataType` and `GEMM_KEY_*` macros are correct (int8→int32_t acc/C; fp8→fp16_t C). - `test_gemm_parity.py` collects 60 cases and skips cleanly without a GPU. - The 16 unrelated failures in `test_examples_integration` / `test_grouped_conv_codegen` / `test_library_caching` are **pre-existing** (verified identical on the base branch; they require a built dispatcher `.a` / GPU). ## Test plan - [x] Merge #8997 (fp16/bf16 bridge), then this reduces to just the fp8/bf8/int8 delta on `develop`. - [x] On an MI300X (gfx942) node: run `python3 tests/test_gemm_parity.py` and confirm fp8/bf8/int8 parity; tune the fp8/bf8 tolerances if needed (current values are first-cut headroom). - [x] FNUZ vs OCP: the fp8/bf8 host codec targets the gfx942 FNUZ format; validate / extend for gfx950 (OCP) before enabling there. --- dispatcher/codegen/codegen_common.py | 25 ++- dispatcher/codegen/unified_gemm_codegen.py | 17 +- .../backends/generated_tile_backend.hpp | 17 +- dispatcher/python/gemm_utils.py | 147 +++++++++++++++++- dispatcher/tests/test_gemm_parity.py | 108 +++++++++---- dispatcher/tests/test_gemm_utils.py | 78 +++++++++- 6 files changed, 341 insertions(+), 51 deletions(-) diff --git a/dispatcher/codegen/codegen_common.py b/dispatcher/codegen/codegen_common.py index a0486da66d..a5f022021f 100644 --- a/dispatcher/codegen/codegen_common.py +++ b/dispatcher/codegen/codegen_common.py @@ -118,6 +118,7 @@ class CommonTypeMappings: "fp8": "fp8_t", "bf8": "bf8_t", "int8": "int8_t", + "int32": "int32_t", } DTYPE_TO_CK_QUALIFIED = { @@ -127,6 +128,7 @@ class CommonTypeMappings: "fp8": "ck_tile::fp8_t", "bf8": "ck_tile::bf8_t", "int8": "int8_t", + "int32": "int32_t", } DTYPE_TO_DISPATCHER = { @@ -136,6 +138,7 @@ class CommonTypeMappings: "fp8": "DataType::FP8", "bf8": "DataType::BF8", "int8": "DataType::INT8", + "int32": "DataType::INT32", } # GEMM-specific layout mappings ("r"/"c" for row/column major). @@ -202,8 +205,26 @@ class CommonTypeMappings: @staticmethod def get_output_dtype(dtype: str) -> str: - """Get output datatype (fp8/bf8 -> fp16).""" - return "fp16" if dtype in ("fp8", "bf8") else dtype + """Get output (C) datatype for an A/B element dtype. + + Low-precision float inputs accumulate into and store as fp16 + (fp8/bf8 -> fp16); int8 stores its int32 accumulator (int8 -> int32). + Everything else stores in its own dtype. + """ + if dtype in ("fp8", "bf8"): + return "fp16" + if dtype == "int8": + return "int32" + return dtype + + @staticmethod + def get_acc_dtype(dtype: str) -> str: + """Get accumulator datatype for an A/B element dtype. + + Integer GEMM accumulates in int32; every float dtype accumulates in + fp32. + """ + return "int32" if dtype == "int8" else "fp32" # ============================================================================ diff --git a/dispatcher/codegen/unified_gemm_codegen.py b/dispatcher/codegen/unified_gemm_codegen.py index f9233d325d..47cb3680f1 100755 --- a/dispatcher/codegen/unified_gemm_codegen.py +++ b/dispatcher/codegen/unified_gemm_codegen.py @@ -414,12 +414,13 @@ using namespace ck_tile; def _kernel_local_types(self, config: KernelConfig) -> str: """Generate data type and layout definitions inside kernel namespace""" output_dtype = self.tm.get_output_dtype(self.datatype) + acc_dtype = self.tm.get_acc_dtype(self.datatype) return f""" // Data types (inside namespace to avoid conflicts across layouts) using ADataType = {self.tm.DTYPE_TO_CK[self.datatype]}; using BDataType = {self.tm.DTYPE_TO_CK[self.datatype]}; - using AccDataType = float; + using AccDataType = {self.tm.DTYPE_TO_CK[acc_dtype]}; using CDataType = {self.tm.DTYPE_TO_CK[output_dtype]}; // Layouts (inside namespace to avoid conflicts when mixing layouts) @@ -452,6 +453,7 @@ using GemmMultiDArgs = GemmMultiDHostArgs; t = config.tile tr = config.trait output_dtype = self.tm.get_output_dtype(self.datatype) + acc_dtype = self.tm.get_acc_dtype(self.datatype) # Generate unique struct name and namespace from kernel name struct_name = f"Kernel_{kernel_name}" @@ -467,7 +469,7 @@ constexpr const char* KERNEL_NAME = "{kernel_name}"; // Data types (inside namespace to avoid conflicts across different kernels) using ADataType = {self.tm.DTYPE_TO_CK[self.datatype]}; using BDataType = {self.tm.DTYPE_TO_CK[self.datatype]}; -using AccDataType = float; +using AccDataType = {self.tm.DTYPE_TO_CK[acc_dtype]}; using CDataType = {self.tm.DTYPE_TO_CK[output_dtype]}; // Layouts (inside namespace to avoid conflicts when mixing layouts like RCR + RRR) @@ -522,8 +524,8 @@ using SelectedKernel = {ns_name}::{struct_name}; constexpr const char* KERNEL_NAME = {ns_name}::KERNEL_NAME; using ADataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]}; using BDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]}; -using CDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.tm.get_output_dtype(self.datatype)]}; -using AccDataType = float; +using CDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[output_dtype]}; +using AccDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[acc_dtype]}; // KernelKey field descriptors for the force-included kernel. // The ctypes library builds the registry KernelKey from these so the @@ -534,7 +536,7 @@ using AccDataType = float; #define GEMM_KEY_DTYPE_A "{self.datatype}" #define GEMM_KEY_DTYPE_B "{self.datatype}" #define GEMM_KEY_DTYPE_C "{output_dtype}" -#define GEMM_KEY_DTYPE_ACC "fp32" +#define GEMM_KEY_DTYPE_ACC "{acc_dtype}" #define GEMM_KEY_LAYOUT_A "{self.layout[0]}" #define GEMM_KEY_LAYOUT_B "{self.layout[1]}" #define GEMM_KEY_LAYOUT_C "{self.layout[2]}" @@ -784,7 +786,7 @@ using AccDataType = float; tuple<>, CLayout, element_wise::PassThrough, TilePartitioner::MPerBlock, TilePartitioner::NPerBlock, WarpPerBlock_M, WarpPerBlock_N, WarpTileM, WarpTileN, WarpTileK, - TransposeC, NumWaveGroups, false, 1, 1, DoubleSmemBuffer>; + TransposeC, NumWaveGroups>; using GemmEpilogue = CShuffleEpilogue;""" else: return """ @@ -815,6 +817,7 @@ class DispatcherWrapperGenerator: """Generate dispatcher wrapper""" kernel_name = KernelNaming.generate(config, self.datatype, self.layout) output_dtype = self.tm.get_output_dtype(self.datatype) + acc_dtype = self.tm.get_acc_dtype(self.datatype) rel_path = kernel_path.relative_to(output_dir) return f"""// SPDX-License-Identifier: MIT @@ -849,7 +852,7 @@ inline KernelInstancePtr make_{kernel_name}(const std::string& gfx_arch = "gfx94 key.signature.dtype_a = {self.tm.DTYPE_TO_DISPATCHER[self.datatype]}; key.signature.dtype_b = {self.tm.DTYPE_TO_DISPATCHER[self.datatype]}; key.signature.dtype_c = {self.tm.DTYPE_TO_DISPATCHER[output_dtype]}; - key.signature.dtype_acc = DataType::FP32; + key.signature.dtype_acc = {self.tm.DTYPE_TO_DISPATCHER[acc_dtype]}; key.signature.layout_a = {self.tm.LAYOUT_TO_DISPATCHER[self.layout[0]]}; key.signature.layout_b = {self.tm.LAYOUT_TO_DISPATCHER[self.layout[1]]}; key.signature.layout_c = {self.tm.LAYOUT_TO_DISPATCHER[self.layout[2]]}; diff --git a/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp b/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp index 4bb1496b9a..d529be3b31 100644 --- a/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp +++ b/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp @@ -108,7 +108,16 @@ class GeneratedTileKernelInstance : public KernelInstance { (void)d_ptrs; // Not used in basic GEMM - // Create arguments using constructor (correct order!) + // Leading dimensions depend on each operand's layout, NOT a fixed + // rcr assumption: RowMajor A/B/C -> inner axis is K/N/N; ColMajor -> + // M/K/M. Hard-coding {K, K, N} only happens to be right for rcr and for + // square problems (M==N==K); it corrupts every non-square rrr/ccr/crr + // launch. Derive each stride from the kernel's real layout instead. + const auto is_row = [](LayoutTag l) { return l == LayoutTag::RowMajor; }; + const auto stride_a = is_row(key_.signature.layout_a) ? problem.K : problem.M; + const auto stride_b = is_row(key_.signature.layout_b) ? problem.N : problem.K; + const auto stride_c = is_row(key_.signature.layout_c) ? problem.N : problem.M; + // Order from GemmHostArgs constructor: a_ptr, b_ptr, e_ptr, k_batch, M, N, K, stride_A, // stride_B, stride_E ck_tile::GemmHostArgs args(a_ptr, // a_ptr @@ -118,9 +127,9 @@ class GeneratedTileKernelInstance : public KernelInstance problem.M, // M problem.N, // N problem.K, // K - problem.K, // stride_A (row-major A: stride = K) - problem.K, // stride_B (column-major B: stride = K) - problem.N // stride_E/C (row-major C: stride = N) + stride_a, // stride_A + stride_b, // stride_B + stride_c // stride_E/C ); const bool bench = this->benchmarking_; diff --git a/dispatcher/python/gemm_utils.py b/dispatcher/python/gemm_utils.py index afe73dcc6a..160c0f0694 100644 --- a/dispatcher/python/gemm_utils.py +++ b/dispatcher/python/gemm_utils.py @@ -398,6 +398,113 @@ def _bf16_u16_to_fp32(u16: np.ndarray) -> np.ndarray: return (u16.astype(np.uint32) << 16).view(np.float32) +# --------------------------------------------------------------------------- +# fp8 (E4M3) / bf8 (E5M2) -- FNUZ ("NANOO") encoding used by gfx942/MI300. +# +# numpy has no native 8-bit float, and the C ABI only cares about the 1-byte +# memory layout (sizeof(fp8_t) == sizeof(bf8_t) == 1). We carry the value as a +# uint8 bit pattern. As with bf16, the DECODE is the load-bearing half: it must +# return the exact value the device's fp8_t/bf8_t represents for a byte, so the +# NumPy reference multiplies bit-for-bit what the GPU multiplies. The ENCODE only +# needs to land on the nearest representable byte. +# +# FNUZ format (gfx942): bias = 2^(exp_bits-1); the all-1s exponent is a normal +# number (no Inf), the sole NaN is the sign=1/exp=0/mant=0 byte (0x80), and there +# is no negative zero. gfx950/MI350 uses the OCP fp8 format instead; this codec +# targets the gfx942 default and the OCP path needs separate handling. +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=None) +def _fnuz_decode_table(exp_bits: int, mant_bits: int) -> np.ndarray: + """Build the 256-entry byte -> fp32 value table for an 8-bit FNUZ float. + + The table is a pure function of (exp_bits, mant_bits), so it is cached; the + returned array is marked read-only because callers share the one instance. + """ + bias = (1 << (exp_bits - 1)) + mant_max = 1 << mant_bits + sign_shift = exp_bits + mant_bits + exp_mask = (1 << exp_bits) - 1 + table = np.zeros(256, dtype=np.float32) + for b in range(256): + sign = (b >> sign_shift) & 1 + exp = (b >> mant_bits) & exp_mask + mant = b & (mant_max - 1) + if exp == 0 and mant == 0: + # +0 (0x00); the negative-zero slot (0x80) is the lone NaN. + table[b] = np.float32(np.nan) if sign else np.float32(0.0) + continue + if exp == 0: + val = (mant / mant_max) * (2.0 ** (1 - bias)) # subnormal + else: + val = (1.0 + mant / mant_max) * (2.0 ** (exp - bias)) # normal + table[b] = np.float32(-val if sign else val) + table.flags.writeable = False # shared cached instance -- do not mutate + return table + + +def _fnuz_encode(x: np.ndarray, exp_bits: int, mant_bits: int) -> np.ndarray: + """Encode fp32 -> nearest 8-bit FNUZ float, returned as a uint8 bit pattern.""" + table = _fnuz_decode_table(exp_bits, mant_bits) + sign_byte = np.uint8(1 << (exp_bits + mant_bits)) # 0x80 + + # Positive half (bytes 0..127) holds every non-negative magnitude, sorted. + # Compare in float64: for very large inputs the gap between the two top + # magnitudes is below fp32 resolution, which would tie and mis-saturate. + pos_mag = table[: int(sign_byte)].astype(np.float64) + order = np.argsort(pos_mag) + sorted_mag = pos_mag[order] + sorted_byte = order.astype(np.uint8) + + xf = np.ascontiguousarray(x, dtype=np.float32) + ax = np.abs(xf).astype(np.float64) + # Both neighbours come from the raw insertion point: raw==size saturates to + # the top magnitude (lo==hi), raw==0 pins to zero, otherwise compare the two. + raw = np.searchsorted(sorted_mag, ax) + hi = np.clip(raw, 0, sorted_mag.size - 1) + lo = np.clip(raw - 1, 0, sorted_mag.size - 1) + pick_lo = np.abs(sorted_mag[lo] - ax) <= np.abs(sorted_mag[hi] - ax) + chosen = np.where(pick_lo, lo, hi) + out = sorted_byte[chosen] + + # Apply sign, but never the 0x80 (-0 == NaN) slot: zeros stay +0. + is_zero = sorted_mag[chosen] == 0 + out = np.where((xf < 0) & ~is_zero, out | sign_byte, out) + out = np.where(np.isnan(xf), sign_byte, out) # NaN inputs -> NaN byte + return out.astype(np.uint8).reshape(np.shape(x)) + + +def _fp32_to_fp8_u8(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> fp8 E4M3 (FNUZ) bit pattern in a uint8 array.""" + return _fnuz_encode(x, exp_bits=4, mant_bits=3) + + +def _fp8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: + """Decode an fp8 E4M3 (FNUZ) bit pattern back to fp32.""" + return _fnuz_decode_table(4, 3)[u8.astype(np.intp)] + + +def _fp32_to_bf8_u8(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> bf8 E5M2 (FNUZ) bit pattern in a uint8 array.""" + return _fnuz_encode(x, exp_bits=5, mant_bits=2) + + +def _bf8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: + """Decode a bf8 E5M2 (FNUZ) bit pattern back to fp32.""" + return _fnuz_decode_table(5, 2)[u8.astype(np.intp)] + + +# Output (C) element dtype for an A/B element dtype, mirroring the codegen's +# CommonTypeMappings.get_output_dtype: fp8/bf8 accumulate into fp16, int8 into +# int32, everything else stores in its own dtype. +_OUTPUT_DTYPE = {"fp8": "fp16", "bf8": "fp16", "int8": "int32"} + + +def _output_dtype(dtype: str) -> str: + return _OUTPUT_DTYPE.get(dtype, dtype) + + def _dtype_from_kernel_name(name: str) -> str: """Extract the dtype token from a kernel name like ``gemm___...``.""" parts = name.split("_") @@ -459,20 +566,47 @@ class GpuGemmRunner: B_lay = B if lb == "r" else B.T C_shape = (M, N) if lc == "r" else (N, M) + # Build A/B host buffers in the kernel's element dtype. The encode + # helpers (bf16/fp8/bf8) already force a contiguous float32 source, so an + # outer ascontiguousarray would only add a redundant copy; the native + # numpy dtypes (fp16/int8) still need it. if dtype == "bf16": - # _fp32_to_bf16_u16 already forces a contiguous float32 buffer, so - # an outer ascontiguousarray here would only add a redundant copy. A_h = _fp32_to_bf16_u16(A_lay) B_h = _fp32_to_bf16_u16(B_lay) - C_h = np.zeros(C_shape, dtype=np.uint16) + elif dtype == "fp8": + A_h = _fp32_to_fp8_u8(A_lay) + B_h = _fp32_to_fp8_u8(B_lay) + elif dtype == "bf8": + A_h = _fp32_to_bf8_u8(A_lay) + B_h = _fp32_to_bf8_u8(B_lay) + elif dtype == "int8": + A_h = np.ascontiguousarray(A_lay, dtype=np.int8) + B_h = np.ascontiguousarray(B_lay, dtype=np.int8) else: # fp16 (default) A_h = np.ascontiguousarray(A_lay, dtype=np.float16) B_h = np.ascontiguousarray(B_lay, dtype=np.float16) - C_h = np.zeros(C_shape, dtype=np.float16) + + # The C buffer's element size must equal sizeof(CDataType): fp8/bf8 + # accumulate into fp16, int8 into int32, otherwise the input dtype. + out_dtype = _output_dtype(dtype) + _C_NP = {"fp16": np.float16, "bf16": np.uint16, "int32": np.int32} + if out_dtype not in _C_NP: + # A silent fp16 fallback would size the host C buffer wrong for an + # unrecognized dtype (sizeof(CDataType) mismatch -> corrupt results + # across the C ABI). Fail loudly so a new dtype is added here. + raise ValueError( + f"unsupported C dtype {out_dtype!r} (from input dtype {dtype!r}); " + "add it to _C_NP so the host buffer matches sizeof(CDataType)" + ) + C_h = np.zeros(C_shape, dtype=_C_NP[out_dtype]) status, time_ms = self.lib.run(A_h, B_h, C_h, M, N, K) - C_dec = _bf16_u16_to_fp32(C_h) if dtype == "bf16" else C_h + # Decode the output back to a comparable numeric array. + if out_dtype == "bf16": + C_dec = _bf16_u16_to_fp32(C_h) + else: # fp16 / int32 are already directly comparable + C_dec = C_h C_out = C_dec if lc == "r" else C_dec.T tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 @@ -859,7 +993,8 @@ def expand_sweep( c = GemmKernelConfig( dtype_a=dtype, dtype_b=dtype, - dtype_c=dtype, + dtype_c=_output_dtype(dtype), + dtype_acc=("int32" if dtype == "int8" else "fp32"), layout_a=_LAYOUT_WORD[la], layout_b=_LAYOUT_WORD[lb], layout_c=_LAYOUT_WORD[lc], diff --git a/dispatcher/tests/test_gemm_parity.py b/dispatcher/tests/test_gemm_parity.py index 308c570067..b9d1cd1cd9 100644 --- a/dispatcher/tests/test_gemm_parity.py +++ b/dispatcher/tests/test_gemm_parity.py @@ -44,21 +44,27 @@ from gemm_utils import ( # noqa: E402 setup_multiple_gemm_dispatchers, _fp32_to_bf16_u16, _bf16_u16_to_fp32, + _fp32_to_fp8_u8, + _fp8_u8_to_fp32, + _fp32_to_bf8_u8, + _bf8_u8_to_fp32, + _output_dtype, ) from ctypes_utils import detect_gpu_arch, get_build_dir # noqa: E402 # (dtype, layout) surface the regular bridge supports. Column-major C is rejected # by ck_tile's universal GEMM at build, so every layout keeps row-major C, which -# leaves exactly the four A/B combinations below. Both dtypes cover all four. +# leaves exactly the four A/B combinations below. Every dtype covers all four. +# +# fp16/bf16 are the PR #8479 surface; fp8 (E4M3), bf8 (E5M2) and int8 are the +# remaining dtypes TE's plain GEMM has MFMA warp tiles for (fp8/bf8 -> fp16 out, +# int8 -> int32 out). int8 only has warp tiles on gfx942; on other arches its +# kernels simply fail to build and the case skips (handled below). +_FLOAT_DTYPES = ("fp16", "bf16", "fp8", "bf8") +_INT_DTYPES = ("int8",) +_LAYOUTS = ("rcr", "rrr", "ccr", "crr") _CASES = [ - ("fp16", "rcr"), - ("fp16", "rrr"), - ("fp16", "ccr"), - ("fp16", "crr"), - ("bf16", "rcr"), - ("bf16", "rrr"), - ("bf16", "ccr"), - ("bf16", "crr"), + (dt, lay) for dt in (*_FLOAT_DTYPES, *_INT_DTYPES) for lay in _LAYOUTS ] # Padded default algorithm: pad_* all True so M/N need not divide the tile, which @@ -81,25 +87,72 @@ _SHAPES = [ ("awkward", 257, 129, 512), ] -# Global-relative-error gates. fp16 measured ~3-4e-4 and bf16 ~8e-3 on gfx942; -# these leave headroom without masking a real regression. -_TOL = {"fp16": 2e-3, "bf16": 1.5e-2} +# Global-relative-error gates. fp16 measured ~3-4e-4 and bf16 ~8e-3 on gfx942. +# fp8/bf8 are far coarser (3- and 2-bit mantissa) so their gates are looser; int8 +# is an exact integer accumulation so it must match bit-for-bit. The fp8/bf8 +# gates are first-cut headroom values and may want tightening once measured on a +# GPU. +_TOL = { + "fp16": 2e-3, + "bf16": 1.5e-2, + "fp8": 1.5e-1, + "bf8": 3.0e-1, + "int8": 0.0, +} _LAYOUT_WORD = {"r": "row", "c": "col"} -def _emulate(x: np.ndarray, dtype: str) -> np.ndarray: - """Round fp32 to the kernel's storage dtype so the CPU reference matches what - the GPU actually multiplies (and stores).""" +def _emulate_input(x: np.ndarray, dtype: str) -> np.ndarray: + """Round an fp32 operand to the kernel's storage dtype so the CPU reference + multiplies exactly what the GPU does. int8 inputs are already integral.""" if dtype == "bf16": return _bf16_u16_to_fp32(_fp32_to_bf16_u16(x)) + if dtype == "fp8": + return _fp8_u8_to_fp32(_fp32_to_fp8_u8(x)) + if dtype == "bf8": + return _bf8_u8_to_fp32(_fp32_to_bf8_u8(x)) + if dtype == "int8": + return x.astype(np.float64) # exact; widened to avoid product overflow return x.astype(np.float16).astype(np.float32) +def _emulate_output(c: np.ndarray, out_dtype: str) -> np.ndarray: + """Round the fp32 accumulator to the kernel's C storage dtype.""" + if out_dtype == "bf16": + return _bf16_u16_to_fp32(_fp32_to_bf16_u16(c)) + if out_dtype == "int32": + return c # integer accumulation is exact + return c.astype(np.float16).astype(np.float32) # fp16 + + +def _make_inputs(dtype, M, N, K, rng): + """Random A (MxK), B (KxN) for a dtype: floats for the float dtypes, small + integers for int8 (kept small so the int32 accumulation cannot overflow).""" + if dtype == "int8": + A = rng.integers(-4, 5, size=(M, K)).astype(np.float32) + B = rng.integers(-4, 5, size=(K, N)).astype(np.float32) + return A, B + A = (rng.standard_normal((M, K)) * 0.1).astype(np.float32) + B = (rng.standard_normal((K, N)) * 0.1).astype(np.float32) + return A, B + + +def _reference(A, B, dtype): + """NumPy reference matching the kernel: round inputs to the storage dtype, + accumulate (fp32 for floats / exact int for int8), then round to C dtype.""" + out_dtype = _output_dtype(dtype) + acc = _emulate_input(A, dtype) @ _emulate_input(B, dtype) + ref = _emulate_output(acc, out_dtype) + return ref.astype(np.int32) if out_dtype == "int32" else ref + + def _config(dtype: str, layout: str, arch: str) -> GemmKernelConfig: la, lb, lc = layout return GemmKernelConfig( - dtype_a=dtype, dtype_b=dtype, dtype_c=dtype, + dtype_a=dtype, dtype_b=dtype, + dtype_c=_output_dtype(dtype), + dtype_acc=("int32" if dtype == "int8" else "fp32"), layout_a=_LAYOUT_WORD[la], layout_b=_LAYOUT_WORD[lb], layout_c=_LAYOUT_WORD[lc], gfx_arch=arch, **_ALGO, ) @@ -161,17 +214,13 @@ class GemmBridgeParity(unittest.TestCase): _, M, N, K = shape problem = GemmProblem(M=M, N=N, K=K) rng = np.random.default_rng(42) - A = (rng.standard_normal((M, K)) * 0.1).astype(np.float32) - B = (rng.standard_normal((K, N)) * 0.1).astype(np.float32) + A, B = _make_inputs(dtype, M, N, K, rng) runner = GpuGemmRunner(lib_path=so) # The .so is the contract endpoint: the name it reports must be the config - # name that drove codegen + the force-include build. - self.assertEqual(runner.kernel_name, GemmKernelConfig( - dtype_a=dtype, dtype_b=dtype, dtype_c=dtype, - layout_a=_LAYOUT_WORD[layout[0]], layout_b=_LAYOUT_WORD[layout[1]], - layout_c=_LAYOUT_WORD[layout[2]], gfx_arch=self.arch, **_ALGO, - ).name) + # name that drove codegen + the force-include build. The kernel name keys + # off the input dtype (dtype_a), not the C/acc dtype. + self.assertEqual(runner.kernel_name, _config(dtype, layout, self.arch).name) result = runner.run(A, B, problem) self.assertTrue( @@ -179,8 +228,8 @@ class GemmBridgeParity(unittest.TestCase): f"{dtype}/{layout} {shape[0]} run failed (status {result.status})", ) - ref = _emulate(_emulate(A, dtype) @ _emulate(B, dtype), dtype) - max_rel = _max_rel(result.output, ref) + ref = _reference(A, B, dtype) + max_rel = _max_rel(result.output.astype(np.float64), ref.astype(np.float64)) self.assertLessEqual( max_rel, _TOL[dtype], f"{dtype}/{layout} {shape[0]} max_rel={max_rel:.2e} > {_TOL[dtype]:.0e}", @@ -237,14 +286,13 @@ def _main() -> int: for sname, M, N, K in _SHAPES: total += 1 problem = GemmProblem(M=M, N=N, K=K) - A = (rng.standard_normal((M, K)) * 0.1).astype(np.float32) - B = (rng.standard_normal((K, N)) * 0.1).astype(np.float32) + A, B = _make_inputs(dtype, M, N, K, rng) result = runner.run(A, B, problem) if not result.success: print(f" {tag:<12} {sname:<12} {'RUN FAILED':>9} status={result.status}") continue - ref = _emulate(_emulate(A, dtype) @ _emulate(B, dtype), dtype) - mr = _max_rel(result.output, ref) + ref = _reference(A, B, dtype) + mr = _max_rel(result.output.astype(np.float64), ref.astype(np.float64)) ok = mr <= _TOL[dtype] passed += ok print(f" {tag:<12} {sname:<12} {result.tflops:>9.1f} " diff --git a/dispatcher/tests/test_gemm_utils.py b/dispatcher/tests/test_gemm_utils.py index 34e07ecfcb..8ed188932e 100644 --- a/dispatcher/tests/test_gemm_utils.py +++ b/dispatcher/tests/test_gemm_utils.py @@ -8,6 +8,9 @@ Locks in the bit-level helpers that the TE -> Dispatcher GEMM bridge relies on: * bf16 <-> uint16 encoding (round-to-nearest-even), since numpy has no native bf16 and the runner carries bf16 as a uint16 bit pattern. + * fp8 (E4M3) / bf8 (E5M2) FNUZ <-> uint8 encoding, used for the gfx942 8-bit + float surface. The decode must be exact to the device format; the encode + only needs to land on the nearest representable byte. * dtype / layout parsing from the compiled kernel name, which drives how the runner lays out host buffers. @@ -29,6 +32,12 @@ from gemm_utils import ( # noqa: E402 GemmKernelConfig, _fp32_to_bf16_u16, _bf16_u16_to_fp32, + _fp32_to_fp8_u8, + _fp8_u8_to_fp32, + _fp32_to_bf8_u8, + _bf8_u8_to_fp32, + _fnuz_decode_table, + _output_dtype, _dtype_from_kernel_name, _layout_from_kernel_name, ) @@ -78,6 +87,70 @@ class TestBf16Encoding(unittest.TestCase): self.assertEqual(u16.itemsize, 2) # must match sizeof(bf16_t) on device +class TestFp8Bf8Encoding(unittest.TestCase): + """fp8 E4M3 / bf8 E5M2 in the FNUZ format used by gfx942. + + The decode is the load-bearing half (it must equal the device value for a + byte); the encode must land on the nearest representable byte and saturate. + """ + + def test_format_ranges(self): + # FNUZ maxima: E4M3 -> 2^7 * 1.875 = 240; E5M2 -> 2^15 * 1.75 = 57344. + t43 = _fnuz_decode_table(4, 3) + t52 = _fnuz_decode_table(5, 2) + self.assertEqual(float(np.nanmax(t43)), 240.0) + self.assertEqual(float(np.nanmin(t43)), -240.0) + self.assertEqual(float(np.nanmax(t52)), 57344.0) + self.assertEqual(float(np.nanmin(t52)), -57344.0) + + def test_zero_and_nan_slots(self): + # 0x00 is +0; the negative-zero slot 0x80 is the lone NaN (FNUZ). + for tab in (_fnuz_decode_table(4, 3), _fnuz_decode_table(5, 2)): + self.assertEqual(float(tab[0x00]), 0.0) + self.assertTrue(np.isnan(tab[0x80])) + + def test_exactly_representable_roundtrip(self): + exact = np.array([0.0, 0.5, 1.0, -1.0, 2.0, -2.0, 1.5, -0.25, 4.0, 8.0], + dtype=np.float32) + np.testing.assert_array_equal( + _fp8_u8_to_fp32(_fp32_to_fp8_u8(exact)), exact) + np.testing.assert_array_equal( + _bf8_u8_to_fp32(_fp32_to_bf8_u8(exact)), exact) + + def test_decode_is_consistent_with_encode(self): + # The parity contract: ref multiplies decode(encode(x)), so the pair must + # be self-consistent and every encoded byte must decode finite. + rng = np.random.default_rng(1) + x = (rng.standard_normal(5000) * 0.1).astype(np.float32) + for enc, dec in ((_fp32_to_fp8_u8, _fp8_u8_to_fp32), + (_fp32_to_bf8_u8, _bf8_u8_to_fp32)): + d = dec(enc(x)) + self.assertTrue(np.all(np.isfinite(d))) + + def test_saturates_no_inf(self): + # FNUZ has no infinity: huge magnitudes clamp to the finite max. + big = np.array([1e30, -1e30], dtype=np.float32) + self.assertEqual(float(_fp8_u8_to_fp32(_fp32_to_fp8_u8(big))[0]), 240.0) + self.assertEqual(float(_bf8_u8_to_fp32(_fp32_to_bf8_u8(big))[1]), -57344.0) + + def test_dtype_and_size(self): + for enc in (_fp32_to_fp8_u8, _fp32_to_bf8_u8): + u8 = enc(np.zeros(4, dtype=np.float32)) + self.assertEqual(u8.dtype, np.uint8) + self.assertEqual(u8.itemsize, 1) # must match sizeof(fp8_t/bf8_t) + + +class TestOutputDtype(unittest.TestCase): + """Output (C) element dtype must mirror the codegen's get_output_dtype.""" + + def test_mapping(self): + self.assertEqual(_output_dtype("fp16"), "fp16") + self.assertEqual(_output_dtype("bf16"), "bf16") + self.assertEqual(_output_dtype("fp8"), "fp16") + self.assertEqual(_output_dtype("bf8"), "fp16") + self.assertEqual(_output_dtype("int8"), "int32") + + class TestKernelNameParsing(unittest.TestCase): """The runner reads dtype + layout straight from the compiled .so name.""" @@ -114,13 +187,14 @@ class TestConfigNameContract(unittest.TestCase): codegen -> runtime; parsing it back must recover dtype and layout.""" def test_name_roundtrips_through_parsers(self): - for dtype in ("fp16", "bf16"): + for dtype in ("fp16", "bf16", "fp8", "bf8", "int8"): for la, lb, lc in (("row", "col", "row"), ("row", "row", "row"), ("col", "col", "row"), ("col", "row", "row")): cfg = GemmKernelConfig( - dtype_a=dtype, dtype_b=dtype, dtype_c=dtype, + dtype_a=dtype, dtype_b=dtype, dtype_c=_output_dtype(dtype), + dtype_acc=("int32" if dtype == "int8" else "fp32"), layout_a=la, layout_b=lb, layout_c=lc, ) name = cfg.name