feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts) by ozturkosu · Pull Request #8997 · ROCm/rocm-libraries

ozturkosu · 2026-07-01T05:11:08Z

Re-opened from #8479 with a compliant branch name (users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16). Supersedes #8479.

Summary

This PR routes the Tile Engine (TE) regular-GEMM sweep through the Dispatcher,
making the Dispatcher the single source of truth for codegen → build → runtime
while the Tile Engine keeps only the config search space and the benchmark
loop. It is the consolidated, single-commit GEMM bridge covering all four
layouts (rcr/rrr/crr/ccr) and both fp16 and bf16.

It is a clean re-roll of the earlier bridge work (previously split across
#8123 + the stacked key/bf16/layouts/parity/example PRs and consolidated in
#8261). Those branches accumulated unrelated cross-project commits through repeated
develop merges; this branch is a single clean commit off the latest develop
containing only the GEMM-bridge files. It supersedes and replaces #8123 / #8261.

Motivation

The Tile Engine historically owned its own codegen/build/runtime for GEMM
(tile_engine/ops/gemm/gemm_universal/). The consolidation goal is for the
Dispatcher to own all of that — exactly as it already does for FMHA and
Grouped Conv — so there is one kernel-generation/build/runtime path and the
TE shrinks to a config+benchmark frontend. This PR brings regular GEMM in line
with that reference binding.

The binding (mirrors the FMHA/Conv reference, six stages)

Config JSON (TE side) — the sweep search space lives in
tile_engine/ops/gemm/configs/ (flat op-root layout, matching the
fmha/ and grouped_conv/ bridges).
Codegen (Dispatcher) — dispatcher/codegen/unified_gemm_codegen.py emits
one fully-typed .hpp per kernel; GemmKernelConfig.name reproduces
KERNEL_NAME byte-for-byte (the thread tying config → kernel → runtime).
Compile to .so — a single static gemm_ctypes_lib.cpp is force-included
(-include <kernel.hpp>); one .so per kernel.
Flat extern "C" ABI — dispatcher_run_gemm(A, B, C, M, N, K, time_ms) +
the kernel-name enumeration entry points. Host-pointer memory model (the C
lib hipMallocs internally) — the FMHA-forward branch of the reference.
Python ctypes wrapper — dispatcher/python/gemm_utils.py
(GemmDispatcherLib + GpuGemmRunner).
TE driver (3 phases) — gemm_full_benchmark.py (parallel codegen+build →
expand_sweep → subprocess-isolated benchmark) + the disposable per-kernel
worker run_one_gemm_kernel.py.

What's included

Bridge core

dispatcher/codegen/unified_gemm_codegen.py — GEMM codegen, byte-exact naming.
dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp — flat C ABI, host-pointer model.
dispatcher/python/gemm_utils.py — GemmKernelConfig, multi-kernel build
(setup_multiple_gemm_dispatchers), expand_sweep, one-.so-per-kernel.
tile_engine/ops/gemm/gemm_full_benchmark.py + run_one_gemm_kernel.py —
3-phase, multi-GPU, subprocess-isolated driver/worker.

Feature surface (the point of this PR)

All four layouts rcr/rrr/crr/ccr (row-major C only — ck_tile rejects
column-major C at build) with layout-aware host transpose.
fp16 + bf16 (bf16 via uint16 byte-encoding; dtype derived from kernel name).
Trait-derived registry KernelKey — replaces the earlier hard-coded
fp16/rcr key so the registry path generalizes across dtype/layout/tile.

Correctness & performance hygiene

--verify opt-in fp32 numpy-reference gate (global max|out-ref|/max|ref|),
verified/max_rel columns in the CSV; a mismatch counts as a failure.
Tile Engine AMDGPU -mllvm codegen-flag parity (without these the kernel
builds with different occupancy and the timing diverges) and
arch-validated tile filtering against the real pipeline/scheduler.
Multi-GPU fan-out across all visible GPUs (--devices, device-pinned
HIP_VISIBLE_DEVICES workers).

Example & tests

dispatcher/examples/gemm/python/12_te_bridge.py — runnable end-to-end example.
dispatcher/tests/test_gemm_parity.py, test_gemm_utils.py, and a parity
regression harness.

Cleanup

Removes the legacy standalone gemm_universal build path
(gemm_universal_instance_builder.py, *_benchmark*.{py,cpp,hpp},
gemm_universal/CMakeLists.txt) and the old test/ck_tile/gemm_tile_engine/
harness; promotes the sweep configs to the flat op-root configs/.

Design decisions (consistent with the reference)

Host-pointer memory ownership (C lib owns device memory) — matches
FMHA-forward; the Python runner passes host numpy arrays straight through.
One .so per kernel — packaging choice; the multi-kernel name ABI is
retained (get_kernel_name_at(0) reports the single kernel), so the Python
enumeration path is unchanged from FMHA/Conv.
Flat configs/ at the op root — matches the fmha//grouped_conv/
convention; the not-yet-bridged variants keep their per-variant configs/
dirs, selected by --variant.

Validation (gfx942 / MI300X)

Bridge build + benchmark + --verify across fp16 and bf16 and all
four layouts, checked against an fp32 numpy reference (A @ B).
Name parity holds end-to-end: each .so's reported runtime name equals
GemmKernelConfig(...).name.
bf16 passes under a widened fp16/bf16 tolerance; fp16 within the standard
max_rel gate.

Test plan

gemm_full_benchmark.py --verify over configs/default_ci_config.json for
fp16 and bf16, each of rcr/rrr/crr/ccr.
unified_gemm_codegen.py emits a header whose stem == GemmKernelConfig.name.
setup_multiple_gemm_dispatchers builds + links each config against
gemm_ctypes_lib.cpp.
pytest dispatcher/tests/test_gemm_parity.py dispatcher/tests/test_gemm_utils.py.
examples/gemm/python/12_te_bridge.py runs end to end.

Notes

Single clean commit off the latest develop; the diff is 35 files, all under
projects/composablekernel/ (dispatcher + tile_engine/ops/gemm + test/ck_tile).
Supersedes [CK_TILE] Add Tile Engine -> Dispatcher bridge for GEMM #8123 and [CK_TILE] TE->Dispatcher GEMM bridge (consolidated: key + bf16 + layouts + parity + example) #8261, which will be closed.
Stream-K ([CK_TILE] Add Tile Engine -> Dispatcher bridge for Stream-K GEMM #8136) and grouped GEMM are separate bridge efforts, not in this PR.

Consolidated, single-commit GEMM bridge routing the Tile Engine regular-GEMM sweep through the Dispatcher (codegen -> build -> runtime), so the Dispatcher is the single source of truth and the Tile Engine owns only the config search space and the benchmark loop. Mirrors the FMHA/Conv reference binding end to end. Scope: - Regular GEMM bridge: unified_gemm_codegen.py, gemm_ctypes_lib.cpp (flat extern "C" ABI, host-pointer model), gemm_utils.py (GemmKernelConfig with byte-exact .name, one-.so-per-kernel build), 3-phase TE driver + subprocess worker (gemm_full_benchmark.py / run_one_gemm_kernel.py). - Trait-derived registry KernelKey (replaces the hard-coded fp16/rcr key). - bf16 support and all four layouts (rcr/rrr/crr/ccr; row-major C only). - Tile Engine AMDGPU -mllvm codegen-flag parity + arch-validated tile filtering. - --verify fp32-reference correctness gate; multi-GPU fan-out. - Runnable example (examples/gemm/python/12_te_bridge.py) and parity/unit tests. - Removes the legacy standalone gemm_universal build path and the old test/ck_tile/gemm_tile_engine harness; promotes sweep configs to the op-root flat configs/ directory (fmha/grouped_conv convention). Validated on gfx942 / MI300X (fp16 + bf16, all four layouts) against an fp32 numpy reference via --verify.

The bridge dispatcher's tile-divisibility gate rejected any problem where M % TileM != 0 for every layout, returning status -2 ("No suitable kernel") at runtime even though the .so built fine. This wrongly excluded bf16 rcr/rrr kernels with a non-power-of-two TileM (e.g. 192) on standard shapes like 1024^3 -- cases Old-TE compiles, runs, and verifies as correct. Root cause: supports() was layout-blind, while the underlying ck_tile::GemmKernel::IsSupportedArgument only constrains a dimension when an operand whose inner axis is that dimension participates without padding: RowMajor A -> K, ColMajor A -> M RowMajor B -> N, ColMajor B -> K RowMajor C -> N, ColMajor C -> M So for rcr (RowMajor A & C) M is never gated, which is why Old-TE runs M=192 tiles on M-indivisible problems. Make supports() compute require_m/n/k from the kernel key's A/B/C layouts so it mirrors IsSupportedArgument exactly (also honoring k_batch in the K grain). Anything it now lets through is still validated by the kernel's own IsSupportedArgument inside launch(), so the bridge stays a strict functional equivalent of Old-TE. Applied to both generated_tile_backend.hpp (the GEMM .so path) and the sibling tile_backend.hpp. Validated on gfx942 (MI300X): 85 previously status-2 rcr/rrr bf16 192-tile .so now run at 1024^3 (Old-TE runs the same, verification correct); the 8 remaining rejects are tile N=192 cases that Old-TE also reports "Arguments not supported" at N=1024 -- parity preserved in both directions.

…oding rcr dispatcher_initialize() in gemm_ctypes_lib.cpp hardcoded the KernelKey layout to rcr (RowMajor/ColMajor/RowMajor) for every kernel. Now that supports() is layout-aware, that wrong key layout makes the dispatcher reject valid problems: a crr kernel does not gate K (neither A=ColMajor nor B=RowMajor has K as its inner axis), but with a hardcoded rcr key supports() applies rcr's K-gate and returns status -2 for TileK=192 problems (e.g. crr 64x64x192 at 1024^3) that Old-TE compiles, runs, and verifies (~87 TFLOPS). Derive signature.layout_a/b/c from the force-included kernel's own ALayout/BLayout/CLayout types via std::is_same_v with tensor_layout::gemm::RowMajor. The key now matches the kernel, so the layout-aware gate is correct for all four layouts. Execution was already layout-correct (the kernel uses its own compile-time layouts); only the host-side selection metadata was wrong. Validated on gfx942 (MI300X): crr 64x64x192 now runs on the bridge (93 TFLOPS), restoring parity with Old-TE.

The >=20% bridge-vs-old-TE perf gaps in the parity sweep are a harness artifact: the sweep timed the bridge in-process but timed old-TE via its separate standalone benchmark binary, which runs the byte-identical kernel at a lower sustained SCLK. Measured through one harness the gap is <1%. ab_same_harness.py removed that artifact but hardcoded the old-TE header dir to fp16/rcr. Derive it per stem as <base>/<dtype>/<layout> so one run covers rcr/rrr/ccr/crr and fp16+bf16, add a --stems-file/--csv resume-aware sweep mode, and use the median (not max) per point.

For a full ~2000-stem sweep on a single GPU: batch all shapes into one worker call per side (5x fewer process startups), cache the compiled old-TE .so, and add a parallel --build-only pre-pass so hipcc compilation uses all CPU cores while GPU measurement stays serial.

… guard) The bridge-vs-old-TE A/B reported phantom regressions from two MEASUREMENT bugs, not real codegen gaps: - ab_same_harness.py built the old-TE side WITHOUT the TE codegen flags the bridge (and real old-TE's own CMake) use, so -enable-post-misched defaulted back on and old-TE ran ~10-40% faster -> the bridge looked regressed when it is at parity. Now both sides build with identical flags. - ab_efficient_sweep.py measured whatever libgemm_<stem>.so existed with no freshness check, so 3-day-old binaries built from an obsolete codegen showed up as -78%/+703% gaps. Added a guard: skip any .so older than its generated header (treated as missing) instead of reporting a phantom gap. With both fixes the 41 former >15% outlier stems measure within +/-10% (median +0.01%); no bridge codegen regression exists. Note: a separate, deliberately UNCOMMITTED perf change in gemm_utils.py (gate -enable-post-misched=0 on persistent) gives non-persistent large tiles ~9-40%; held back pending a broader persistent-kernel no-regression sweep.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

The bridge compiles each kernel .so with a hand-maintained hipcc flag list (dispatcher/python/gemm_utils.py) that had drifted from Tile Engine's CMake flags, so the bridge .so and the TE benchmark were not compiled apples-to-apples: * MISSING -mllvm -amdgpu-coerce-illegal-types=1 (TE's CMakeLists.txt adds it when the compiler accepts it; the bridge build never did) * EXTRA -mllvm -enable-noalias-to-md-conversion=0 (not a TE GEMM flag; it only appears in standalone CK examples/tests, never the TE gemm path) Align the bridge's backend codegen flags with the exact set the TE gemm_universal benchmark TU is built with. The coerce flag is added through a cached hipcc probe that mirrors TE's check_cxx_compiler_flag, so the bridge stays matched to TE on every toolchain (present where TE has it, skipped where TE's CMake would skip it too). The generated kernel source was already identical between the two engines; this makes their compilation identical as well.

…y_diag/regression Old-TE must remain until the dispatcher bridge implements every datatype Old-TE supports, so revert the Old-TE removal from the bridge commit and re-wire its build: * restore test/ck_tile/gemm_tile_engine/* (10 files) * restore tile_engine/ops/gemm/gemm_universal/* (6 files: benchmark / instance builder / profiler / single-bench / CMakeLists) * re-add `add_subdirectory(gemm_universal EXCLUDE_FROM_ALL)` in tile_engine/ops/gemm/CMakeLists.txt; restore test/ck_tile/CMakeLists.txt to the develop state (gemm_tile_engine entry kept commented, as in develop) Also drop the parity_diag/regression dev scripts that should not ship in the PR: * dispatcher/parity_diag/regression/ab_efficient_sweep.py * dispatcher/parity_diag/regression/ab_same_harness.py

…whitespace - Add AMD copyright/SPDX header to gemm_full_benchmark.py and run_one_gemm_kernel.py (CK requires a header on every source file). - Remove a trailing-whitespace blank line in generated_tile_backend.hpp that would trip the whitespace/clang-format CI gate.

….g. 192) The CShuffle epilogue stores the accumulator back through LDS in power-of-two MRepeat/NRepeat chunks, where MRepeat = tile_m / (wave_m * warp_tile_m) (and likewise N). A tile whose per-wave repeat is not a power of two (or whose tile dim is not divisible by wave*warp_tile) is mis-stored and produces numerically WRONG results at runtime -- yet it still passes the ctypes validator and the epilogue's static_asserts, so it compiles and silently returns garbage. Observed on MI350 for tile_m=192 (MRepeat = 192/(2*32) = 3) and tile_n=192 (e.g. 64x192x64_1x4x1, 192 not divisible by 4*32): both verified incorrect (fp32 reference, max_rel ~1.2-1.4) on the bridge AND Tile Engine, at every shape including shapes divisible by 192. Power-of-two tiles (64/128/256) are unaffected; a control 256-tile verifies cleanly (max_rel ~4e-4). Add a validity gate in both tile-expansion paths: * unified_gemm_codegen.py::_get_tile_configs (codegen CLI path) * gemm_utils.py::expand_sweep (bridge .so build path; this path only ran the ctypes validate_kernel_config, which does not catch this) so invalid tiles are dropped instead of emitted/run. tile_k is unaffected (the K reduction has no CShuffle store constraint).

therock-pr-bot · 2026-07-01T05:15:29Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

therock-pr-bot · 2026-07-01T05:15:30Z

🎉 All checks passed! This PR is ready for review.

…r layer Add the full Tile Engine driver layer for BQuantGrouped GEMM, mirroring the regular-GEMM bridge from PR #8997. This includes the 3-phase benchmark driver (compile -> load problems -> multi-GPU subprocess benchmark), subprocess worker, sweep configs, and fixes to codegen, ctypes lib, CMakeLists, and test assertions for the permute_n epilogue. - Codegen: emit GroupSizeK in generated kernel struct and CK_TILE_SINGLE_KERNEL_INCLUDE block for TE binary compatibility - Ctypes lib: BQ preshuffle for BPreshuffleQuant=true kernels; correct KernelKey registration from SelectedKernel compile-time members; dtype detection via constexpr type traits - CMakeLists: derive GFX_ARCH from CMAKE_HIP_ARCHITECTURES instead of hardcoding gfx950 - New gemm_bquant_full_benchmark.py: 3-phase driver with multi-GPU subprocess isolation, dtype/arch/compile-only flags - New run_one_bquant_kernel.py: subprocess worker with input caching - New dispatcher sweep configs (CI + full) and example problems - Tests: GroupSizeK codegen validation, expand_bquant_sweep coverage, fix 12 test assertions for permute_n epilogue (N_repeat even) Co-Authored-By: Claude <noreply@anthropic.com>

…9000) > Re-opened from #8130 with a policy-compliant branch name (`users/muozturk/ck-tile/dispatcher-te-bridge-grouped-gemm`). Supersedes #8130. ## What this PR does Routes the **grouped_gemm** variant through the Tile Engine (TE) → Dispatcher **bridge**: TE only generates configs and benchmarks; the Dispatcher owns codegen, build, and runtime. This is the grouped counterpart of the regular-GEMM bridge (#8123/#8479), the fp8/bf8/int8 bridge (#8887), and the Stream-K bridge (#8136). **This PR now also contains the grouped Dispatcher codegen** that previously lived in #8075 — that PR has been **closed in favor of this one** to keep the grouped codegen in a single place (it was otherwise duplicated across both). ## Why grouped needs special handling Grouped GEMM is **multi-problem**: one launch runs a *list* of `(M, N, K)` sub-problems with arrays of A/B/C device pointers. 1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`) cannot express a list of problems. 2. The generated registry wrapper (`generated_tile_backend.hpp::run()`) hard-codes the single-problem launch and won't compile against a grouped `SelectedKernel`. So the grouped path **bypasses the registry**: a dedicated ctypes lib calls the generated `SelectedKernel::launch(descs, stream)` directly and reports the name from the compile-time `KERNEL_NAME` macro. ## Changes **Codegen (absorbed from #8075)** - `codegen/arch_filter.py` — `GEMM_GROUPED` operator tile constraints. - `codegen/unified_gemm_codegen.py` — `GemmVariant.GROUPED`, the grouped launch generator (DeviceMem internal workspace via `MakeKargs`, persistent/non-persistent grid), `grouped` in `--variants`. - `examples/gemm/cpp/02_grouped_gemm_driver.cpp` — standalone, layout/dtype-generic grouped driver with per-group reference verification. - `codegen/README.md` + `examples/gemm/cpp/README.md` — grouped sections. **Bridge** - `bindings/ctypes/grouped_gemm_ctypes_lib.cpp` — multi-problem, registry-bypass C ABI; per-group device alloc/copy; strides derived from the compile-time `ALayout/BLayout/CLayout`; warmup/repeat timing matched to Old-TE (`CK_TILE_BENCH_WARMUP/REPEAT`). - `python/gemm_utils.py` — `GroupedGemmProblem`/`GroupedGemmResult`, `GpuGroupedGemmRunner`, `run_grouped`, fp16/bf16/fp8(E4M3 FNUZ)/bf8(E5M2 FNUZ) codecs, output-dtype-aware C buffer. - `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` + `run_one_grouped_gemm_kernel.py` — TE driver + worker for the parity sweep. - `bindings/ctypes/GROUPED_GEMM_BRIDGE.md` — design README. ## Coverage (= Old-TE grouped runnable set on develop) | Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) | |---|---|---|---|---| | rcr / rrr / ccr / crr | ✓ | ✓ | ✓ | ✓ | C is always row-major. `int8` (rejected by the TE grouped builder) and `fp32`/`fp64` (no MFMA warp tiles) are excluded on both sides. ## Parity vs Old-TE (MI300X / gfx942) Apples-to-apples (same warmup=50/repeat=100 both sides, A/B interleaved, single GPU, both engines rebuilt fresh, stale-`.so` guard, matched compile flags): - **Correctness: 64/64 PASS.** - **Performance: 64/64 within ±15%.** - The 5 small-shape (1024³ fp8/bf8) rows that initially read >15% were proven by `rocprof` to be a **measurement-harness artifact** (Old-TE's JSON `latency(ms)` rounded to 2 decimals → 30–50% TFLOPS swing on ~0.02 ms kernels), **not** a kernel/codegen difference — bridge and Old-TE launch byte-identical kernels (same grid/VGPR/SGPR, duration ≤3.22%); full-precision re-measure collapses all 5 to <3%. ## Notes - Targets `develop`. Depends on #8997 (fp16/bf16 bridge) and #8998 (fp8/bf8/int8 bridge) merging to `develop` first; until then this PR's diff also shows their content, after which it reduces to the grouped-only files. - Supersedes #8075 (closed). --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Muhammed Ozturk <muozturk@ctr2-alola-ctrl-01.amd.com> Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>

## Summary Adds a microscaling-GEMM (`mx_gemm`) bridge from the Old Tile-Engine into the dispatcher's ctypes path for gfx950/MI350. Supports **fp4** (`pk_fp4_t`, e2m1) and **fp8** (e4m3), **rcr** layout, `comp_async` + `cshuffle` + `intrawave`, fixed `16x16x128` warp tile, `k_batch == 1`. Per-32-K `e8m0` block scales are pre-shuffled on-host exactly as the Old-TE profiler does. ## Motivation Extend the TE→Dispatcher bridge family to the microscaling GEMM op so the dispatcher can launch byte-identical `mx_gemm` kernels with block scaling. Sibling to the other bridge PRs (see below). ## Design note - **Byte-exact kernels:** the codegen reuses Old-TE `MxGemmKernelBuilder._generate_kernel_instance` directly rather than re-implementing header assembly, guaranteeing the emitted kernel matches Old-TE. It only strips the stale `ck_tile/ops/gemm_mx.hpp` umbrella include (absent on develop; the mx pipeline is pulled in via `ck_tile/ops/gemm.hpp`). - **Registry bypass:** `mx_gemm`'s `launch` takes `ck_tile::MxGemmHostArgs` with per-32-K `e8m0` scales that the generic dispatcher backend cannot express, so the ctypes lib builds `MxGemmHostArgs` from plain C arrays and calls `SelectedKernel::launch()` directly — the same direct-launch pattern as the batched/multi-D bridges. - **Scale pre-shuffle** mirrors `mx_gemm_profiler.hpp`; pack params are derived from `SelectedKernel` tile dims at compile time (not hardcoded). - **Packing-correct byte accounting:** device buffers are sized via `HostTensor::get_element_space_size_in_bytes()`, which divides by `numeric_traits<T>::PackedSize`, so fp4 (`PackedSize==2`, two e2m1 per byte) and fp8 (`PackedSize==1`) are both correct. ## Verification - fp8 derisk: PASS (`max_rel = 0.0`, kernel-name match, output fully non-zero) on gfx950 - fp4: GPU-verified on gfx950 - `clang-format-18 -style=file` clean; Python `py_compile` clean ## Sibling PRs - #8997 — TE→dispatcher GEMM bridge (fp16/bf16, all layouts) — foundational (merged) - #8998 — fp8/bf8/int8 bridge (merged) - #9000 — grouped GEMM - #9028 — stream-K GEMM - #9305 — multi-ABD GEMM - #9306 — batched GEMM - #9307 — preshuffle GEMM - #9308 — multi-D GEMM - #9328 — batched-contraction GEMM --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>

ISSUE ID: #8997 ## Motivation The CK Tile dispatcher could already generate and launch regular GEMM through the TileEngine → Dispatcher bridge, but it had no path for the multi-tensor **gemm_multi_abd** op. Multi-ABD is used when a GEMM needs to combine several A and B operands and fuse several D operands in the epilogue (`E = cde_op(a_op(As) @ b_op(Bs), {Ds})`), which is a real Old-TE capability with no dispatcher equivalent. This PR closes that gap so Python callers can drive multi_abd through the dispatcher at parity with the legacy Tile Engine version, without touching C++. It follows the divergent-ABI pattern established by the grouped bridge (#9000) because multi_abd needs **arrays** of A/B/D device pointers, not the single-pointer regular GEMM ABI. The capability set matches the Old-TE `gemm_multi_abd_instance_builder.py` exactly: `fp16`, `rcrr` layout, configurable A/B/D tensor counts, and the element-wise op set `{PassThrough, AddScale, MultiDMultiply, MultiDAdd}`. ## Test Plan - Run the CPU-only unit tests (no GPU required): `python3 -m pytest dispatcher/tests/test_multi_abd_bridge.py -v` - On-GPU numeric verification through the bridge launch path (gfx942 / MI300X), 512x512x512 fp16 rcrr, across the default 2/2/2 all-PassThrough config and non-PassThrough element-wise ops. - Confirm the CI and default config expansions yield the expected kernel counts. ## Test Result - CPU-only unit tests pass (10 passed). - Numeric verification (bridge launch path), 512x512x512 fp16 rcrr: - default 2/2/2 all-PassThrough: `max_rel = 2.9e-4` - CDE = MultiDAdd: `max_rel = 5.7e-4` - A-op = MultiDAdd: `max_rel = 4.1e-4` - all far below the fp16 tolerance (2e-2); 0 failed measurements. - CI config expands to 16 arch-valid kernels; `default_config.json` → 8896. - `standard` variant `expand_sweep` regression clean. - clang-format-18 (18.1.8) clean on `gemm_multi_abd_ctypes_lib.cpp`. - Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16 rcrr, 16 stems × 5 shapes = 80 rows, interleaved, fair 50/100/flush/rotating both sides): **at parity** — median gap -0.24%, mean -0.66%, 100% within ±15%, 87.5% within ±5% (range [-9.57%, +5.35%]). See the parity comment for details. --- **Related PRs / references (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in the same bridge effort tracked across those PRs. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>

ISSUE ID: #8997 ## Motivation The TileEngine → Dispatcher bridge had no path for the **gemm_multi_d** op, which fuses one or more extra D operands into the GEMM epilogue (`E = elementwise_op(A@B, D0, D1, ...)`). This is a real Old-TE capability used for fused bias/residual-style epilogues with no dispatcher equivalent, so this PR adds a complete bridge so the dispatcher can generate, build, and launch multi_d at parity with the legacy Tile Engine version. The capability set matches the Old-TE `gemm_multi_d_instance_builder.py` exactly: `fp16`, the 4-char layouts `{rcrr, rrrr, ccrr, crrr}` (A/B vary, C and D row-major), the element-wise ops `{MultiDAdd, MultiDMultiply, PassThrough}`, and a swept number of D tensors (1 and 2). It follows the registry-bypass bridge pattern used by the grouped (#9000) and stream-K (#9028) bridges. ## Test Plan - Run the CPU-only unit tests (no GPU required): `python3 -m pytest dispatcher/tests/test_multi_d_bridge.py -v` - On-GPU numeric verify over the full capability matrix (fp16 × {rcrr, rrrr, ccrr, crrr} × {MultiDAdd, MultiDMultiply} × {num_d 1, 2} = 16 combos) at M=N=K=1024 against an fp32 reference, gate 2e-2. - Confirm the CI config builds real kernels and the sweep covers all ops × D counts. ## Test Result - CPU-only unit tests pass (10 passed). - On-GPU numeric verify: 16/16 combos pass at M=N=K=1024, worst-case `max_rel = 6.16e-4` (~30x under the 2e-2 gate). Col-major `ccrr` / `crrr` have real on-GPU numeric evidence. - CI config now builds real kernels (was zero); the sweep expands evenly across `{MultiDAdd, MultiDMultiply} × {num_d 1, 2}` per layout. - clang-format (18.1.8) clean on `multi_d_gemm_ctypes_lib.cpp`. - Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16, 4 layouts × 2 ops × num_d=1 = 8 stems × 5 shapes = 40 rows, interleaved, fair 50/100/flush/rotating both sides): **at parity, bridge consistently faster** — median gap +9.44%, 100% within ±15% (range [+3.76%, +14.99%]; positive = bridge faster, from the registry-bypass direct launch avoiding the Old-TE profiler's per-call overhead). num_d=1 is the fair slice since the Old-TE `gemm_multi_d` benchmark is single-D. See the parity comment for details. --- **Related PRs / references (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in the same bridge effort tracked across those PRs. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

ISSUE ID: #8997 ## Motivation The CK Tile dispatcher could already generate and launch regular GEMM through the TileEngine → Dispatcher bridge, but it had no path for the multi-tensor **gemm_multi_abd** op. Multi-ABD is used when a GEMM needs to combine several A and B operands and fuse several D operands in the epilogue (`E = cde_op(a_op(As) @ b_op(Bs), {Ds})`), which is a real Old-TE capability with no dispatcher equivalent. This PR closes that gap so Python callers can drive multi_abd through the dispatcher at parity with the legacy Tile Engine version, without touching C++. It follows the divergent-ABI pattern established by the grouped bridge (#9000) because multi_abd needs **arrays** of A/B/D device pointers, not the single-pointer regular GEMM ABI. The capability set matches the Old-TE `gemm_multi_abd_instance_builder.py` exactly: `fp16`, `rcrr` layout, configurable A/B/D tensor counts, and the element-wise op set `{PassThrough, AddScale, MultiDMultiply, MultiDAdd}`. ## Test Plan - Run the CPU-only unit tests (no GPU required): `python3 -m pytest dispatcher/tests/test_multi_abd_bridge.py -v` - On-GPU numeric verification through the bridge launch path (gfx942 / MI300X), 512x512x512 fp16 rcrr, across the default 2/2/2 all-PassThrough config and non-PassThrough element-wise ops. - Confirm the CI and default config expansions yield the expected kernel counts. ## Test Result - CPU-only unit tests pass (10 passed). - Numeric verification (bridge launch path), 512x512x512 fp16 rcrr: - default 2/2/2 all-PassThrough: `max_rel = 2.9e-4` - CDE = MultiDAdd: `max_rel = 5.7e-4` - A-op = MultiDAdd: `max_rel = 4.1e-4` - all far below the fp16 tolerance (2e-2); 0 failed measurements. - CI config expands to 16 arch-valid kernels; `default_config.json` → 8896. - `standard` variant `expand_sweep` regression clean. - clang-format-18 (18.1.8) clean on `gemm_multi_abd_ctypes_lib.cpp`. - Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16 rcrr, 16 stems × 5 shapes = 80 rows, interleaved, fair 50/100/flush/rotating both sides): **at parity** — median gap -0.24%, mean -0.66%, 100% within ±15%, 87.5% within ±5% (range [-9.57%, +5.35%]). See the parity comment for details. --- **Related PRs / references (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in the same bridge effort tracked across those PRs. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>

ISSUE ID: #8997 ## Motivation The TileEngine → Dispatcher bridge had no path for the **gemm_multi_d** op, which fuses one or more extra D operands into the GEMM epilogue (`E = elementwise_op(A@B, D0, D1, ...)`). This is a real Old-TE capability used for fused bias/residual-style epilogues with no dispatcher equivalent, so this PR adds a complete bridge so the dispatcher can generate, build, and launch multi_d at parity with the legacy Tile Engine version. The capability set matches the Old-TE `gemm_multi_d_instance_builder.py` exactly: `fp16`, the 4-char layouts `{rcrr, rrrr, ccrr, crrr}` (A/B vary, C and D row-major), the element-wise ops `{MultiDAdd, MultiDMultiply, PassThrough}`, and a swept number of D tensors (1 and 2). It follows the registry-bypass bridge pattern used by the grouped (#9000) and stream-K (#9028) bridges. ## Test Plan - Run the CPU-only unit tests (no GPU required): `python3 -m pytest dispatcher/tests/test_multi_d_bridge.py -v` - On-GPU numeric verify over the full capability matrix (fp16 × {rcrr, rrrr, ccrr, crrr} × {MultiDAdd, MultiDMultiply} × {num_d 1, 2} = 16 combos) at M=N=K=1024 against an fp32 reference, gate 2e-2. - Confirm the CI config builds real kernels and the sweep covers all ops × D counts. ## Test Result - CPU-only unit tests pass (10 passed). - On-GPU numeric verify: 16/16 combos pass at M=N=K=1024, worst-case `max_rel = 6.16e-4` (~30x under the 2e-2 gate). Col-major `ccrr` / `crrr` have real on-GPU numeric evidence. - CI config now builds real kernels (was zero); the sweep expands evenly across `{MultiDAdd, MultiDMultiply} × {num_d 1, 2}` per layout. - clang-format (18.1.8) clean on `multi_d_gemm_ctypes_lib.cpp`. - Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16, 4 layouts × 2 ops × num_d=1 = 8 stems × 5 shapes = 40 rows, interleaved, fair 50/100/flush/rotating both sides): **at parity, bridge consistently faster** — median gap +9.44%, 100% within ±15% (range [+3.76%, +14.99%]; positive = bridge faster, from the registry-bypass direct launch avoiding the Old-TE profiler's per-call overhead). num_d=1 is the fair slice since the Old-TE `gemm_multi_d` benchmark is single-D. See the parity comment for details. --- **Related PRs / references (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in the same bridge effort tracked across those PRs. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

…idges (#10612) ISSUE ID: #8997 ## Summary This PR combines two sibling TileEngine → Dispatcher bridge ops into a single PR: - **batched GEMM** (previously #9306) - **batched_contraction** (previously #9328) Both follow the same **direct-launch, registry-bypass** pattern (as the stream-K bridge #9028), because their launch ABIs carry variable-length / batch-specific arguments the single-pointer registry backend cannot express. The two ops touch disjoint files except `dispatcher/tests/CMakeLists.txt`, where both GPU-correctness test registrations are kept. ## Motivation The dispatcher had no path for batched GEMM (same GEMM across many independent problems with per-batch strides) or batched_contraction (generalized batched tensor contraction `E[G.., M.., N..] = sum_K A[G.., M.., K..] * B[G.., N.., K..]` with multi-dim G/M/N/K index groups). Both are real Old-TE ops; this bridge lets Python callers generate, build, and launch them at parity with the legacy Tile Engine — without writing C++. ## Test Plan / Result - CPU-only unit tests for both bridges + gemm_utils: **74 passed** (`pytest dispatcher/tests/test_batched_bridge.py dispatcher/tests/test_batched_contraction_bridge.py dispatcher/tests/test_gemm_utils.py`). - Batched GEMM: end-to-end name-parity + correctness across batch counts 1/2/4/8 (`max_rel ~5e-4`), non-packed strides, split-K; full `default_config` codegen 6672 kernels / 0 failures. - Batched contraction: on-GPU verify (gfx950) across dtype × layout × shape × multi-dim × pipeline + D-tensor epilogue, all PASS. - clang-format-18 clean on both ctypes libs. ### Perf parity vs Old-TE - Batched GEMM (MI300X, fp16 rcr, batch=8): at parity / slightly ahead — median gap +4.00%, 100% within ±15%. - Batched contraction (MI350X, fp16 rcr): at parity — median gap -0.95%, 100% within ±15%. ## Scope / known limitations - Batched contraction: `rcr` only, `k_batch==1` only (split-K is a shared Old-TE kernel defect — hard-rejected, never silently-wrong), non-tile-multiple M/N/K rejected by `IsSupportedArguments`. --- Supersedes and closes #9306 (batched GEMM) and #9328 (batched_contraction). **Related PRs (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM), #9000 (grouped), #9028 (stream-K), #8887 (fp8/bf8/int8), #9305 (multi-ABD), #9307 (preshuffle), #9308 (multi-D), #10439 (block-scale quant, 5 ops). --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>

feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts) > Re-opened from #8479 with a compliant branch name (users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16). Supersedes #8479. ## Summary This PR routes the **Tile Engine (TE) regular-GEMM sweep through the Dispatcher**, making the Dispatcher the single source of truth for **codegen → build → runtime** while the Tile Engine keeps only the **config search space** and the **benchmark loop**. It is the consolidated, **single-commit** GEMM bridge covering **all four layouts (`rcr`/`rrr`/`crr`/`ccr`)** and **both `fp16` and `bf16`**. It is a clean re-roll of the earlier bridge work (previously split across #8123 + the stacked key/bf16/layouts/parity/example PRs and consolidated in #8261). Those branches accumulated unrelated cross-project commits through repeated `develop` merges; **this branch is a single clean commit off the latest `develop`** containing only the GEMM-bridge files. It supersedes and replaces #8123 / #8261. ## Motivation The Tile Engine historically owned its own codegen/build/runtime for GEMM (`tile_engine/ops/gemm/gemm_universal/`). The consolidation goal is for the **Dispatcher** to own all of that — exactly as it already does for **FMHA** and **Grouped Conv** — so there is one kernel-generation/build/runtime path and the TE shrinks to a config+benchmark frontend. This PR brings regular GEMM in line with that reference binding. ## The binding (mirrors the FMHA/Conv reference, six stages) 1. **Config JSON (TE side)** — the sweep search space lives in `tile_engine/ops/gemm/configs/` (flat op-root layout, matching the `fmha/` and `grouped_conv/` bridges). 2. **Codegen (Dispatcher)** — `dispatcher/codegen/unified_gemm_codegen.py` emits one fully-typed `.hpp` per kernel; `GemmKernelConfig.name` reproduces `KERNEL_NAME` **byte-for-byte** (the thread tying config → kernel → runtime). 3. **Compile to `.so`** — a single static `gemm_ctypes_lib.cpp` is force-included (`-include <kernel.hpp>`); one `.so` per kernel. 4. **Flat `extern "C"` ABI** — `dispatcher_run_gemm(A, B, C, M, N, K, time_ms)` + the kernel-name enumeration entry points. **Host-pointer** memory model (the C lib `hipMalloc`s internally) — the FMHA-forward branch of the reference. 5. **Python ctypes wrapper** — `dispatcher/python/gemm_utils.py` (`GemmDispatcherLib` + `GpuGemmRunner`). 6. **TE driver (3 phases)** — `gemm_full_benchmark.py` (parallel codegen+build → `expand_sweep` → subprocess-isolated benchmark) + the disposable per-kernel worker `run_one_gemm_kernel.py`. ## What's included **Bridge core** - `dispatcher/codegen/unified_gemm_codegen.py` — GEMM codegen, byte-exact naming. - `dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp` — flat C ABI, host-pointer model. - `dispatcher/python/gemm_utils.py` — `GemmKernelConfig`, multi-kernel build (`setup_multiple_gemm_dispatchers`), `expand_sweep`, one-`.so`-per-kernel. - `tile_engine/ops/gemm/gemm_full_benchmark.py` + `run_one_gemm_kernel.py` — 3-phase, multi-GPU, subprocess-isolated driver/worker. **Feature surface (the point of this PR)** - **All four layouts** `rcr`/`rrr`/`crr`/`ccr` (row-major C only — ck_tile rejects column-major C at build) with layout-aware host transpose. - **`fp16` + `bf16`** (bf16 via uint16 byte-encoding; dtype derived from kernel name). - **Trait-derived registry `KernelKey`** — replaces the earlier hard-coded fp16/rcr key so the registry path generalizes across dtype/layout/tile. **Correctness & performance hygiene** - **`--verify`** opt-in fp32 numpy-reference gate (global `max|out-ref|/max|ref|`), `verified`/`max_rel` columns in the CSV; a mismatch counts as a failure. - **Tile Engine AMDGPU `-mllvm` codegen-flag parity** (without these the kernel builds with different occupancy and the timing diverges) and **arch-validated tile filtering** against the real pipeline/scheduler. - **Multi-GPU** fan-out across all visible GPUs (`--devices`, device-pinned `HIP_VISIBLE_DEVICES` workers). **Example & tests** - `dispatcher/examples/gemm/python/12_te_bridge.py` — runnable end-to-end example. - `dispatcher/tests/test_gemm_parity.py`, `test_gemm_utils.py`, and a parity regression harness. **Cleanup** - Removes the legacy standalone `gemm_universal` build path (`gemm_universal_instance_builder.py`, `*_benchmark*.{py,cpp,hpp}`, `gemm_universal/CMakeLists.txt`) and the old `test/ck_tile/gemm_tile_engine/` harness; promotes the sweep configs to the flat op-root `configs/`. ## Design decisions (consistent with the reference) - **Host-pointer memory ownership** (C lib owns device memory) — matches FMHA-forward; the Python runner passes host numpy arrays straight through. - **One `.so` per kernel** — packaging choice; the multi-kernel name ABI is retained (`get_kernel_name_at(0)` reports the single kernel), so the Python enumeration path is unchanged from FMHA/Conv. - **Flat `configs/`** at the op root — matches the `fmha/`/`grouped_conv/` convention; the not-yet-bridged variants keep their per-variant `configs/` dirs, selected by `--variant`. ## Validation (gfx942 / MI300X) - Bridge build + benchmark + `--verify` across **`fp16` and `bf16`** and **all four layouts**, checked against an fp32 numpy reference (`A @ B`). - **Name parity** holds end-to-end: each `.so`'s reported runtime name equals `GemmKernelConfig(...).name`. - bf16 passes under a widened fp16/bf16 tolerance; fp16 within the standard `max_rel` gate. ## Test plan - [ ] `gemm_full_benchmark.py --verify` over `configs/default_ci_config.json` for `fp16` and `bf16`, each of `rcr`/`rrr`/`crr`/`ccr`. - [ ] `unified_gemm_codegen.py` emits a header whose stem == `GemmKernelConfig.name`. - [ ] `setup_multiple_gemm_dispatchers` builds + links each config against `gemm_ctypes_lib.cpp`. - [ ] `pytest dispatcher/tests/test_gemm_parity.py dispatcher/tests/test_gemm_utils.py`. - [ ] `examples/gemm/python/12_te_bridge.py` runs end to end. ## Notes - Single clean commit off the latest `develop`; the diff is **35 files, all under `projects/composablekernel/`** (dispatcher + tile_engine/ops/gemm + test/ck_tile). - **Supersedes #8123 and #8261**, which will be closed. - Stream-K (#8136) and grouped GEMM are separate bridge efforts, not in this PR. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Muhammed Ozturk <muozturk@ctr2-alola-ctrl-01.amd.com>

ozturkosu and others added 16 commits June 16, 2026 00:25

Potential fix for pull request finding

a741d14

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

1af77db

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

fa7aa1b

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

polish comments

29c8cd5

Fix stale checkout ref in libraries PR bot workflow

b45fadb

Merge branch 'develop' into muozturk/gemm-bridge-all-layouts-bf16

2361d88

ozturkosu requested review from a team as code owners July 1, 2026 05:11

github-actions Bot added documentation project: composablekernel github actions labels Jul 1, 2026

ozturkosu self-assigned this Jul 1, 2026

ozturkosu requested review from ThruptiRajLakshmanaGowda and yraparti July 1, 2026 05:12

ozturkosu mentioned this pull request Jul 1, 2026

feat(ck_tile): TE -> Dispatcher GEMM bridge (all layouts, fp16/bf16) #8479

Closed

5 tasks

ozturkosu mentioned this pull request Jul 1, 2026

feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts) #8998

Merged

3 tasks

ozturkosu changed the title ~~feat(ck_tile): TE -> Dispatcher GEMM bridge (all layouts, fp16/bf16)~~ feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts) Jul 1, 2026

ozturkosu mentioned this pull request Aug 10, 2026

feat(ck-tile): batched GEMM + batched-contraction TE to dispatcher bridges #10612

Merged

Navigation Menu

Uh oh!

feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts) - #8997

feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts)#8997
ozturkosu merged 25 commits into
developfrom
users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16

ozturkosu commented Jul 1, 2026

Uh oh!

therock-pr-bot Bot commented Jul 1, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Jul 1, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

ozturkosu commented Jul 1, 2026

Summary

Motivation

The binding (mirrors the FMHA/Conv reference, six stages)

What's included

Design decisions (consistent with the reference)

Validation (gfx942 / MI300X)

Test plan

Notes

Uh oh!

therock-pr-bot Bot commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

therock-pr-bot Bot commented Jul 1, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 1, 2026 •

edited

Loading