Skip to content

Commit 6e8e366

Browse files
committed
[CK_TILE] Add grouped_gemm variant to the TE -> Dispatcher bridge
Grouped GEMM variant of the bridge, stacked on the fp8/bf8/int8 bridge (#8887, itself on #8479). The grouped kernel is multi-problem, so it uses a dedicated registry-bypass ctypes ABI; TE only generates configs and benchmarks. - codegen: grouped variant + launch generator (arch_filter.py, unified_gemm_codegen.py), standalone 02_grouped_gemm_driver.cpp, README. - bridge: grouped_gemm_ctypes_lib.cpp (multi-problem ABI), GpuGroupedGemmRunner + dtype/layout codecs (gemm_utils.py), TE driver/worker harness, GROUPED_GEMM_BRIDGE.md. Coverage: {rcr,rrr,ccr,crr} x {fp16,bf16,fp8,bf8}; validated at Old-TE parity on MI300X/gfx942 (64/64 correct, 64/64 within +/-15%).
1 parent e30be76 commit 6e8e366

11 files changed

Lines changed: 1530 additions & 21 deletions

File tree

projects/composablekernel/dispatcher/bindings/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ This directory contains language bindings for the CK Tile Dispatcher.
88
bindings/
99
|---- ctypes/ # Python ctypes bindings (C API)
1010
| |---- gemm_ctypes_lib.cpp # GEMM dispatcher C API
11+
| |---- grouped_gemm_ctypes_lib.cpp # Grouped (multi-problem) GEMM bridge C API -- see GROUPED_GEMM_BRIDGE.md
1112
| |---- conv_ctypes_lib.cpp # Grouped conv dispatcher C API (fwd + bwd_data)
1213
| |---- conv_bwdw_ctypes_lib.cpp # Grouped conv backward weight C API (separate library)
1314
| |---- fmha_ctypes_lib.cpp # FMHA dispatcher C API (fwd + bwd)
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
<!--
2+
Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
3+
SPDX-License-Identifier: MIT
4+
-->
5+
6+
# Grouped GEMM: Tile Engine -> Dispatcher Bridge
7+
8+
This document describes the **grouped_gemm** variant of the Tile Engine (TE) ->
9+
Dispatcher bridge (PR #8130). It is the grouped counterpart of the regular-GEMM
10+
bridge (#8123/#8479), the fp8/bf8/int8 bridge (#8887), and the Stream-K bridge
11+
(#8136).
12+
13+
## What the bridge is
14+
15+
In the bridge model the **Dispatcher is the single source of truth** for
16+
codegen, build, and runtime; **Tile Engine only generates configs and
17+
benchmarks** them. TE no longer carries its own kernel-instance build path —
18+
it shells out to the dispatcher codegen and runs the resulting kernel.
19+
20+
For most variants the dispatcher runs a kernel through its registry/backend.
21+
Grouped GEMM cannot use that path (see below), so the grouped bridge takes the
22+
same approach as Stream-K: a dedicated ctypes library that **bypasses the
23+
registry** and calls the generated `SelectedKernel::launch(...)` directly.
24+
25+
## Why grouped needs special handling
26+
27+
Grouped GEMM is **multi-problem**: a single launch runs a *list* of `(M, N, K)`
28+
sub-problems, each with its own A/B/C device pointers. Two consequences:
29+
30+
1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`) cannot
31+
express a list of problems.
32+
2. The generated registry wrapper (`generated_tile_backend.hpp::run()`)
33+
hard-codes the single-problem `SelectedKernel::launch(GemmHostArgs, ...)`
34+
signature and will not compile against a grouped `SelectedKernel`.
35+
36+
So the grouped kernel header exposes a different launch signature
37+
38+
```cpp
39+
static float launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>& descs,
40+
const stream_config& stream);
41+
```
42+
43+
and the grouped ctypes lib force-includes one generated kernel header
44+
(`-include ..._grouped.hpp` with `CK_TILE_SINGLE_KERNEL_INCLUDE`), calls that
45+
`launch` directly, and reports the kernel name from the compile-time
46+
`KERNEL_NAME` macro.
47+
48+
## Components
49+
50+
| Layer | File | Role |
51+
|---|---|---|
52+
| Codegen | `dispatcher/codegen/unified_gemm_codegen.py` | `GemmVariant.GROUPED`; `_launch_function_grouped` (DeviceMem internal workspace, `MakeKargs`, persistent/non-persistent grid). Kept in lockstep with PR #8075. |
53+
| Codegen | `dispatcher/codegen/arch_filter.py` | `GEMM_GROUPED` operator tile constraints. |
54+
| C API | `dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp` | Multi-problem ABI; per-group device alloc/copy; layout-derived strides; warmup/repeat timing. |
55+
| Python | `dispatcher/python/gemm_utils.py` | `GroupedGemmProblem` / `GroupedGemmResult`, `GpuGroupedGemmRunner`, `run_grouped`, `build_grouped`, dtype/layout codecs. |
56+
| Python | `dispatcher/python/ctypes_utils.py` | Threads the `grouped` variant into the codegen `--variants` flag. |
57+
| TE driver | `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` | Generates configs, builds `.so`s in parallel, benchmarks in disposable workers. |
58+
| TE worker | `tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py` | Runs one grouped kernel; dtype/layout-aware operand generation. |
59+
60+
## C ABI
61+
62+
```c
63+
int dispatcher_init(void); // lightweight no-op (no registry)
64+
int dispatcher_run_grouped_gemm(
65+
int group_count,
66+
const int64_t* Ms, // [group_count]
67+
const int64_t* Ns, // [group_count]
68+
const int64_t* Ks, // [group_count]
69+
const void** A_ptrs, // host A buffers, one per group
70+
const void** B_ptrs, // host B buffers, one per group
71+
void** C_ptrs, // host C out buffers, one per group
72+
float* time_ms); // out: average kernel time
73+
// returns 0 ok, -1 HIP/throw, -2 arguments unsupported by the kernel
74+
```
75+
76+
The lib `hipMalloc`s A/B/C per group, copies A and B host->device, memsets C,
77+
builds `std::vector<ck_tile::GroupedGemmHostArgs<>>` with **strides derived from
78+
the compile-time `ALayout`/`BLayout`/`CLayout`** of the `-include`d header
79+
(`std::is_same_v<…, RowMajor>`), launches once, then copies each C back. The ABI
80+
is `void*` + element-size, so it is dtype-agnostic; the Python runner owns the
81+
numpy codecs.
82+
83+
## Coverage
84+
85+
The bridge runnable set is exactly the Old-TE grouped_gemm runnable set on
86+
`develop` — no more, no less:
87+
88+
| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) |
89+
|---|---|---|---|---|
90+
| rcr |||||
91+
| rrr |||||
92+
| ccr |||||
93+
| crr |||||
94+
95+
- **Matrix C is always row-major** (grouped builder constraint), so the layout
96+
string varies A/B only.
97+
- **Excluded:** `int8` (rejected by the TE grouped builder), `fp32`/`fp64`
98+
(no MFMA warp tiles). These are excluded on both sides.
99+
- fp8/bf8 use the **FNUZ** encoding on gfx942 (matches the regular #8887 path);
100+
the Python codecs require `ml_dtypes`.
101+
102+
## Building and running
103+
104+
Generate + build one grouped `.so` and run the A/B parity sweep vs Old-TE:
105+
106+
```bash
107+
# Codegen smoke (no GPU): one variant/dtype/layout
108+
python3 dispatcher/codegen/unified_gemm_codegen.py \
109+
--output-dir /tmp/grp --datatype bf16 --layout ccr \
110+
--variants grouped --config dispatcher/codegen/default_config.json
111+
112+
# Full TE-driven parity sweep (build + benchmark)
113+
python3 tile_engine/ops/gemm/grouped_gemm_full_benchmark.py <config.json> \
114+
--arch gfx942 --dtype fp16 --layout rcr --csv grouped_results.csv
115+
```
116+
117+
Timing knobs `CK_TILE_BENCH_WARMUP` (default 50) and `CK_TILE_BENCH_REPEAT`
118+
(default 100) are honored by **both** the grouped ctypes lib and the registry
119+
backend, so bridge-vs-Old-TE A/B comparisons stay matched. For fair parity keep
120+
`flush_cache=false`, `rotating_count=1`, run on a single GPU, and re-measure any
121+
`|gap|>15%` outlier standalone.

0 commit comments

Comments
 (0)