From a88d8333d0f3738475b5857284637d8121186de8 Mon Sep 17 00:00:00 2001 From: Muhammed Ozturk Date: Mon, 8 Jun 2026 18:12:23 -0400 Subject: [PATCH 1/2] [CK_TILE] Add bf16 support to GEMM Tile-Engine -> Dispatcher bridge Closes the regular-path dtype-coverage gap (fp16-only) identified when measuring the GEMM bridge against the FMHA reference (PR #5260). The C ABI is already dtype-agnostic (sizes buffers from sizeof(ADataType)) and the codegen already emits bf16 kernels; the only fp16 hard-coding lived in the Python host buffers. - gemm_utils.py: GpuGemmRunner.run() is now dtype-aware. It detects the kernel's real dtype from the compiled kernel name and encodes the host buffers to match. bf16 has no native numpy dtype, so it is carried as a uint16 bit pattern (sizeof(bf16_t) == 2): _fp32_to_bf16_u16 encodes fp32 -> bf16 with round-to-nearest-even, _bf16_u16_to_fp32 decodes the output back to fp32. fp16 path is unchanged. - run_one_gemm_kernel.py: worker now generates fp32 source data so the runner owns all dtype encoding (no double rounding). Stacked on muozturk/dispatcher-gemm-registry-key (PR #8187): bf16 needs the trait-derived KernelKey so the registry reports the correct dtype. Validated on gfx942 (MI300X): 3 bf16 kernels x 2 problems = 6/6 OK, 46-219 TFLOPS; numeric check vs CPU reference max_rel_err 7.7e-3 (within bf16 tolerance). Co-Authored-By: Claude Opus 4.7 --- .../dispatcher/python/gemm_utils.py | 48 +++++++++++++++++-- .../ops/gemm/run_one_gemm_kernel.py | 6 ++- 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 6d5112048586..bd78bb868439 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -377,6 +377,30 @@ def cleanup(self) -> None: # ============================================================================ +def _fp32_to_bf16_u16(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> bfloat16 bit pattern in a uint16 array (round-to-nearest-even). + + numpy has no native bf16, but the C ABI only cares about the 2-byte memory + layout (sizeof(bf16_t) == 2 == sizeof(uint16)). Truncating the low 16 bits of + the fp32 representation with round-to-nearest-even matches ck_tile's bf16. + """ + u32 = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32) + # round-to-nearest-even: add (lsb-of-kept-bits + 0x7FFF) before truncating + rounding = ((u32 >> 16) & 1) + np.uint32(0x7FFF) + return ((u32 + rounding) >> 16).astype(np.uint16) + + +def _bf16_u16_to_fp32(u16: np.ndarray) -> np.ndarray: + """Decode a uint16 bf16 bit pattern back to fp32 (low 16 mantissa bits zero).""" + return (u16.astype(np.uint32) << 16).view(np.float32) + + +def _dtype_from_kernel_name(name: str) -> str: + """Extract the dtype token from a kernel name like ``gemm___...``.""" + parts = name.split("_") + return parts[1] if len(parts) > 1 else "fp16" + + class GpuGemmRunner: """High-level runner: construct from a .so path, call run(A, B, problem). @@ -402,16 +426,30 @@ def run( M, N, K = problem.M, problem.N, problem.K # A is row-major MxK; B is supplied KxN and stored column-major (the - # 'c' in rcr), matching how the kernel expects its operands. - A_h = np.ascontiguousarray(A, dtype=np.float16) - B_h = np.ascontiguousarray(B.T, dtype=np.float16) - C_h = np.zeros((M, N), dtype=np.float16) + # 'c' in rcr), matching how the kernel expects its operands. The element + # dtype is dictated by the compiled kernel (encoded in its name); the C + # ABI sizes its buffers from sizeof(ADataType), so the host buffers must + # match it byte-for-byte. + dtype = _dtype_from_kernel_name(self._kernel_name) + if dtype == "bf16": + A_h = _fp32_to_bf16_u16(A) + B_h = _fp32_to_bf16_u16(np.ascontiguousarray(B.T)) + C_h = np.zeros((M, N), dtype=np.uint16) + else: # fp16 (default) + A_h = np.ascontiguousarray(A, dtype=np.float16) + B_h = np.ascontiguousarray(B.T, dtype=np.float16) + C_h = np.zeros((M, N), dtype=np.float16) status, time_ms = self.lib.run(A_h, B_h, C_h, M, N, K) + if dtype == "bf16": + C_out = _bf16_u16_to_fp32(C_h) + else: + C_out = C_h + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 return GemmResult( - output=C_h, + output=C_out, time_ms=time_ms, status=status, tflops=tflops, diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py index 23fca859d4a4..e0c71d768b12 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py @@ -37,8 +37,10 @@ def _run_one(idx, so_path, prob_dict, kernel_name): problem = GemmProblem.from_dict(prob_dict) np.random.seed(42) - A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float16) - B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float16) + # Generate fp32 source; the runner encodes to the kernel's real dtype + # (fp16 or bf16) based on the compiled kernel name. + A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float32) + B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float32) # CRITICAL: load the library ONLY inside this subprocess. runner = GpuGemmRunner(lib_path=so_path) From b3f43db46730287f0917f4d816146aa8ac00bda6 Mon Sep 17 00:00:00 2001 From: Muhammed Ozturk Date: Tue, 9 Jun 2026 13:33:27 -0400 Subject: [PATCH 2/2] [CK_TILE] Widen GEMM benchmark --dtype choices to include bf16 This branch's runner encodes bf16 from the kernel name, so add bf16 to SUPPORTED_DTYPES. Layout stays rcr-only here. Co-Authored-By: Claude Opus 4.7 --- .../tile_engine/ops/gemm/gemm_full_benchmark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py index 4b90a99cc2ea..ed296e35d6f8 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -47,7 +47,7 @@ # column-major B (rcr). Constrain the CLI to what the runner actually supports so # a mismatched --dtype/--layout fails fast instead of silently benchmarking the # wrong thing. Later phases widen these sets as the runner gains signatures. -SUPPORTED_DTYPES = ("fp16",) +SUPPORTED_DTYPES = ("fp16", "bf16") SUPPORTED_LAYOUTS = ("rcr",)