[CK_TILE] Add bf16 support to GEMM TE -> Dispatcher bridge by ozturkosu · Pull Request #8190 · ROCm/rocm-libraries

ozturkosu · 2026-06-08T22:12:49Z

Summary

Closes the regular-path dtype-coverage gap (fp16/rcr only) identified when
measuring our GEMM Tile-Engine → Dispatcher bridge against the FMHA reference
(PR #5260). FMHA shipped fp16 and bf16; our regular path was fp16-only.
This PR adds bf16, bringing the regular path to dtype parity for the two
core floating-point types.

Draft, stacked on muozturk/dispatcher-gemm-registry-key (PR #8187).
bf16 depends on the trait-derived KernelKey from #8187: without it a bf16
kernel would be registered under the old hard-coded fp16 key and the
registry would misreport its dtype.

Why this is small

The bridge was already 90% dtype-agnostic:

The flat C ABI (gemm_ctypes_lib.cpp) sizes every device buffer from
sizeof(ADataType)/sizeof(BDataType)/sizeof(CDataType), which come from
the force-included kernel header — so it already handles bf16 byte layout.
unified_gemm_codegen.py already emits bf16 kernels (gemm_bf16_rcr_...).
expand_sweep already threads the --dtype argument through to the configs.

The only fp16 hard-coding was in the Python host buffers in the runner.

The fix

dispatcher/python/gemm_utils.py — GpuGemmRunner.run() is now
dtype-aware. It detects the kernel's real dtype from the compiled kernel
name (gemm_<dtype>_...) and encodes the host buffers to match. numpy has
no native bf16, so bf16 is carried as a uint16 bit pattern
(sizeof(bf16_t) == 2 == sizeof(uint16)):
- _fp32_to_bf16_u16 — fp32 → bf16 with round-to-nearest-even
- _bf16_u16_to_fp32 — decodes the kernel's output back to fp32
  The fp16 path is unchanged.
tile_engine/ops/gemm/run_one_gemm_kernel.py — the worker now generates
fp32 source data so the runner owns all dtype encoding (avoids
double-rounding for bf16).

Design decisions

Dependency-free bf16 via uint16 byte-encoding rather than adding
ml_dtypes/bfloat16 as a new dependency. The C ABI only cares about the
2-byte memory layout, so a correct bit pattern is sufficient and keeps the
bridge dependency-light.
Dtype detected from the kernel name, not passed as a separate flag, so
the runner cannot disagree with the kernel it actually loaded. Single source
of truth = the compiled .so's reported name.

Validation (gfx942 / MI300X)

3 bf16 kernels × 2 problems = 6/6 OK:

Kernel	M=N=K=512	M=N=K=1024
compv3_cshuffle_intrawave	46.2 TFLOPS	210.2 TFLOPS
compv3_cshuffle_intrawave (persistent)	43.4 TFLOPS	209.0 TFLOPS
compv3_default_intrawave	46.9 TFLOPS	219.3 TFLOPS

Numeric correctness vs CPU reference (bf16-emulated inputs, fp32 accumulate):
max_rel_err 7.7e-3, mean 2.8e-3 — within bf16 precision. Output buffers
fully populated (non_zero == M*N).

Remaining (tracked, not in this PR)

Layouts beyond rcr (rrr/rcc/…): the same dtype-agnostic mechanism applies;
needs a layout-aware host-buffer transpose, follow-up.

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 <noreply@anthropic.com>

Copilot

Pull request overview

Adds bfloat16 (bf16) support to the CK Tile GEMM Tile-Engine → Dispatcher bridge by making the Python runner dtype-aware and ensuring host buffers match the compiled kernel’s dtype.

Changes:

Generate fp32 source matrices in the GEMM worker so dtype encoding is handled centrally in the runner.
Add dependency-free bf16 host-buffer encoding/decoding in GpuGemmRunner.run() using uint16 bit patterns.
Infer dtype from the compiled kernel name (gemm_<dtype>_...) to prevent runner/kernel mismatches.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py	Switches worker input generation to fp32 so the runner performs dtype-specific encoding (fp16/bf16).
projects/composablekernel/dispatcher/python/gemm_utils.py	Adds bf16 bit-pattern encode/decode helpers and makes `GpuGemmRunner.run()` select host buffer dtypes based on the compiled kernel name.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+            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)


+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)


…spatcher-gemm-bf16

Addresses the Copilot review requests on the bf16 (#8190) and layout (#8191) changes: the bit-level bf16 round-to-nearest-even encoder and the dtype/layout name parsers had no coverage. Adds dispatcher/tests/test_gemm_utils.py (10 CPU-only tests): bf16 exact round-trip, <=2^-8 relative error, tie-to-even, Inf/NaN, uint16 size; dtype/layout parsing with fallbacks; and a GemmKernelConfig.name -> parser round-trip locking the config/codegen/runtime naming contract. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

ozturkosu · 2026-06-08T22:43:27Z

bf16 RNE encoder now has unit-test coverage

Addressing the review request to lock in the bit-level _fp32_to_bf16_u16 / _bf16_u16_to_fp32 rounding: CPU tests were added in dispatcher/tests/test_gemm_utils.py (on the stacked layouts branch #8191, which contains these helpers). Coverage includes exact round-trip, <= 2^-8 relative error, round-to-nearest-even ties in both directions, Inf/NaN handling, and uint16/2-byte sizing. 10 passed.

…spatcher-gemm-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 <noreply@anthropic.com>

Addresses the three review items on the TE->Dispatcher GEMM bridge driver, scoped to this foundation PR's fp16/rcr surface (bf16/layouts follow in the #8190/#8191 stack): 1. Example configs to sweep - gemm_full_benchmark.py defaults to the selected variant's configs/default_ci_config.json (small CI sweep) when no config is passed, and to configs/example_problems.json when --problems is omitted; configs/default_config.json remains the full sweep. - New gemm_universal/configs/example_problems.json (square / rectangular / large M,N,K). Nightly-test JSON drops into the same configs/ dir -- no driver change needed. 2. Multi-GPU launch in parallel (supersedes grouped_conv's serial-GPU design) - Phase 3 fans the (kernel x problem) work across every visible GPU: one worker thread per device pulls batches from a shared queue and spawns a disposable subprocess pinned with HIP_VISIBLE_DEVICES, so an N-GPU box runs ~Nx faster while keeping per-batch fault isolation. - Devices auto-detected (HIP_VISIBLE_DEVICES, then rocm-smi/amd-smi); override with --devices (count, explicit ids, or all). 3. Variant organization + README + deprecation note - --variant selects the per-variant configs/ directory. - New README "Dispatcher Bridge Workflow" section: scripts, per-variant config layout, run examples, multi-GPU explanation, supported surface (fp16/rcr here), and a deprecation note for the legacy *_instance_builder.py generators. Driver --dtype/--layout choices stay fp16/rcr to match this PR's dispatcher host path; run_one_gemm_kernel.py (fp16 host gen) is unchanged.

ozturkosu · 2026-06-17T05:17:31Z

Closing this draft PR as its changes are fully superseded by #8479 ([CK_TILE] TE -> Dispatcher GEMM bridge — all layouts + fp16/bf16).

PR #8479 is the consolidated, single-commit re-roll of the entire GEMM bridge work that was previously split across the stacked draft series (#8187 → #8190 → #8191 → #8193). Every file touched by this PR is included in #8479 at the same or an updated path. Please track further work in #8479.

ozturkosu mentioned this pull request Jun 8, 2026

[CK_TILE] Add Tile Engine -> Dispatcher bridge for GEMM #8123

Closed

4 tasks

github-actions Bot added the project: composablekernel label Jun 8, 2026

ozturkosu self-assigned this Jun 8, 2026

ozturkosu requested a review from Copilot June 8, 2026 22:14

Copilot started reviewing on behalf of ozturkosu June 8, 2026 22:14 View session

ozturkosu mentioned this pull request Jun 8, 2026

[CK_TILE] Add multi-layout support to GEMM TE -> Dispatcher bridge #8191

Closed

Copilot AI reviewed Jun 8, 2026

View reviewed changes

ozturkosu mentioned this pull request Jun 8, 2026

[CK_TILE] Add TE -> Dispatcher bridge GEMM example #8193

Closed

2 tasks

Merge branch 'muozturk/dispatcher-gemm-registry-key' into muozturk/di…

514ecdd

…spatcher-gemm-bf16

ozturkosu mentioned this pull request Jun 8, 2026

[CK_TILE] Derive GEMM Dispatcher registry key from codegen traits #8187

Closed

3 tasks

Merge branch 'muozturk/dispatcher-gemm-registry-key' into muozturk/di…

b361af2

…spatcher-gemm-bf16

assistant-librarian Bot added the organization: ROCm label Jun 8, 2026

Muhammed Ozturk and others added 3 commits June 8, 2026 19:02

Merge branch 'muozturk/dispatcher-gemm-registry-key' into muozturk/di…

50a742d

…spatcher-gemm-bf16

Merge branch 'muozturk/dispatcher-gemm-registry-key' into muozturk/di…

c685192

…spatcher-gemm-bf16

[CK_TILE] Widen GEMM benchmark --dtype choices to include bf16

b3f43db

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 <noreply@anthropic.com>

ozturkosu mentioned this pull request Jun 10, 2026

[CK_TILE] TE->Dispatcher GEMM bridge (consolidated: key + bf16 + layouts + parity + example) #8261

Closed

3 tasks

ozturkosu closed this Jun 17, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CK_TILE] Add bf16 support to GEMM TE -> Dispatcher bridge#8190

[CK_TILE] Add bf16 support to GEMM TE -> Dispatcher bridge#8190
ozturkosu wants to merge 6 commits into
muozturk/dispatcher-gemm-registry-keyfrom
muozturk/dispatcher-gemm-bf16

ozturkosu commented Jun 8, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

ozturkosu commented Jun 8, 2026

Uh oh!

ozturkosu commented Jun 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ozturkosu commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why this is small

The fix

Design decisions

Validation (gfx942 / MI300X)

Remaining (tracked, not in this PR)

Next

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

ozturkosu commented Jun 8, 2026

bf16 RNE encoder now has unit-test coverage

Uh oh!

ozturkosu commented Jun 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ozturkosu commented Jun 8, 2026 •

edited

Loading