feat(ck-tile): multi-D GEMM TE to dispatcher bridge#9308
Conversation
Extend the bridge to the gemm_multi_d GEMM op so the Dispatcher generates, builds, and launches it at full parity with the legacy Tile Engine version, mirroring gemm_universal (#8997), grouped (#9000), and stream-K (#9028). Capability set matches Old-TE exactly: fp16; layouts rcrr/rrrr/ccrr/crrr (A/B vary, C and D row-major); elementwise ops MultiDAdd/MultiDMultiply/ PassThrough; num D tensors swept (default 1,2). E = op(A@B, D0, D1, ...). - unified_gemm_codegen.py: _multi_d_single_include re-exports NumDTensor/ DsDataType/DsLayout/DLayout/ElementWiseFn/GemmMultiDArgs and ALayout/BLayout/ CLayout plus GEMM_KEY_MULTI_D/NUM_D_TENSORS/ELEMENTWISE_OP/D_LAYOUT macros under CK_TILE_SINGLE_KERNEL_INCLUDE (leverages the existing MULTI_D codegen path: _multi_d_types, _launch_function_multi_d, CShuffle multi-D epilogue). - multi_d_gemm_ctypes_lib.cpp: dedicated C ABI (dispatcher_run_multi_d_gemm + dispatcher_get_num_d_tensors) that force-includes one kernel, allocs A/B/C and each D, derives strides from the compile-time layouts, and calls SelectedKernel::launch(GemmMultiDArgs, ...) directly (registry-bypass). - gemm_utils.py: MultiDGemmProblem/Result, GpuMultiDGemmRunner, run_multi_d; GemmKernelConfig multi_d fields + 4-char codegen_layout + _multid_ name suffix (byte-for-byte name parity); _ctypes_source_name; expand_sweep variant="multi_d". - gemm_multi_d_full_benchmark.py / run_one_gemm_multi_d_kernel.py: 3-phase driver + isolated GPU worker (num-D read off the .so, fp32 verify reference). - MULTI_D_GEMM_BRIDGE.md: bridge documentation. Smoke (gfx942 MI300X): codegen+compile+link 0 failures across all 4 layouts x 2 ops x 2 num_d; end-to-end driver 6/6 kernels VERIFY (max_rel 4.9e-4).
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. |
Address review findings on the gemm_multi_d TileEngine->Dispatcher bridge:
- Add multi_d_config (elementwise_ops x num_d_tensors) to default_config.json
and default_ci_config.json so expand_sweep covers the real capability set
(MultiDAdd + MultiDMultiply, num_d in {1,2}) instead of falling back to
MultiDAdd/num_d=1 only. Schema matches _expand_values ({"values":[...]}).
- Fix default_ci_config.json tile point: the old 64x64x64 / warp_tile_k=64
entry was rejected by the dispatcher validator and expanded to ZERO kernels,
so CI built nothing. Use a validated 128x128x64 / 2x2 / 32x32x16 point that
yields kernels across all four op x num_d combos.
- Trim gemm_multi_d_full_benchmark.py: correct the false fp16+bf16+fp8+bf8
dtype claim to fp16-only, fix copy-paste stream-K docstrings
(streamk_gemm_full_benchmark -> gemm_multi_d_full_benchmark, GemmProblem ->
MultiDGemmProblem), and note MultiDMultiply error amplification in --verify-tol
help.
Verified fp16 x {rcrr,rrrr,ccrr,crrr} x {MultiDAdd,MultiDMultiply} x {num_d 1,2}
on gfx942 (MI300X): 16/16 combos pass, worst max_rel 6.16e-4 << 2e-2 gate.
|
Posting the round-2 review fixes here to keep the description focused on what the PR does. Here is what I addressed from review:
|
A/B perf-parity sweep vs Old-TE (MI300X / gfx942)Ran the serialized bridge-vs-Old-TE performance parity sweep for multi_d GEMM on a single MI300X (gfx942), apples-to-apples:
Result (40 rows = 8 stems × 5 shapes)
Full per-(layout, op, shape) CSV is attached to the PR working set ( |
A/B perf-parity sweep vs Old-TE (MI350X / gfx950)MI350X (CDNA4 / gfx950) re-run of the multi-D GEMM TE-to-dispatcher bridge perf
Result (num_d=1, as-requested reproduction of the MI300X scope)
Important fairness caveat — num_d mismatch. The Old-TE multi_d benchmark bakes Result (num_d=2, byte-identical to Old-TE — the fair comparison)
Verdict: at parity. On the byte-identical num_d=2 comparison the bridge is centered on CSV (num_d=1, requested): |
There was a problem hiding this comment.
Pull request overview
Adds end-to-end Tile Engine → Dispatcher bridge support for the multi-D GEMM variant (gemm_multi_d), enabling dispatcher-driven codegen/build/runtime for GEMM kernels that fuse one or more D tensors into the epilogue.
Changes:
- Introduces a dedicated multi-D ctypes C ABI (
dispatcher_run_multi_d_gemm) and corresponding Python-side runner/config expansion support. - Adds Tile Engine sweep driver + isolated worker for building/benchmarking multi-D kernels (with optional fp32 reference verification).
- Adds CPU-only unit tests and extends TE sweep configs to cover
{elementwise_op} × {num_d_tensors}for multi-D.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_multi_d_kernel.py | New subprocess worker to run/verify one multi-D kernel from a compiled .so. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/configs/default_config.json | Adds multi_d_config sweep dimensions (elementwise op × D count). |
| projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/configs/default_ci_config.json | CI config updated to a validated tile point and adds multi_d_config. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d_full_benchmark.py | New 3-phase multi-D sweep driver (codegen/build → problems → multi-GPU benchmark). |
| projects/composablekernel/dispatcher/tests/test_multi_d_bridge.py | New CPU-only tests for multi-D naming and codegen JSON projection/contracts. |
| projects/composablekernel/dispatcher/python/gemm_utils.py | Adds multi-D config fields, name/layout encoding, build selection, and a multi-D runner. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Exports multi-D types/macros under CK_TILE_SINGLE_KERNEL_INCLUDE for single-include builds. |
| projects/composablekernel/dispatcher/bindings/ctypes/multi_d_gemm_ctypes_lib.cpp | New dedicated ctypes library implementing the multi-D ABI and direct-kernel launch path. |
| projects/composablekernel/dispatcher/bindings/ctypes/MULTI_D_GEMM_BRIDGE.md | New documentation describing the bridge design and ABI. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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>
Resolve conflicts in the dispatcher codegen + gemm_utils with the grouped/ stream-K bridge changes that landed on develop: - unified_gemm_codegen.py: keep the multi_d single-include export block; the global ALayout/BLayout/CLayout aliases develop added are already emitted once earlier in the same block. - gemm_utils.py: canonicalize the regular-GEMM flag on _has_single (dropping _has_gemm), keep _has_grouped + _has_multi_d and both run_grouped / run_multi_d methods, and merge the variant-name / ctypes-source / expand_sweep docstrings so grouped and multi_d coexist. Also address Copilot review feedback: - gemm_utils.run(): raise a descriptive RuntimeError when the .so lacks dispatcher_run_gemm (grouped/multi_d libs) instead of AttributeError. - run_one_gemm_multi_d_kernel.py: cache host A/B/D per (shape, num_d) so batch mode stops regenerating inputs per kernel; fix the input-JSON docstring to list only the consumed fields (num_d from the .so, op from the kernel name). - gemm_multi_d_full_benchmark.py: correct the docstring to state multi_d uses its dedicated dispatcher_run_multi_d_gemm ABI, not the regular one.
|
Merged latest 1. 2. Batch mode regenerates A/B/D per kernel ( 3. Worker docstring lists ignored 4. "same C ABI" wording ( 5. Shebang not on line 1 (both TE scripts) — Already on line 1 in the current revision (addressed in an earlier commit before this review batch); no change needed. 6. Test docstring run path ( CPU-only unit tests still green (10 passed) and both conflict files compile clean after the merge. |
…ault C-ABI timing Old-TE gemm_multi_d_benchmark_single.cpp bakes DsDataType=tuple<D0,D1> (DsDataType::size()==2), so num_d=2 is the byte-identical Old-TE comparison. - MultiDGemmProblem.num_d default 1 -> 2 - expand_sweep multi_d num_d_tensors fallback [1] -> [2] - multi_d_gemm_ctypes_lib.cpp: default flush_cache=true, rotating_count=1000 (matching Old-TE) via new env_bool helper; add CK_TILE_BENCH_FLUSH and CK_TILE_BENCH_ROTATING knobs so the committed benchmark is fair and reproducible without an un-committed patched .so - MULTI_D_GEMM_BRIDGE.md: document num_d=2 as the headline parity slice (correcting the false num_d=1 fair-slice claim), the new timing knobs, and a rocprof TODO for the 14/640 4096^3 mem-pipeline outliers
…bridge_gemm_multi_d # Conflicts: # projects/composablekernel/dispatcher/python/gemm_utils.py
The dispatcher multi_d codegen emitted UsePersistentKernel from the swept 'persistent' trait unconditionally, but Old-TE (gemm_instance_builder.py) wires that flag only for universal/preshuffle/grouped/mx/batched and leaves gemm_multi_d at the template default (false). For '_True' stems this made the bridge kernel genuinely persistent, flipping AccumVGPR 224->384 and spilling 32B to scratch -> ~32% slower on small register-bound shapes (e.g. 1024x512x256: 8860ns vs Old-TE 6695ns); large shapes were unaffected. Gate the flag on the operator so multi_d always emits UsePersistentKernel=false, matching Old-TE; other variants are unchanged. After the fix the multi_d '_True' kernel is byte-identical to Old-TE (AccumVGPR 224, Scratch 0) and A/B parity collapses from ~-34% to -0.13%; a gemm_universal '_True' kernel still emits UsePersistentKernel=true (no regression).
Fix: multi_d bridge no longer forces
|
metric (1024×512×256, _True) |
before | after | Old-TE |
|---|---|---|---|
| AccumVGPR | 384 | 224 | 224 |
| Scratch | 32 B | 0 | 0 |
| GPU kernel duration | 8860 ns | 5492 ns | 5572 ns |
| A/B gap (median-15, TFLOPS) | ~-34% | -0.13% | — |
- Post-fix mangled kernel name is byte-identical to Old-TE.
- Regression check: a
gemm_universal _Truekernel still emitsUsePersistentKernel = true(non-multi_d variants unaffected). - The large-shape 4096³
_Truecase that looked like -11.4% in an earlier sweep was compile-time contention; a clean re-measure is -1.0% (parity).
Net: multi_d bridge is now at kernel-level parity with Old-TE across the shape range, including the small _True stems.
… not default to gfx942
Architecture-default fix (per review feedback)Addressed the feedback "don't default to gfx942; use get_arch and throw an error for unsupported arch types", scoped to this PR's own multi_d files. Changed — #ifndef GFX_ARCH
#error "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=<arch>); do not default to a specific GPU architecture."
#endifThe build harness already passes Why this is the whole change here: the requested Python Verified: |
Whole-config perf parity — multi_d bridge vs Old Tile-Engine (MI300X / gfx942)Ran the comprehensive default-config config-space (not just a shape sweep): Verdict: AT PARITY.
gap = (bridge − old)/old × 100 (+ = bridge faster). Shapes: 512³, 1024³, 2048³, Per-layout medians: rcrr +0.38 %, rrrr +0.21 %, crrr +0.15 %, ccrr +0.03 % (all 100 % Fair comparison: Old-TE multi_d bakes Correctness note (separate from perf): 44/640 rows are non-running kernels CSV + full summary attached. |
| return f"""// Multi-D symbol exports for the single-include ctypes lib. | ||
| using ALayout = {ns_name}::ALayout; | ||
| using BLayout = {ns_name}::BLayout; | ||
| using CLayout = {ns_name}::CLayout; | ||
| using DsDataType = {ns_name}::DsDataType; | ||
| using DsLayout = {ns_name}::DsLayout; | ||
| using DLayout = {ns_name}::DLayout; | ||
| using ElementWiseFn = {ns_name}::ElementWiseFn; | ||
| static constexpr ck_tile::index_t NumDTensor = {ns_name}::NumDTensor; | ||
| using GemmMultiDArgs = {ns_name}::GemmMultiDArgs; | ||
| // Multi-D signature descriptors (consumed by the ctypes lib / KernelKey). | ||
| #define GEMM_KEY_MULTI_D 1 | ||
| #define GEMM_KEY_NUM_D_TENSORS {config.num_d_tensors} | ||
| #define GEMM_KEY_ELEMENTWISE_OP "{config.elementwise_op}" | ||
| #define GEMM_KEY_D_LAYOUT "{config.d_layout}" | ||
| """ |
| from gemm_utils import ( # noqa: E402 | ||
| MultiDGemmProblem, | ||
| GpuMultiDGemmRunner, | ||
| _multi_d_layout_from_kernel_name, | ||
| ) |
MI350 / gfx950 build break — diagnosis (blocks the gfx950 parity run)Building this branch (HEAD Root cause
On this branch those specializations are gated only by SFINAE ( Already fixed on developCurrent FixRebase / merge this branch onto current ControlThe sibling mx_gemm branch, which still carries the older |
…ti_d benchmark (no gfx942 default)
Arch-agnostic fix — no more silent gfx942 default (commit
|
## 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>
gfx950 / MI350 A/B parity + functional equivalence (arch-fixed head ec3dc0e)Built FRESH from PR tip on smci355 (gfx950, ROCm 7.2). This closes the second-arch gap (gfx942/MI300X already posted). VERDICT: AT PARITY. Bridge builds, runs, and is numerically correct on gfx950; A/B timing at parity with Old-TE. Functional equivalence
Performance (80 rows = 16 kernels x {512,1024,2048,1024x512x256,4096}^3, A/B interleaved median-3)
gap = (bridge-old)/old (negative = bridge faster). |gap|>15% outliers — standalone median-5 re-measure: ALL PERSIST, bridge FASTER
Fairness
|
…bridge_gemm_multi_d
Conflict-proofing update (merged develop, made shared regions byte-identical)Merged What changed here (append-only identical-superset strategy):
Why: so that after ANY sibling bridge PR merges to develop, the others do not re-conflict on these shared files. Note (transparency): the shared Validated: |
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 forfused 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.pyexactly: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 ofD tensors (1 and 2). It follows the registry-bypass bridge pattern used by the grouped
(#9000) and stream-K (#9028) bridges.
Test Plan
python3 -m pytest dispatcher/tests/test_multi_d_bridge.py -v(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.
Test Result
max_rel = 6.16e-4(~30x under the 2e-2 gate). Col-majorccrr/crrrhave realon-GPU numeric evidence.
{MultiDAdd, MultiDMultiply} × {num_d 1, 2}per layout.multi_d_gemm_ctypes_lib.cpp.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_dbenchmark 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.