SIGN IN SIGN UP

[rocm-libraries] ROCm/rocm-libraries#9308 (commit 7fcb5f3)

feat(ck-tile): multi-D 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 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 for
fused 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.py`
exactly:
`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 of
D tensors (1 and 2). It follows the registry-bypass bridge pattern used
by the grouped
(#9000) and stream-K (#9028) bridges.

## Test Plan

- Run the CPU-only unit tests (no GPU required):
  `python3 -m pytest dispatcher/tests/test_multi_d_bridge.py -v`
- On-GPU numeric verify over the full capability matrix
(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.
- Confirm the CI config builds real kernels and the sweep covers all ops
× D counts.

## Test Result

- CPU-only unit tests pass (10 passed).
- On-GPU numeric verify: 16/16 combos pass at M=N=K=1024, worst-case
`max_rel = 6.16e-4` (~30x under the 2e-2 gate). Col-major `ccrr` /
`crrr` have real
  on-GPU numeric evidence.
- CI config now builds real kernels (was zero); the sweep expands evenly
across
  `{MultiDAdd, MultiDMultiply} × {num_d 1, 2}` per layout.
- clang-format (18.1.8) clean on `multi_d_gemm_ctypes_lib.cpp`.
- Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16, 4 layouts
× 2 ops ×
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_d` benchmark is
single-D.
  See the parity comment for details.
M
Muhammed Emin Ozturk committed
5a74dec07a894484b9489d0c0e00cd3b52652d18
Parent: 0845868
Committed by assistant-librarian[bot] <assistant-librarian[bot]@users.noreply.github.com> on 7/23/2026, 9:01:47 PM