feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts)#8998
feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts)#8998ozturkosu wants to merge 20 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.
✅ 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. |
ozturkosu
changed the title
…ffle epilogue unified_gemm_codegen forced CShuffleEpilogueProblem trailing template args (false, 1, 1, DoubleSmemBuffer) for the gemm_universal cshuffle epilogue, while Old-TE's gemm_universal_instance_builder stops at NumWaveGroups (letting the epilogue defaults apply). For RowMajor-A those forced values equal the defaults (parity), but for ColMajor-A and 4x1x1 block-maps they yield a higher-VGPR kernel (120/128 vs Old-TE 92/100) -> lower occupancy -> 30-75% slower. Drop the forced args so the bridge emits the same epilogue as Old-TE. On MI300X all 18 affected fp8 stems recover: 50/51 formerly >15% (M/N/K sweep) rows now within +/-15% (median 0.01%). multi_d variant left unchanged.
Fix: fp8 ColMajor-A /
|
bf8 / all-layout revalidation of the cshuffle-epilogue fix (
|
Revalidation results — all
|
| median | mean | worst | rows >15% |
|
|---|---|---|---|---|
| Before (previous 4-GPU sharded sweep) | −1.72% | −5.61% | −83.7% | 198 / 556 |
| After (fixed + median-of-3, isolated) | −0.16% | −1.00% | −12.9% | 0 / 556 |
By layout (after):
| dtype | layout | rows | median | worst | >15% |
|---|---|---|---|---|---|
| bf8 | ccr | 132 | −0.06% | −12.9% | 0 |
| bf8 | crr | 76 | +0.01% | +6.9% | 0 |
| bf8 | rcr | 128 | −0.46% | −10.7% | 0 |
| bf8 | rrr | 216 | −0.29% | −11.7% | 0 |
| fp8 | ccr | 4 | −0.83% | −3.1% | 0 |
The residual spread (worst −12.9%) is on the smallest/most-noise-prone shapes and is well within the ±15% bar. Merging these remeasured rows back into the full sweep yields 0 / 11,592 baseline rows >15% (overall median −0.35%).
Correctness is unchanged from baseline: bridge_verified and max_rel (bf8 ≈ 0.07, fp8 ≈ 0.04) are identical to the pre-fix sweep — the expected lossy-fp8/bf8 tolerance and the known verify-harness artifact on non-square 1024x512x256 ColMajor-A (numpy reference transpose), not a kernel change.
Conclusion: the cshuffle-epilogue fix resolves the real ColMajor-A / 4x1x1 regression across all layouts and both fp8 and bf8; the remaining large gaps in the earlier sweep were 4-GPU measurement contention, confirmed collapsed under isolated remeasurement.
…figs
The bridge moved the gemm sweep config JSONs from gemm_universal/configs/ up
to the flat op-root gemm/configs/ (single source of truth shared with the
dispatcher bridge). Old-TE's gemm_universal/CMakeLists.txt still read from
gemm_universal/configs/, so CMake configure aborted with "Config file not
found: .../gemm_universal/configs/default_ci_config.json" (fires even under
EXCLUDE_FROM_ALL, breaking the whole build).
Add GEMM_UNIVERSAL_CONFIG_DIR (= ../configs) and repoint the three config
path constructions at it. Old-TE and the bridge now read the same
gemm/configs/{default_ci_config,default_config,user_provided_config}.json,
keeping Old-TE alive as the baseline until the bridge covers all layouts/dtypes.
Resolve conflicts after PR #8997 (fp16/bf16 TE->Dispatcher bridge) merged: - gemm_ctypes_lib.cpp, backends/{tile_backend,generated_tile_backend}.hpp: take develop (clang-format-only churn; verified whitespace-identical). - tests/test_gemm_{parity,utils}.py: keep HEAD (fp8/bf8/int8 supersets; develop's fp16/bf16 coverage is subsumed). - gemm_utils.py: keep the fp8/bf8 FNUZ codecs, run() buffer branching and generic dtype_c/dtype_acc; adopt develop's pinned _resolve_hipcc() and the epilogue-specific pow2 CShuffle gate; drop now-unused os/shutil imports. - unified_gemm_codegen.py: keep the dtype-generic C/Acc datatype emission; adopt develop's epilogue-specific pow2 gate (_cshuffle_repeat_ok). Also remove the pre-rename duplicate example 12_te_bridge.py; the canonical tile_engine_dispatcher_bridge.py from #8997 supersedes it.
There was a problem hiding this comment.
Pull request overview
This PR extends the existing Tile Engine ↔ Dispatcher GEMM bridge surface beyond fp16/bf16 to support fp8 (E4M3), bf8 (E5M2), and int8, including deriving output/accumulator dtypes (fp8/bf8 → fp16 output, fp32 accumulate; int8 → int32 output/accumulate) and updating the host runner + parity tests accordingly.
Changes:
- Extend codegen type maps and key emission to support
int32C/accumulator types and dtype-deriveddtype_c/dtype_acc. - Add host-side fp8/bf8 FNUZ uint8 codecs and make the Python runner allocate/encode buffers based on parsed kernel dtype and derived output dtype.
- Expand unit + parity coverage to include fp8/bf8/int8 cases with dtype-aware references and tolerances.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/tests/test_gemm_utils.py | Adds CPU-only fp8/bf8 codec tests and output-dtype mapping tests. |
| projects/composablekernel/dispatcher/tests/test_gemm_parity.py | Extends parity matrix to fp8/bf8/int8 and updates reference generation + tolerances. |
| projects/composablekernel/dispatcher/python/gemm_utils.py | Implements fp8/bf8 FNUZ codecs; makes runner/expand_sweep dtype-aware for C/acc types. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Emits AccDataType/CDataType and GEMM_KEY dtype macros based on derived output/acc dtypes. |
| projects/composablekernel/dispatcher/codegen/codegen_common.py | Adds int32 to dtype maps; adds get_acc_dtype; updates get_output_dtype for int8→int32. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Resolve conflicts after PR #8997 (fp16/bf16 all-layout bridge) advanced develop. Land Stream-K as an addition on top of develop's current bridge; KEEP gemm_universal (dropped this branch's "retire legacy gemm_universal" change, since #8997/#8998 rely on it as the Old-TE parity baseline). - gemm_universal/* + test/ck_tile/gemm_tile_engine/* + gemm/CMakeLists.txt: restored from develop (un-retire). - gemm_ctypes_lib.cpp, backends/{tile_backend,generated_tile_backend}.hpp, gemm/README.md, gemm_full_benchmark.py, run_one_gemm_kernel.py: take develop (whitespace/comment-only or superseded; Stream-K ships its own driver/worker). - gemm_utils.py: develop base + re-applied the Stream-K delta (reduction_strategy field, _ctypes_source_name -> streamk_gemm_ctypes_lib.cpp, per-variant compile flags matching TE gemm_streamk, static-lib-skip link, expand_sweep variant/ reduction threading). Regular path keeps develop's _tile_engine_codegen_flags(). - unified_gemm_codegen.py: Stream-K codegen (STREAM_K variant, atomic/linear/tree backend, _launch_function_streamk) auto-merged; resolved 2 hunks -- emit both the ALayout/BLayout/CLayout re-exports (Stream-K lib) AND develop's GEMM_KEY_* registry-key macros (regular lib); adopt develop's epilogue-specific pow2 gate. Static validation: no conflict markers, all touched Python compiles, Stream-K codegen smoke emits a *_streamk.hpp, and GemmKernelConfig.name == codegen stem (name parity, incl. linear/tree suffixes). NOTE: static resolution only -- requires a GPU build + Stream-K A/B parity re-run on gfx942/MI300X to validate before merge.
…dtype - _fnuz_decode_table: memoize with functools.lru_cache (pure function of exp_bits/mant_bits) and mark the returned array read-only, so fp8/bf8 encode/decode no longer rebuild the 256-entry table on every call. - GpuGemmRunner.run: the C host buffer dtype map now raises on an unrecognized out_dtype instead of silently defaulting to fp16 (a wrong element size would corrupt results across the C ABI).
… assumption
GeneratedTileKernelInstance::run() hard-coded GemmHostArgs strides as
{K, K, N}, which is only correct for rcr (RowMajor-A / ColMajor-B / RowMajor-C)
and for square problems (M==N==K). For any non-square rrr/ccr/crr launch the
leading dimensions were wrong, producing grossly incorrect results (GPU output
max_rel ~1.5 vs reference; square shapes masked it because all three strides
coincide).
Derive each stride from the kernel's real layout instead:
stride_A = RowMajor-A ? K : M
stride_B = RowMajor-B ? N : K
stride_C = RowMajor-C ? N : M
Verified on gfx942/MI300X: non-square rrr/ccr/crr now match the NumPy reference
for every dtype (fp16 ~4e-4, int8 bit-exact 0.0) across square and rectangular
shapes; rcr and square were already correct and remain so. Dtype-independent
fix. (Non-tile-aligned "awkward" shapes still hit a separate, dtype-independent
launch-support limitation, tracked separately.)
int8 support — validity check (gfx942/MI300X) + a stride bug found & fixedOld-TE
|
| shape | int8 | fp16 (control) |
|---|---|---|
| square 512³ (all layouts) | ✅ bit-exact (0.0) | ✅ 4.8e-4 |
| rectangular 1024×512×256 (all layouts) | ✅ bit-exact (0.0) | ✅ ~4e-4 |
int8 is bit-exact wherever fp16/bf16 are correct.
Bug found while validating: non-square rrr/ccr/crr gave wrong results (all dtypes)
Initially every non-square rrr/ccr/crr case returned garbage (max_rel ~1.5) for all dtypes (fp16/bf16/int8 alike). Root cause: GeneratedTileKernelInstance::run() in generated_tile_backend.hpp hard-coded the GemmHostArgs leading dimensions as {K, K, N}, which is only right for rcr and for square problems (M==N==K makes the three strides coincide, which is why it was never caught). Fixed to derive each stride from the kernel's real layout:
stride_A = RowMajor-A ? K : M
stride_B = RowMajor-B ? N : K
stride_C = RowMajor-C ? N : MAfter the fix (commit in this PR), non-square rrr/ccr/crr match the reference for every dtype (fp16 ~4e-4, int8 bit-exact) across square and rectangular shapes; rcr/square were already correct and stay correct. This is a dtype-independent correctness fix.
Known, separate caveat
Non-tile-aligned "awkward" shapes (e.g. 257×129×512) still fail to launch (status -1) — but this reproduces identically for fp16, so it's a separate, dtype-independent kernel launch-support limitation for non-aligned dims, not an int8 issue. Tracked separately.
Bottom line: int8 works properly (bit-exact, on par with fp16/bf16). The one real bug surfaced during validation (rcr-only hard-coded strides) is fixed in this PR and benefits all layouts.
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:
All four A/B layout combinations per dtype (row-major C only, matching #8479).
fp32/fp64are intentionally excluded — they appear in TE's dtype-string map but have no MFMA warp tiles inGEMM_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 targetsdevelop, so until #8997 merges its diff also includes the base bridge changes; please merge #8997 first.Changes
codegen_common.py,unified_gemm_codegen.py): addint32to the dtype maps;get_output_dtypeint8→int32; newget_acc_dtype(int8→int32, else fp32); deriveAccDataType/CDataType, theGEMM_KEY_DTYPE_{C,ACC}macros, and the registrydtype_c/dtype_accfrom the dtype instead of hard-codingfloat/fp32.gemm_utils.py): fp8/bf8 FNUZ (gfx942) uint8 codecs — exact decode (matches devicefp8_t/bf8_t), nearest-representable saturating encode (same pattern as the existing bf16 helper);GpuGemmRunner.runencodes A/B and sizes the C buffer per dtype;expand_sweepsetsdtype_c/dtype_acc.test_gemm_utils.pyadds CPU-only fp8/bf8 codec + output-dtype tests (all green);test_gemm_parity.pyadds 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).ADataType/CDataType/AccDataTypeandGEMM_KEY_*macros are correct (int8→int32_t acc/C; fp8→fp16_t C).test_gemm_parity.pycollects 60 cases and skips cleanly without a GPU.test_examples_integration/test_grouped_conv_codegen/test_library_cachingare pre-existing (verified identical on the base branch; they require a built dispatcher.a/ GPU).Test plan
develop.python3 tests/test_gemm_parity.pyand confirm fp8/bf8/int8 parity; tune the fp8/bf8 tolerances if needed (current values are first-cut headroom).