[rocm-libraries] ROCm/rocm-libraries#9305 (commit 34cb16f)
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. M
Muhammed Emin Ozturk committed
6fdb9c3f416c4c5b243ff98aa731c73b9acd28e9
Parent: a267ae9
Committed by assistant-librarian[bot] <assistant-librarian[bot]@users.noreply.github.com>
on 7/23/2026, 6:11:05 PM