Repository navigation
feat(ck-tile): multi-ABD GEMM TE to dispatcher bridge - #9305
Merged
Merged
Conversation
Bridge the multi-tensor gemm_multi_abd op (fp16, rcrr layout, arrays of A/B/D tensors with per-group element-wise ops) so the Dispatcher generates and launches it at parity with the legacy Tile Engine version. Mirrors the gemm_universal (#8997) bridge, using the grouped/divergent-ABI pattern (#9000) since multi_abd needs array-of-pointer arguments. - unified_gemm_codegen.py: new GemmVariant.MULTI_ABD (tuple As/Bs/Ds types and layouts, three element-wise fns, GemmKernelMultiABD launch, CShuffle epilogue); tensor counts + all three ops folded into key_name()/kernel name; multi_abd_config search block; global tensor-count re-exports for the registry-bypass ctypes lib. - gemm_multi_abd_ctypes_lib.cpp: new registry-bypass C ABI (dispatcher_run_multi_abd) that calls SelectedKernel::launch directly on the force-included kernel with GemmMultiABDHostArgs; owns GPU memory. - gemm_utils.py: multi_abd fields on GemmKernelConfig (byte-exact name), 4-char layout handling, multi_abd codegen json, variant-aware build (own ctypes source, no static lib), MultiABDDispatcherLib + GpuMultiABDRunner, expand_sweep variant support. - gemm_full_benchmark.py / run_one_gemm_kernel.py: gemm_multi_abd variant wired into the 3-phase driver and routed to GpuMultiABDRunner by name. - ctypes CMakeLists: dispatcher_gemm_multi_abd_lib target. Validated on gfx942 (MI300X): fp16/rcrr default_config kernel codegen (0 failures), ctypes .so compiles + links, all C-ABI symbols exported; standard/multi_d/multi_abd codegen regression clean; name parity between GemmKernelConfig.name and the generated header confirmed byte-for-byte.
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. |
…s, valid CI tiles
Round-2 fixes for the gemm_multi_abd TileEngine->Dispatcher bridge:
- B1: verify the BRIDGE launch path numerically. GpuMultiABDRunner now owns a
numpy reference E = CDE( AB(As) @ BB(Bs), {Ds} ) that mirrors
ck_tile::reference_gemm_multiple_abd byte-for-byte (A/B groups combined
element-wise into one matrix, single GEMM, then the CDE op folds the D
tensors) and reports max_rel on the result; run_one_gemm_kernel routes
--verify through it. This also uncovered + fixed a real launch bug: the
GemmMultiABDKernel does NOT derive strides from a 0 sentinel, so the runner
now passes explicit leading strides for A/B/D/E (get_default_stride
semantics) instead of 0, which had collapsed the whole output onto row 0.
- B2: drive the element-wise ops + tensor counts. expand_sweep reads a
multi_abd_config block (lists) and iterates the op/count combos, and
gemm_full_benchmark exposes --multi-abd-{a,b,cde}-op / --multi-abd-num-{a,b,d}
CLI overrides, so non-PassThrough kernels are actually generated and verified.
- B3: default_ci_config.json warp tile 32x32x64 -> 32x32x16 (32x32x64 is not in
the gfx942 fp16 XDL allow-list; CI now yields 16 kernels, was 0).
- N1: comment the fp16_fp16_fp16 vs fp16_fp16_fp32 arch-filter key mismatch in
ctypes_utils as a latent bridge-wide parity narrowing.
- N2: derive layout + per-group ops from the config object (passed via the
worker payload) instead of parsing the kernel name with a silent rcrr default.
- Ensure build/examples exists before hipcc -o so the self-contained multi_abd
.so builds without a prior full CMake configure.
dtype x layout unchanged: exactly fp16 x rcrr (Old-TE parity set).
Validated on gfx942 (MI300X, GPU0): default 2/2/2 PassThrough max_rel=2.9e-4,
CDE=MultiDAdd max_rel=5.7e-4, A-op=MultiDAdd max_rel=4.1e-4 (all << fp16 tol);
CI config yields 16 kernels; standard-variant expand_sweep regression clean.
Contributor
Author
|
Posting the round-2 review fixes here to keep the description focused on what the PR does. Here is what I addressed from review:
|
ozturkosu
marked this pull request as ready for review
July 12, 2026 22:34
Contributor
Author
A/B perf-parity sweep vs Old-TE (MI300X / gfx942)Ran the full serialized bridge-vs-Old-TE performance parity sweep for multi_abd on a single MI300X (gfx942), apples-to-apples:
Result (80 rows = 16 stems × 5 shapes)
Full per-(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
This PR extends the CK Tile “TileEngine → Dispatcher bridge” to support the multi-tensor gemm_multi_abd variant (arrays of A/B/D tensors + fused epilogue), aiming for parity with the legacy Tile Engine (Old-TE) capability without requiring C++ changes from Python callers.
Changes:
- Adds a new registry-bypass, divergent C ABI (
dispatcher_run_multi_abd) and a Python runner (GpuMultiABDRunner) to launch compiled multi-ABD kernels via ctypes. - Extends unified GEMM codegen and TE benchmark plumbing to generate/build/run multi-ABD kernels, including kernel-name contract enforcement.
- Adds CPU-only unit tests to lock in the multi-ABD kernel naming and config JSON projection invariants.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Adds multi_abd variant support to codegen (naming, types, direct-launch generator). |
| projects/composablekernel/dispatcher/python/gemm_utils.py | Extends config model/naming/JSON projection; adds multi-ABD ctypes wrapper + GPU runner; integrates variant-aware build. |
| projects/composablekernel/dispatcher/python/ctypes_utils.py | Adds notes around arch-filter validation interactions (no functional bridge change here). |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp | New ctypes C ABI shim implementing dispatcher_run_multi_abd (device alloc/copy/launch/copyback). |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | Adds a CMake target intended to build the multi-ABD ctypes library in CI/CMake flows. |
| projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py | New CPU-only tests for multi-ABD naming + codegen JSON + shipped config validity. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py | Wires benchmark driver to select multi-ABD variant, propagate layout/ops/counts to worker, adds CLI overrides. |
| projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py | Routes multi-ABD kernels to GpuMultiABDRunner and uses runner-driven verification. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json | Adjusts CI warp tile K to an arch-valid value for gfx942 fp16 XDL. |
💡 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 multi-ABD GEMM TE-to-dispatcher bridge perf
Result (num_a=2 / num_b=2 / num_d=2 — byte-identical to Old-TE)
Verdict: at parity. On the byte-identical num_a/b/d=2 comparison the bridge is CSV: |
Resolve conflicts in unified_gemm_codegen.py and gemm_utils.py with the grouped-GEMM bridge that landed on develop. All conflicts were the multi_abd (this branch) and grouped (develop) variants touching the same extension points; kept both: - GemmVariant enum, KernelNaming suffix, includes, _get_configs_for_variant, variant->operator map, and --variants CLI choices: keep MULTI_ABD and GROUPED. - gemm_utils: fold multi_abd into develop's _ctypes_source_name helper; merge the setup-batch preflight so it checks per-config ctypes sources (develop) while preserving the multi_abd registry-bypass skip of the static lib (branch); keep the multi_abd expand_sweep kwargs + config construction. Also fix two pre-existing parse errors on the branch (the file did not compile before this merge): the multi_abd launch f-string was missing its closing `return ave_time; }}"""`, and the multi_abd_config default dict was missing its closing brace. Address Copilot review feedback: - gemm_multi_abd_ctypes_lib.cpp: strides are forwarded verbatim (the kernel does NOT derive defaults from 0), so validate stride_as/bs/ds and stride_e are non-null and strictly positive, reject otherwise (-1), and correct the ABI doc comment that wrongly claimed <=0 derived a default. - gemm_utils.py: reject non-fp16 multi_abd in expand_sweep (before codegen/build) and fail fast at the top of GpuMultiABDRunner.run, so bf16 no longer fails only after kernels are built. - ctypes/CMakeLists.txt: skip the dispatcher_gemm_multi_abd_lib target when no kernel header exists instead of building a "without kernel" lib that cannot compile (it unconditionally references force-included kernel symbols).
…bridge_gemm_multi_abd # Conflicts: # projects/composablekernel/dispatcher/python/gemm_utils.py
…do not default to gfx942
Contributor
Author
Arch handling fix — do not default to gfx942Addresses the review feedback "don't default to gfx942. Use get_arch and throw an error for unsupported arch types." What changed (
Why: the old default meant a build with no Scope note: the get_arch / clang-format-18 clean. Commit: 18a1a9f. |
yraparti
approved these changes
Jul 23, 2026
ozturkosu
enabled auto-merge (squash)
July 23, 2026 02:06
ozturkosu
disabled auto-merge
July 23, 2026 02:48
…raint expand_sweep emitted the full wave_m x wave_n x wave_k cartesian product but never applied Old-TE's validate_warp_configuration gate (WARP_SUPPORTED_COMBINATIONS[arch], which allows only [1,4,1]/[2,2,1]/[4,1,1] on CDNA). The bridge therefore generated warps-per-block triples (product 1/2/16) that Old-TE never emits, diverging the two instance sets (gfx950 whole-config: bridge 5308 vs Old-TE 4752 buildable, overlap only 1920). Add _warp_config_supported() (reusing the canonical table from gemm_validation_utils with a hardcoded fallback) and gate the emit loop on it, right after the tile-divisibility check, so the bridge and Old-TE produce the same warp configurations for a given config+arch.
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Jul 23, 2026
feat(ck-tile): multi-ABD GEMM TE to dispatcher bridge
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
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.
This was referenced Jul 24, 2026
bghimireamd
pushed a commit
that referenced
this pull request
Jul 28, 2026
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>
This was referenced Jul 29, 2026
ozturkosu
added a commit
that referenced
this pull request
Aug 17, 2026
…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>
This was referenced Aug 18, 2026
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 18, 2026
feat(ck-tile): multi-ABD GEMM TE to dispatcher bridge
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>
ThruptiRajLakshmanaGowda
added a commit
that referenced
this pull request
Aug 19, 2026
…ops (#10132) ISSUE ID: #8997 ## Summary Adds **on-device GPU correctness tests** for the merged bridge operators that had **no in-tree GPU test**: grouped (#9000), multi-ABD (#9305), multi-D (#9308), mx (#9329). Closes part of the GPU-test coverage gap (11 of 15 bridge ops previously shipped only CPU-only unit tests). ## What each test does Builds the op's dispatcher `.so`, runs it **on-device** via the op's `Gpu*Runner`, and compares output to an independent **numpy fp32 reference** with a scale-relative tolerance (`max|err| / max|ref|`). Real math — reference is `A@B` per group / `A@B + D0 + D1` / the runner's built-in multi-ABD verify. No trivial always-pass asserts. They **SKIP cleanly** (not fail) when no GPU / hipcc / static-lib is present, so they add **zero cost on a CPU-only runner** and only execute on a GPU runner. ## Verification (MI300X / gfx942) - grouped — PASS (max_rel 7.6e-3 bf16 / 4.8e-4 fp16) - multi-ABD — PASS (6.4e-4) - multi-D — PASS (1.8e-3) - mx — **gfx950-only**, so it arch-gates and SKIPs on gfx942; verify-pending on a gfx950 runner. Registered in `dispatcher/tests/CMakeLists.txt` under label `dispatcher;python;bridge;gpu`. ## Scope Merged ops only (their runners are on develop). The 7 open bridge ops (batched/preshuffle/batched-contraction + 5 block-scale) get their GPU tests in their own PR branches. ## Test plan - [x] grouped / multi-ABD / multi-D PASS on gfx942 - [x] all SKIP cleanly with no GPU (CPU gate stays green) - [ ] mx verified on a gfx950 runner - [ ] CI runs the new `dispatcher_*_gpu_correctness` entries on a GPU runner --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>
ozturkosu
added a commit
that referenced
this pull request
Aug 24, 2026
Fold the grouped (#9000), multi-D (#9308), and multi-ABD (#9305) variant enablement onto the shared gfx1250 foundation so gemm_universal + these three "full" variants ship in one PR: - grouped: gfx1250 CI config (WMMA 16x16x32) + CPU tests - multi-D: gfx1250 CI config (WMMA 16x16x32) + CPU tests + --arch help - multi-ABD: gfx1250 CI config (WMMA 16x16x32, fp16-only) + CPU tests Each merged variant pinned the CDNA MFMA 32x32x16 tile in its CI config, which does not run on gfx1250 WMMA; the new configs use 16x16x32. No variant codegen change was needed beyond the shared arch enablement. Validated on MI400/gfx1250 (ROCm 10.1): grouped 64/64 fp16 OK; multi-D 80 measurements VERIFY (max_rel ~5e-4); multi-ABD 32/32 fp16 VERIFY. 63 CPU tests pass.
ozturkosu
added a commit
that referenced
this pull request
Aug 27, 2026
Fold the grouped (#9000), multi-D (#9308), and multi-ABD (#9305) variant enablement onto the shared gfx1250 foundation so gemm_universal + these three "full" variants ship in one PR: - grouped: gfx1250 CI config (WMMA 16x16x32) + CPU tests - multi-D: gfx1250 CI config (WMMA 16x16x32) + CPU tests + --arch help - multi-ABD: gfx1250 CI config (WMMA 16x16x32, fp16-only) + CPU tests Each merged variant pinned the CDNA MFMA 32x32x16 tile in its CI config, which does not run on gfx1250 WMMA; the new configs use 16x16x32. No variant codegen change was needed beyond the shared arch enablement. Validated on MI400/gfx1250 (ROCm 10.1): grouped 64/64 fp16 OK; multi-D 80 measurements VERIFY (max_rel ~5e-4); multi-ABD 32/32 fp16 VERIFY. 63 CPU tests pass.
ozturkosu
added a commit
that referenced
this pull request
Sep 2, 2026
JIRA ID : AICK-2101 JIRA ID : AICK-2133 ## Summary This PR enables the **multi-ABD** (`a2_b2_d2`) GEMM TE→dispatcher bridge (#9305) on **gfx1250 / MI400**. It was split out of #10921 so the other three ops (`gemm_universal`, `grouped_gemm`, `gemm_multi_d`) weren't held back. Comparison of this PR's original file set vs `develop` (branch head `56e7bc3`): **So the substantive multi-ABD gfx1250 enablement is already on `develop`; the remaining unique delta in this PR is a single config file (`example_problems.json`).** ## Motivation (why the enablement was needed) gfx1250 runs **WMMA**, not MFMA: the valid fp16/bf16 warp tile is **16x16x32**, whereas the CDNA MFMA **32x32x16** tile in `default_ci_config.json` instantiates zero runnable kernels there. ## Supported surface (gfx1250) The `a2_b2_d2` variant is intentionally narrow — the codegen, ctypes lib, and runner are all **fp16-only** and **rcrr-only** (`expand_sweep` rejects other dtypes/layouts). So **fp16 × rcrr is the complete supported matrix**. ## Validation (on-device, MI400 / gfx1250) Old-TE has no gfx1250 back end, so this is **correctness-only** (GPU output vs an fp32 CPU reference), not a performance baseline. Full `default_config` sweep of the fp16/rcrr surface, 3 tile-divisible shapes (768³ / 1536³ / 3072³): | op | supported matrix | kernels | result rows | fail | |----|------------------|--------:|------------:|-----:| | `gemm_multi_abd` | fp16 × rcrr (only) | 1,640 | 4,920 | **0** | Every kernel in the supported surface is correct on gfx1250 (rel err ≈ 4e-4). ## Note on #11161 (compv3 + intrawave + 8-warp) The full sweep initially surfaced 552 failures, all in the `compv3 + intrawave + 8-warp (2x4x1 / 4x2x1)` class — the **#11161** defect (a wave64 assumption in the `comp_v3` pipeline, wrong on gfx1250 wave32). An interim gate excludes that class so the supported surface reports clean; the real fix, **#11175** (`Use8WarpSchedule`), is already on `develop`. ## Notes / dependencies - Built on #10921's shared gfx1250 foundation (now on `develop` via the sibling series). - Supersedes the earlier closed attempt **#10917**. - Sibling gfx1250 bridge PRs: #10921 (universal / grouped / multi-D), #10927 (grouped quant), #10928 (batched GEMM + contraction), #11043 (rowcol / tensor quant), #10934 (contraction multi-ABD). --------- 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.
JIRA ID : AICK-2086
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 nodispatcher 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.pyexactly:
fp16,rcrrlayout, configurable A/B/D tensor counts, and the element-wiseop set
{PassThrough, AddScale, MultiDMultiply, MultiDAdd}.Test Plan
python3 -m pytest dispatcher/tests/test_multi_abd_bridge.py -v512x512x512 fp16 rcrr, across the default 2/2/2 all-PassThrough config and
non-PassThrough element-wise ops.
Test Result
max_rel = 2.9e-4max_rel = 5.7e-4max_rel = 4.1e-4default_config.json→ 8896.standardvariantexpand_sweepregression clean.gemm_multi_abd_ctypes_lib.cpp.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.