feat(ck-tile): preshuffle GEMM TE to dispatcher bridge#9307
Open
ozturkosu wants to merge 12 commits into
Open
Conversation
Extend the bridge so the Dispatcher generates and launches the weight- preshuffled GEMM (gemm_preshuffle) at parity with the legacy Tile Engine op, for all four supported dtypes (fp16/bf16/fp8/bf8), rcr layout only. - unified_gemm_codegen.py: complete the PRESHUFFLE variant branch so it works for the bridge single-config path (accept a pinned preshufflev2 pipeline, sweep the epilogue, de-dup on the sweep anchor); emit a PermuteN kernel-struct member and thread permute_n through key_name / KernelNaming; export ALayout/BLayout/CLayout globally for single-kernel consumers; fix the wrapper double_buffer flag for preshufflev2. - gemm_ctypes_lib.cpp: preshuffle host B before upload using the same ck_tile::shuffle_b transform Old-TE applies, built from the force- included kernel geometry (BridgePreshuffleConfig) and its BLayout, so the device reads byte-for-byte identical packed weights. B_host stays logical so the numpy reference (A @ B) remains valid. - gemm_utils.py: add permute_n to GemmKernelConfig, thread variant into expand_sweep, and pin permute_n=False (the emitted WeightPreshufflePipelineAGmemBGmemCRegV2 matches the non-permuteN shuffle; GPU-verified). - ctypes_utils.py: validate preshuffle configs against the preshuffle- specific warp-tile whitelist and key the warp-tile tables on the accumulator dtype (fixes fp8/bf8 admitting tiles the codegen rejects). - gemm_full_benchmark.py: map gemm_preshuffle to the preshuffle codegen variant, allow fp8/bf8 + rcr-only, and pass the variant to expand_sweep. GPU smoke on MI300X (gfx942): codegen + .so build succeed for all four dtypes; output verified vs fp32 numpy reference at M=N=K=512 (max_rel fp16 5e-4, bf16 5.5e-3, fp8 3.4e-2, bf8 6.5e-2).
ozturkosu
changed the title
✅ 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. |
… permute_n - Cache the host B-preshuffle per (B pointer, K, N) in the ctypes lib so the reorder runs once per callable instead of on every dispatcher_run_gemm call. The kernel-timed region (g_dispatcher->run) never contained the shuffle, but this keeps the repeated per-iteration path launch-only so the upcoming A/B perf sweep is apples-to-apples with Old-TE (which shuffles B once). - Make permute_n a single source of truth (BRIDGE_PERMUTE_N in gemm_utils); the pin now reads that constant instead of a bare hardcoded False, with a note that both default_config.json and default_ci_config.json ship permute_n=true and must not be honored (host-marker only, no distinct bridged device kernel). Mirror the note at the codegen read site. Correctness re-verified on GPU (rcr, M=N=K=512): fp16 5.15e-04, bf16 5.53e-03, fp8 3.37e-02, bf8 6.54e-02 -- byte-identical to pre-hoist results.
Contributor
Author
A/B perf-parity sweep vs Old-TE (MI300X / gfx942)Ran the full serialized bridge-vs-Old-TE performance parity sweep for preshuffle GEMM on a single MI300X (gfx942), apples-to-apples, across all four dtypes:
Result (80 rows = 16 stems × 5 shapes)
Per-dtype median gap: fp16 +2.63%, bf16 +2.92%, fp8 +2.98%, bf8 +3.30% (20 rows each).
Full per-(dtype, stem, shape) CSV is attached to the PR working set ( |
This was referenced Jul 13, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Adds TileEngine → Dispatcher bridging support for weight-preshuffled GEMM (gemm_preshuffle) so the Dispatcher can codegen/name/run preshuffle kernels consistently with the existing GEMM bridge infrastructure, including host-side B reordering for preshuffle correctness.
Changes:
- Adds CPU-only preshuffle bridge tests to lock down kernel naming and codegen JSON (
permute_n) projection. - Extends GEMM bridge config/name plumbing to carry a preshuffle variant (including
permute_n) through Python config objects and unified codegen naming. - Updates the ctypes GEMM runtime to apply the host-side B preshuffle transform when compiling/running a preshuffle kernel.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py | Threads a variant token into expand_sweep and adds preshuffle dtype/layout guards. |
| projects/composablekernel/dispatcher/tests/test_preshuffle_bridge.py | New CPU-only unit tests for preshuffle naming + codegen JSON + shipped config validation. |
| projects/composablekernel/dispatcher/python/gemm_utils.py | Adds permute_n to GemmKernelConfig, pins bridge permute behavior, threads variant through sweep expansion. |
| projects/composablekernel/dispatcher/python/ctypes_utils.py | Fixes warp-tile whitelist lookup to key by accumulator dtype and adds preshuffle-specific whitelist selection. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Completes preshuffle variant naming/config generation (incl. permute_n) and fixes preshuffle double-buffer flag emission. |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp | Adds host-side B preshuffle (and caching) for preshuffle kernels before H2D upload. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
Author
A/B perf-parity sweep vs Old-TE (MI350X / gfx950)MI350X (CDNA4 / gfx950) re-run of the preshuffle (weight-preshuffle) GEMM
Result (fp16 + bf16, default epilogue -- the buildable matched set)
Verdict: at parity. On the byte-identical matched set the bridge is centered Correctness / coverage exclusions (NOT perf gaps)Unlike the MI300X run (80 rows, all 4 dtypes, both epilogues), on gfx950 with
Net timable, byte-identical, correctness-gated set on this node: fp16/bf16 x CSV: |
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>
…bridge_gemm_preshuffle # Conflicts: # projects/composablekernel/dispatcher/python/gemm_utils.py
…out guard, dtype key)
Address reviewer correctness items on the preshuffle GEMM TE->Dispatcher bridge:
1. get_shuffled_b stale-read foot-gun: the shuffle cache is keyed on
(B pointer, K, N) and cannot detect in-place mutation of the bytes behind an
unchanged pointer. Document a loud immutability contract at the public
dispatcher_run_gemm boundary AND add an opt-in recompute path via the
CK_DISPATCHER_PRESHUFFLE_NO_CACHE env var so a correctness caller that mutates
B is never served a stale shuffle.
2. BLayout byte-identity assumption: the host B-shuffle only reproduces Old-TE's
b_k_n byte-for-byte when BLayout is column-major. Add a static_assert(!kBRowMajor)
in preshuffle_host_b so a future row-major-B layout expansion fails loudly at
compile time instead of silently mis-shuffling B.
3. dtype_key mixed-A/B latent bug: validate_kernel_config lives on the shared
standard path and keyed warp-tile lookup on dtype_a repeated. Thread dtype_b
through so a future mixed-A/B-dtype caller (e.g. fp8_bf8) forms the correct
{dtype_a}_{dtype_b}_{dtype_acc} key instead of silently hitting the wrong one.
No behaviour change for preshuffle scope (dtype_a == dtype_b).
Old-TE ships persistent=[true,false] for the preshuffle op, so both are legitimate configs the bridge must support. On gfx942 both verify against an fp32 reference at rel ~4e-4. Add CPU regression tests that pin the host-launch invariant guarding correctness of the non-persistent path: a non-persistent kernel must use the full one-block-per-tile GridSize (not the occupancy-sized MaxOccupancyGridSize used by persistent kernels), and UsePersistentKernel must track the trait so the host grid and the device entry point agree. Also assert the shipped CI/default configs keep persistent=False in the sweep so the non-persistent path stays exercised.
…bridge_gemm_preshuffle # Conflicts: # projects/composablekernel/dispatcher/python/gemm_utils.py
ozturkosu
requested review from
Copilot and
yraparti
and removed request for
yraparti
July 21, 2026 02:35
Contributor
Author
Whole-config perf parity — preshuffle GEMM bridge vs Old Tile-Engine (MI300X / gfx942)Extends the earlier fixed-CI sweep to the full Verdict: AT PARITY.
gap = (bridge − old)/old × 100 (+ = bridge faster). Shapes: 512³, 1024³, 2048³, 1024×512×256, 4096³. permute_n scope: Coverage findings (not perf gaps, flagged separately):
Fairness: same warmup=50/repeat=100/flush/rotating both sides; Old-TE via its own gemm_preshuffle CMake (real TE |
…_n guard, cache cleanup) Address Copilot + internal review on the preshuffle GEMM bridge: - Copilot / F6: preshuffle_host_b() copied from the caller's raw byte buffer (b_host comes from Python/ctypes, e.g. fp8/bf8 as numpy uint8) with std::copy, which does typed reads through the pointer -- strict-aliasing / object-lifetime UB. Replaced with a byte-wise std::memcpy of the exact K*N*sizeof(T) span into the tensor's flat storage (byte-identical, no typed reads). - Copilot / F5: a preshuffle config with permute_n=True would build a "_permuteN" kernel whose device pipeline is not bridged yet (mis-shuffles B -> wrong results). setup_multiple_gemm_dispatchers now hard-fails on such a config before codegen/compile. expand_sweep still pins permute_n to BRIDGE_PERMUTE_N, so the normal flow (and the shipped configs that carry permute_n=true as a TE host-marker) is unaffected; name/JSON projection is untouched so the existing permute_n name/codegen tests still pass. Added a CPU test. - F3: dispatcher_cleanup() now releases the process-lifetime shuffled-B cache so a library consumer gets its memory back. Verified: 13/13 CPU tests pass; clang-format-18 + py_compile clean; GPU build+run of a preshuffle fp16/rcr kernel (gfx942, ROCm 7.2) matches the numpy reference (max_rel=0.000).
Contributor
Author
Review pass: Copilot threads + internal review (commit
|
| # | Item | Fix |
|---|---|---|
| F6 (Copilot) | preshuffle_host_b() used std::copy, doing typed reads through the caller's raw byte buffer (fp8/bf8 numpy uint8) → strict-aliasing/lifetime UB |
Byte-wise std::memcpy of the exact K*N*sizeof(T) span into the tensor's flat storage — byte-identical, no typed reads |
| F5 (Copilot) | a permute_n=True config would build a _permuteN kernel whose device pipeline isn't bridged (mis-shuffles B → wrong results) |
Hard-fail in setup_multiple_gemm_dispatchers (build path) before codegen/compile; expand_sweep still pins permute_n=BRIDGE_PERMUTE_N so shipped configs / normal flow are unaffected; name/JSON untouched (tests still pass). Added test_permute_n_build_hard_fails |
| F3 | process-lifetime shuffled-B cache never released | dispatcher_cleanup() now resets g_shuffled_b_cache |
Reviewed, intentionally not changed (with rationale)
| # | Item | Disposition |
|---|---|---|
| F1 | #define GFX_ARCH "gfx942" default in gemm_ctypes_lib.cpp |
Shared file — changing to #error affects every sibling bridge; tracked for the central arch-cleanup PR, not this one (same as F7 in the codegen) |
| F2 | shuffled-B cache keyed on (ptr,K,N) can serve stale bytes if B is mutated in place |
Kept: documented at 3 sites + CK_DISPATCHER_PRESHUFFLE_NO_CACHE=1 opt-out; benchmark-scoped. Making no-cache the default would break the A/B-parity purpose |
| F4 | preshuffle_cache_disabled() env parse (v[0] != '0') is looser than the doc |
Kept: documented as "set to 1"; cosmetic |
Verification
- 13/13 CPU unit tests pass (incl. the new build-guard test);
clang-format-18+py_compileclean. - GPU build+run (gfx942, ROCm 7.2): a preshuffle fp16/rcr kernel built with the
memcpychange matches the numpy reference (max_rel=0.000) — the shuffle path is unchanged in behavior, only the UB is removed.
…s (no gfx942 default)
Contributor
Author
Arch-agnostic fix — no more silent gfx942 default (commit
|
ozturkosu
added a commit
that referenced
this pull request
Jul 22, 2026
## 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>
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.
ISSUE ID: #8997
Motivation
The TileEngine → Dispatcher bridge had no path for the weight-preshuffled GEMM op
gemm_preshuffle, which pre-permutes the B (weight) tensor into the pipeline's
packed layout for higher throughput on weight-heavy GEMMs. This is a real Old-TE
capability with no dispatcher equivalent, so this PR adds the bridge so the dispatcher
can generate and launch preshuffle GEMM at parity with the legacy Tile Engine version.
The capability set matches the Old-TE op exactly:
fp16,bf16,fp8,bf8,rcrlayout only (A row-major, B col-major, C row-major), with the B tensor preshuffled
into the packed layout. It follows the same bridge scheme as #8997 (regular GEMM),
#9000 (grouped), and #9028 (stream-K).
Test Plan
python3 -m pytest dispatcher/tests/test_preshuffle_bridge.py -v.sofor all 4 dtypes and verify output vs an fp32 NumPyreference at M=N=K=512 on gfx942 / MI300X.
ck_tile::shuffle_bbyte-for-byte.Test Result
.sobuild succeed for all 4 dtypes (0 codegen failures). Output verifiedvs fp32 reference at M=N=K=512:
ck_tile::shuffle_b; the permuteN shuffle gives ~1.25 rel error and is deliberatelynot used, since the bridge does not emit the permuteN pipeline.
-style=file.stems × 5 shapes = 80 rows, interleaved, fair 50/100/flush/rotating both sides):
at parity / slightly ahead — median gap +3.03%, mean +2.24%, 100% within
±15%, 97.5% within ±5% (range [-3.79%, +5.24%]; per-dtype medians fp16 +2.63%,
bf16 +2.92%, fp8 +2.98%, bf8 +3.30%). 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.