SIGN IN SIGN UP

[rocm-libraries] ROCm/rocm-libraries#9329 (commit f66c060)

feat(ck-tile): microscaling (mx) GEMM TE to dispatcher bridge
 (#9329)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 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
M
Muhammed Emin Ozturk committed
85ea951f0e0b2a6eab36fe4005a5286b760c08b5
Parent: 5305cdd
Committed by assistant-librarian[bot] <assistant-librarian[bot]@users.noreply.github.com> on 7/22/2026, 7:17:56 PM