|
| 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