mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-18 01:28:27 +00:00
58c1bcbe2e
feat(ck-tile): stream-K GEMM TE to dispatcher bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Re-opened from #8136 with a policy-compliant branch name. Supersedes #8136. ## Summary Routes the **stream_k** GEMM variant through the same Tile Engine (TE) → Dispatcher bridge already landed for regular GEMM (#8123) and grouped GEMM (#8130). The Dispatcher stays the single source of truth for codegen/build/runtime; TE only produces configs and benchmarks. ## Design note Stream-K is a single-problem GEMM with the **same C ABI** as regular GEMM, so the Python runner (`GpuGemmRunner`/`GemmProblem`) and the GPU worker are reused unchanged. The one twist: the registry path can't compile against a Stream-K `SelectedKernel`, so the Stream-K ctypes lib **bypasses the registry** and calls `SelectedKernel::launch(args, stream)` directly — the same approach grouped GEMM uses. The generated launch owns the reduction workspace internally (`DeviceMem`) and uses the Atomic strategy. ## Changes **New** - `bindings/ctypes/streamk_gemm_ctypes_lib.cpp` — single-problem C ABI (`dispatcher_run_gemm`), builds `StreamKHostArgs`, direct launch; returns `0` / `-1` (HIP/throw) / `-2` (args unsupported). - `tile_engine/ops/gemm/streamk_gemm_full_benchmark.py` + `run_one_streamk_gemm_kernel.py` — 3-phase driver (expand → build → subprocess-isolated benchmark) and disposable GPU worker. - `tile_engine/ops/gemm/gemm_streamk/configs/default_config.json` — small sweep config. **Modified** - `dispatcher/python/gemm_utils.py`, `ctypes_utils.py` — thread `variant="stream_k"` through codegen/build and `.so` selection. ## Validation fp16/rcr on gfx942/MI300X: numeric parity vs an fp32 numpy reference (widened fp16-atomic tolerance) and a full driver run of **16/16 OK** with end-to-end name parity; unsupported tiny shapes are reported gracefully (`status -2`), not crashes. fp8/bf8/bf16 now supported via `ml_dtypes` FNUZ codecs. Full tables and the bridge-vs-old-TE comparison are in the comments. ## Test plan - [x] codegen emits `*_streamk.hpp` with stem == `GemmKernelConfig(variant="stream_k").name` - [x] build/link against `streamk_gemm_ctypes_lib.cpp` - [x] numeric parity passes (fp16 atomic tolerance) - [x] full driver run 16/16 OK, name parity end-to-end - [x] unsupported shape → `status -2` handled gracefully ## Next Land #8123, then this; afterwards remove the legacy `tile_engine/ops/gemm_streamk/` machinery (Phase 4).