mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-18 01:28:27 +00:00
6648115aed
feat(ck-tile): add grouped GEMM variant to TE to dispatcher bridge (#9000) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Re-opened from #8130 with a policy-compliant branch name (`users/muozturk/ck-tile/dispatcher-te-bridge-grouped-gemm`). Supersedes #8130. ## What this PR does Routes the **grouped_gemm** variant through the Tile Engine (TE) → Dispatcher **bridge**: TE only generates configs and benchmarks; the Dispatcher owns codegen, build, and runtime. This is the grouped counterpart of the regular-GEMM bridge (#8123/#8479), the fp8/bf8/int8 bridge (#8887), and the Stream-K bridge (#8136). **This PR now also contains the grouped Dispatcher codegen** that previously lived in #8075 — that PR has been **closed in favor of this one** to keep the grouped codegen in a single place (it was otherwise duplicated across both). ## Why grouped needs special handling Grouped GEMM is **multi-problem**: one launch runs a *list* of `(M, N, K)` sub-problems with arrays of A/B/C device pointers. 1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`) cannot express a list of problems. 2. The generated registry wrapper (`generated_tile_backend.hpp::run()`) hard-codes the single-problem launch and won't compile against a grouped `SelectedKernel`. So the grouped path **bypasses the registry**: a dedicated ctypes lib calls the generated `SelectedKernel::launch(descs, stream)` directly and reports the name from the compile-time `KERNEL_NAME` macro. ## Changes **Codegen (absorbed from #8075)** - `codegen/arch_filter.py` — `GEMM_GROUPED` operator tile constraints. - `codegen/unified_gemm_codegen.py` — `GemmVariant.GROUPED`, the grouped launch generator (DeviceMem internal workspace via `MakeKargs`, persistent/non-persistent grid), `grouped` in `--variants`. - `examples/gemm/cpp/02_grouped_gemm_driver.cpp` — standalone, layout/dtype-generic grouped driver with per-group reference verification. - `codegen/README.md` + `examples/gemm/cpp/README.md` — grouped sections. **Bridge** - `bindings/ctypes/grouped_gemm_ctypes_lib.cpp` — multi-problem, registry-bypass C ABI; per-group device alloc/copy; strides derived from the compile-time `ALayout/BLayout/CLayout`; warmup/repeat timing matched to Old-TE (`CK_TILE_BENCH_WARMUP/REPEAT`). - `python/gemm_utils.py` — `GroupedGemmProblem`/`GroupedGemmResult`, `GpuGroupedGemmRunner`, `run_grouped`, fp16/bf16/fp8(E4M3 FNUZ)/bf8(E5M2 FNUZ) codecs, output-dtype-aware C buffer. - `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` + `run_one_grouped_gemm_kernel.py` — TE driver + worker for the parity sweep. - `bindings/ctypes/GROUPED_GEMM_BRIDGE.md` — design README. ## Coverage (= Old-TE grouped runnable set on develop) | Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) | |---|---|---|---|---| | rcr / rrr / ccr / crr | ✓ | ✓ | ✓ | ✓ | C is always row-major. `int8` (rejected by the TE grouped builder) and `fp32`/`fp64` (no MFMA warp tiles) are excluded on both sides. ## Parity vs Old-TE (MI300X / gfx942) Apples-to-apples (same warmup=50/repeat=100 both sides, A/B interleaved, single GPU, both engines rebuilt fresh, stale-`.so` guard, matched compile flags): - **Correctness: 64/64 PASS.** - **Performance: 64/64 within ±15%.** - The 5 small-shape (1024³ fp8/bf8) rows that initially read >15% were proven by `rocprof` to be a **measurement-harness artifact** (Old-TE's JSON `latency(ms)` rounded to 2 decimals → 30–50% TFLOPS swing on ~0.02 ms kernels), **not** a kernel/codegen difference — bridge and Old-TE launch byte-identical kernels (same grid/VGPR/SGPR, duration ≤3.22%); full-precision re-measure collapses all 5 to <3%. ## Notes - Targets `develop`. Depends on #8997 (fp16/bf16 bridge) and #8998 (fp8/bf8/int8 bridge) merging to `develop` first; until then this PR's diff also shows their content, after which it reduces to the grouped-only files. - Supersedes #8075 (closed).