feat(ck-tile): add grouped GEMM variant to TE to dispatcher bridge#9000
Open
ozturkosu wants to merge 19 commits into
Open
feat(ck-tile): add grouped GEMM variant to TE to dispatcher bridge#9000ozturkosu wants to merge 19 commits into
ozturkosu wants to merge 19 commits into
Conversation
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>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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).
…(all layouts) Adds the remaining data types Tile Engine's plain GEMM has MFMA warp tiles for beyond the fp16/bf16 surface of PR #8479: fp8 (E4M3) and bf8 (E5M2) accumulating into fp16, and int8 accumulating into int32 (gfx942). Covers all four A/B layout combinations per dtype (row-major C only, as ck_tile rejects column-major C). Codegen (codegen_common.py, unified_gemm_codegen.py): - add int32 to the CK / qualified / dispatcher dtype maps - get_output_dtype: int8 -> int32 (fp8/bf8 -> fp16 unchanged) - 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, mirroring the existing bf16 helper - GpuGemmRunner.run encodes A/B and sizes the C buffer per dtype (fp16 for fp8/bf8, int32 for int8) - expand_sweep sets dtype_c/dtype_acc from the input dtype Tests: - test_gemm_utils.py: fp8/bf8 codec round-trip, format ranges, NaN/zero slots, saturation, byte size; output-dtype mapping (CPU-only) - test_gemm_parity.py: fp8/bf8/int8 cases with dtype-aware inputs, references and tolerances (int8 exact); GPU-gated like the existing fp16/bf16 cases GPU parity validation deferred to a follow-up run on an MI300X node.
Grouped GEMM variant of the bridge, stacked on the fp8/bf8/int8 bridge (#8887, itself on #8479). The grouped kernel is multi-problem, so it uses a dedicated registry-bypass ctypes ABI; TE only generates configs and benchmarks. - codegen: grouped variant + launch generator (arch_filter.py, unified_gemm_codegen.py), standalone 02_grouped_gemm_driver.cpp, README. - bridge: grouped_gemm_ctypes_lib.cpp (multi-problem ABI), GpuGroupedGemmRunner + dtype/layout codecs (gemm_utils.py), TE driver/worker harness, GROUPED_GEMM_BRIDGE.md. Coverage: {rcr,rrr,ccr,crr} x {fp16,bf16,fp8,bf8}; validated at Old-TE parity on MI300X/gfx942 (64/64 correct, 64/64 within +/-15%).
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
…ge-grouped-gemm Resolve conflicts from #8997 (reviewed fp16/bf16 base bridge) landing on develop while this branch carried its own bridge lineage + fp8/bf8/int8 + grouped-gemm on top. Resolution: - C++ backends/ctypes: took develop's clang-formatted lines (content identical). - codegen: kept grouped/fp8-int8 dtype+layout emission (acc_dtype-aware CDataType/AccDataType, global ALayout re-export); adopted develop's epilogue-specific CShuffle pow2-repeat gate (via _cshuffle_repeat_ok). - gemm_utils.py: kept the grouped + fp8/bf8/int8 superset; adopted develop's pinned hipcc (_resolve_hipcc) and its epilogue-aware pow2 gate. - tests: kept the fp8/bf8/int8 + grouped test superset.
The example is a regular fp16/bf16 GEMM gallery (matrix/shapes/sweep) with nothing grouped-specific -- its only mention of grouped GEMM is a docstring note that grouped rides the same bridge on a separate branch. Unreferenced elsewhere, so removing it is self-contained.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the CK Tile Engine (TE) → Dispatcher “bridge” to support the grouped GEMM variant, where one launch executes a list of independent (M, N, K) sub-problems. It adds dispatcher-side grouped codegen, a dedicated grouped ctypes C ABI (registry-bypass), Python/TE runners, and updates tests/docs/examples to cover additional dtype behavior (fp8/bf8/int8 mappings and codecs).
Changes:
- Add grouped GEMM variant to dispatcher GEMM codegen (naming, arch filtering, grouped launch generation).
- Introduce a grouped GEMM ctypes bridge + Python runner and TE benchmark driver/worker for subprocess-isolated benchmarking.
- Expand tests/docs/examples to cover fp8/bf8 codecs, output-dtype mapping, and grouped usage.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py | New subprocess worker to run grouped kernels and emit JSON results. |
| projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_full_benchmark.py | New 3-phase TE driver to build grouped kernels and benchmark via subprocess batching. |
| projects/composablekernel/dispatcher/tests/test_gemm_utils.py | Add fp8/bf8 codec and output-dtype mapping unit tests. |
| projects/composablekernel/dispatcher/tests/test_gemm_parity.py | Extend parity tests across fp16/bf16/fp8/bf8/int8 and correct output/acc dtype handling. |
| projects/composablekernel/dispatcher/python/gemm_utils.py | Add grouped problem/result/runner, grouped ABI wiring, fp8/bf8 codecs, and output dtype handling. |
| projects/composablekernel/dispatcher/examples/gemm/python/12_te_bridge.py | New Python gallery example for the regular GEMM bridge. |
| projects/composablekernel/dispatcher/examples/gemm/cpp/README.md | Document the new grouped GEMM standalone driver example. |
| projects/composablekernel/dispatcher/examples/gemm/cpp/02_grouped_gemm_driver.cpp | New standalone grouped GEMM driver that launches SelectedKernel::launch(descs, stream) directly. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Add GemmVariant.GROUPED and grouped launch generation. |
| projects/composablekernel/dispatcher/codegen/README.md | Document the grouped GEMM variant and usage. |
| projects/composablekernel/dispatcher/codegen/codegen_common.py | Add int32 mapping + output/acc dtype mapping helpers. |
| projects/composablekernel/dispatcher/codegen/arch_filter.py | Add GEMM_GROUPED operator constraints for arch filtering. |
| projects/composablekernel/dispatcher/bindings/README.md | Document the new grouped GEMM ctypes binding. |
| projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp | New grouped GEMM ctypes C ABI that bypasses the registry. |
| projects/composablekernel/dispatcher/bindings/ctypes/GROUPED_GEMM_BRIDGE.md | Design doc for the grouped GEMM bridge. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| // Workspace allocated internally (dispatcher idiom, mirrors grouped_conv stream-K). | ||
| const std::size_t ws_size = kargs.size() * sizeof(ck_tile::GemmTransKernelArg<>); | ||
| DeviceMem workspace_dev(ws_size); |
Comment on lines
+33
to
+37
| #include <exception> | ||
| #include <string> | ||
| #include <type_traits> | ||
| #include <vector> | ||
|
|
Comment on lines
+615
to
+635
| 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.""" | ||
| 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) | ||
| return table |
Run clang-format 18.1.4 (projects/composablekernel/.clang-format) over the grouped-GEMM bridge sources so the CK Jenkins Static checks stage passes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.g_dispatcher->run/GemmHostArgs) cannot express a list of problems.generated_tile_backend.hpp::run()) hard-codes the single-problem launch and won't compile against a groupedSelectedKernel.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-timeKERNEL_NAMEmacro.Changes
Codegen (absorbed from #8075)
codegen/arch_filter.py—GEMM_GROUPEDoperator tile constraints.codegen/unified_gemm_codegen.py—GemmVariant.GROUPED, the grouped launch generator (DeviceMem internal workspace viaMakeKargs, persistent/non-persistent grid),groupedin--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-timeALayout/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)
C is always row-major.
int8(rejected by the TE grouped builder) andfp32/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-
.soguard, matched compile flags):rocprofto be a measurement-harness artifact (Old-TE's JSONlatency(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
develop. Depends on feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts) #8997 (fp16/bf16 bridge) and feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts) #8998 (fp8/bf8/int8 bridge) merging todevelopfirst; until then this PR's diff also shows their content, after which it reduces to the grouped-only files.