ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 09:40:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#8998 (commit 5501ef1)

Browse Source 
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.
This commit is contained in:
Muhammed Emin Ozturk
2026-07-09 23:58:00 +00:00
committed by assistant-librarian[bot]
parent b0f200713a
commit 1c455b1bf5
6 changed files with 341 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
dispatcher/codegen/codegen_common.py
View File

@@ -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"
# ============================================================================

17
dispatcher/codegen/unified_gemm_codegen.py
View File

@@ -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<NumDTensor>;
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<EpilogueProblem>;"""
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]]};

17
dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp
View File

@@ -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_;

147
dispatcher/python/gemm_utils.py
View File

@@ -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_<dtype>_<layout>_...``."""
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],

108
dispatcher/tests/test_gemm_parity.py
View File

@@ -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} "

78
dispatcher/tests/test_gemm_utils.py
View File

@@ -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