Skip to content

Commit 683dd0c

Browse files
ozturkosuMuhammed Ozturk
ozturkosu
and
Muhammed Ozturk
authored
[rocm-libraries] ROCm/rocm-libraries#8985 (commit 3d4cbef)
feat(ck-tile): add stream_k variant to GEMM Dispatcher codegen > Supersedes #8094 (closed when its branch was renamed to a policy-compliant path). Same commits, same head SHA. ## Motivation This is the next slice of the Tile Engine → Dispatcher consolidation, following the same pattern as the grouped_gemm PR (#8075). It adds the **stream-K** GEMM variant to the unified GEMM codegen, implemented **the dispatcher way** (workspace owned internally via `DeviceMem`, clean `launch(args, stream)` signature), and proves numeric + performance parity against Tile Engine. Branch is based on `develop` and contains **only** the stream-K work (no grouped_gemm commits). ## Technical Details - **`codegen/arch_filter.py`** — added `OperatorType.GEMM_STREAMK` and its tile constraints. - **`codegen/unified_gemm_codegen.py`**: - Added `GemmVariant.STREAM_K`, made it reachable from the CLI (`--variants stream_k`), wired naming (`_streamk` suffix), includes, and the variant→operator map. - New `_launch_function_streamk`: builds a single `StreamKHostArgs`, `MakeKernelArgs` → `GetWorkSpaceSize` → allocate `DeviceMem` workspace **internally** + `SetZero` → `SetWorkSpacePointer` → `IsSupportedArgument` check → `make_kernel` via `launch_kernel_time_mask` with an Atomic-reduction preprocess that zeros C between timed iterations. No external `kargs_ptr` (not the Tile Engine way). - Exported `A/B/CLayout` in the `CK_TILE_SINGLE_KERNEL_INCLUDE` block so a single-kernel driver is layout-generic. - Restricted stream_k configs to the `cshuffle` epilogue (only one the kernel supports). - **`examples/gemm/cpp/03_streamk_gemm_driver.cpp`** (NEW) — minimal standalone driver: `-include`s one generated stream-K header, builds a single A/B/C tensor, calls `SelectedKernel::launch(args, stream)`, verifies against `ck_tile::reference_gemm`, prints TFLOPS/GB/s. The generated GPU kernel (`StreamKKernel<StreamKTilePartitioner, GemmPipeline, GemmEpilogue>`) is identical to TE's; only host-side workspace ownership differs (internal `DeviceMem` vs TE's external pointer). Numerics match. ## Test Plan - **Config:** `fp16_rcr_compv3_cshuffle_intrawave_..._128x128x64_2x2x1_32x32x16` (atomic reduction; exists identically in TE and the dispatcher). - **Shape:** `M=3840, N=4096, K=2048`, `warmup=10`, `repeat=50`, MI300X (gfx942), ROCm 7.1.1. - Run the `03_streamk_gemm_driver` and verify against `ck_tile::reference_gemm`; compare latency/TFLOPS/GB/s against the matching Tile Engine config. > Methodology note: TE's benchmark forces `repeat=1, warmup=0` whenever `verify=1` (the atomic kernel accumulates into C, so it can only verify a single run). A `verify=1` invocation therefore reports a single cold iteration (~0.30 ms), which is **not** a representative perf number. The table below uses TE `verify=0` (so warmup/repeat are honored) for the perf row and a separate TE `verify=1` run for correctness. The dispatcher driver times (warmup=10/repeat=50) and verifies in the same run because it re-zeros C between timed iterations via the masked preprocess. ## Test Result Performance + numerical verification (Dispatcher vs Tile Engine): | | latency (ms) | TFLOPS | GB/s | verify | |---|---|---|---|---| | **Tile Engine** (warmup=10, repeat=50) | 0.24 | 266.7 | 264.8 | correct | | **Dispatcher** (warmup=10, repeat=50) | 0.242 | 266.1 | 264.2 | PASS | | **Δ** | ~0% | ~0% | ~0% | identical | ## Next - Once signed off, delete `tile_engine/ops/gemm_streamk/`. - Continue toward a first-class `dispatcher` GEMM interface folder (roadmap step 5). --- *Deep-core integration details and the multi-datatype / all-layout / scope-correction updates are posted as a separate comment on this PR.* --------- Co-authored-by: Muhammed Ozturk <muozturk@ctr2-alola-ctrl-01.amd.com> Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
1 parent 04caf5b commit 683dd0c

15 files changed

Lines changed: 1770 additions & 11 deletions

‎dispatcher/README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ A unified kernel dispatch system for AMD GPUs with C++ and Python frontends, sup
44

55
**Validated Platform:** AMD Instinct MI300 series (gfx942)
66

7+
> **Stream-K GEMM:** see [STREAMK.md](STREAMK.md) for how to generate, build, run, and
8+
> test the Stream-K deep-core path (atomic/linear/tree reductions).
79
810
---
911

‎dispatcher/STREAMK.md‎

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
# Stream-K GEMM (Dispatcher Deep-Core Path)
2+
3+
Stream-K is a single GEMM that splits the **K** dimension across compute units (CUs)
4+
and reduces the partial results, instead of giving each CU a whole output tile. It
5+
keeps every CU busy on shapes where a classic data-parallel tiling would leave some
6+
idle (tall-skinny / large-K problems), at the cost of a reduction step.
7+
8+
This document explains how to **generate**, **build**, **run**, and **test** the
9+
Stream-K kernels through the CK Tile dispatcher.
10+
11+
> **Validated platform:** AMD Instinct MI300X (gfx942). See [Known limitations](#known-limitations)
12+
> for gfx950 (MI350) status.
13+
14+
---
15+
16+
## Why Stream-K needs its own path
17+
18+
A plain GEMM rides `Dispatcher::run(A, B, C, problem)`. Stream-K cannot use that
19+
signature unchanged: it needs a **reduction workspace** and a **reduction strategy**,
20+
so its host args type (`ck_tile::StreamKHostArgs`) is ABI-incompatible with the
21+
regular `GemmHostArgs`. The deep-core path makes Stream-K ride the registry anyway:
22+
23+
```
24+
codegen (unified_gemm_codegen.py)
25+
-> generated Stream-K kernel + dispatcher wrapper
26+
-> Registry::register_kernel(GeneratedStreamKKernelInstance)
27+
-> Dispatcher::select_kernel(Problem.streamk + reduction_strategy)
28+
-> GeneratedStreamKKernelInstance::run() (Dispatcher owns the workspace)
29+
-> SelectedKernel::launch(StreamKHostArgs, cfg, workspace)
30+
```
31+
32+
### Reduction strategies
33+
34+
The reduction strategy is a **compile-time** property, so each strategy is a
35+
*distinct kernel*. The registry holds all three side by side and the dispatcher
36+
selects by `Problem::reduction_strategy`:
37+
38+
| Strategy | Workspace | Identifier suffix | Notes |
39+
|---|---|---|---|
40+
| `atomic` | none | `_streamk` | partials accumulate directly into C via atomics |
41+
| `linear` | yes | `_streamk_linear` | partials reduced through a device workspace, in order |
42+
| `tree` | yes | `_streamk_tree` | tree reduction through a device workspace |
43+
44+
### Supported datatypes / layouts
45+
46+
- **Datatypes:** `fp16`, `bf16`, `fp8`, `bf8`. (`fp32`/`fp64` have no MFMA warp tiles;
47+
`int8` Stream-K is out of scope for this path.)
48+
- **Layouts:** `rcr`, `rrr`, `ccr`, `crr` — A/B in either order, **C is row-major**
49+
(the atomic C-reset relies on it).
50+
51+
---
52+
53+
## Prerequisites
54+
55+
A full ROCm toolchain with HIP headers (`hip/hip_runtime.h`) and `hipcc`. Bare SLURM
56+
compute nodes on the cluster often ship an incomplete ROCm, so build inside the CK
57+
ROCm container, e.g.:
58+
59+
```bash
60+
# on a GPU node (pyxis/enroot), mounting your home:
61+
srun --jobid=<JOBID> --overlap \
62+
--container-image=/cluster/images/ck/ck_rocm7.1.1_therock_<date>.sqsh \
63+
--container-mounts=$HOME:$HOME \
64+
bash -lc '<commands below>'
65+
```
66+
67+
---
68+
69+
> All commands below are run from the dispatcher root
70+
> (`projects/composablekernel/dispatcher`).
71+
72+
## 1. Generate a Stream-K kernel
73+
74+
The codegen emits all three reduction-strategy headers from one tile config:
75+
76+
```bash
77+
python3 codegen/unified_gemm_codegen.py \
78+
--datatype fp16 --layout rcr \
79+
--gpu-target gfx942 \
80+
--variants stream_k \
81+
--tile-config-json '{
82+
"tile_config": {"tile_m":[128],"tile_n":[128],"tile_k":[64],
83+
"warp_m":[2],"warp_n":[2],"warp_k":[1],
84+
"warp_tile_m":[32],"warp_tile_n":[32],"warp_tile_k":[16],
85+
"block_size":[256]},
86+
"trait_config": {"pipeline":["compv3"],"epilogue":["cshuffle"],"scheduler":["intrawave"],
87+
"pad_m":[false],"pad_n":[false],"pad_k":[false],"persistent":[false]},
88+
"streamk_config": {"reduction_strategy":["atomic","linear","tree"]}
89+
}' \
90+
--output-dir ./gen_fp16_rcr
91+
```
92+
93+
This produces, per strategy, a header named:
94+
95+
```
96+
gemm_<dtype>_<layout>_compv3_cshuffle_intrawave_<padM>_<padN>_<padK>_<persistent>_<TILE>_<variant>.hpp
97+
# variant ∈ { streamk, streamk_linear, streamk_tree }
98+
```
99+
100+
Each header force-includes into the global namespace: `SelectedKernel`,
101+
`ADataType/BDataType/CDataType/AccDataType`, `ALayout/BLayout/CLayout`, `KERNEL_NAME`.
102+
103+
Omit `--tile-config-json` to generate the full arch-filtered tile set instead of a
104+
single config. Use `--show-arch-info` to print what a target GPU supports.
105+
106+
---
107+
108+
## 2a. Run via the standalone driver (`03_streamk_gemm_driver.cpp`)
109+
110+
Calls `SelectedKernel::launch()` **directly** (bypasses the dispatcher). Use this for
111+
apple-to-apple performance measurement against Tile Engine.
112+
113+
```bash
114+
HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp
115+
116+
hipcc -std=c++17 --offload-arch=gfx942 -O3 \
117+
-DCK_TILE_SINGLE_KERNEL_INCLUDE \
118+
-I ../include -I gen_fp16_rcr \
119+
-include "$HDR" \
120+
examples/gemm/cpp/03_streamk_gemm_driver.cpp -o streamk_gemm_driver
121+
122+
# performance (cold cache, TE-matched defaults):
123+
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 0
124+
# correctness (single cold shot so C matches the reference):
125+
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 1
126+
```
127+
128+
| Option | Default | Meaning |
129+
|---|---|---|
130+
| `--m/--n/--k` | 3840/4096/2048 | GEMM dims |
131+
| `--warmup` | 50 | warmup iterations (timing) |
132+
| `--repeat` | 100 | timed iterations |
133+
| `--validate` | 1 | verify vs `reference_gemm`; forces 1 cold shot, no rotation |
134+
| `--timer` | 1 | use the GPU timer |
135+
| `--flush_cache` | 1 | flush L2 each iter (cold measurement, like Tile Engine) |
136+
| `--rotating_count` | 1000 | rotating input copies to defeat cache (Tile Engine default) |
137+
138+
> **Methodology:** leaving the cache warm over-reports TFlops and is the entire
139+
> source of spurious "dispatcher vs Tile Engine" perf gaps. Always measure perf with
140+
> the cold-cache defaults (`--validate 0`); run correctness separately (`--validate 1`).
141+
142+
---
143+
144+
## 2b. Run via the registry/dispatcher (`04_streamk_registry_driver.cpp`)
145+
146+
Exercises the **full deep-core path**: registers the kernel, lets the dispatcher
147+
select it by `Problem::reduction_strategy`, runs it (dispatcher owns the workspace),
148+
and verifies vs the reference with a **split-K-aware tolerance**.
149+
150+
```bash
151+
HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp
152+
153+
# core objects (once, no force-include):
154+
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/dispatcher.cpp -o dispatcher.o
155+
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/registry.cpp -o registry.o
156+
157+
# driver (force-include one strategy's header):
158+
hipcc -std=c++17 --offload-arch=gfx942 -O3 \
159+
-DCK_TILE_SINGLE_KERNEL_INCLUDE -DGFX_ARCH='"gfx942"' \
160+
-I ../include -I include -I gen_fp16_rcr -include "$HDR" \
161+
-c examples/gemm/cpp/04_streamk_registry_driver.cpp -o drv04.o
162+
hipcc --offload-arch=gfx942 drv04.o dispatcher.o registry.o -o streamk_registry_driver
163+
164+
./streamk_registry_driver --m 3840 --n 4096 --k 2048 --strategy atomic --validate 1
165+
```
166+
167+
| Option | Default | Meaning |
168+
|---|---|---|
169+
| `--m/--n/--k` | 3840/4096/2048 | GEMM dims |
170+
| `--strategy` | atomic | `atomic` / `linear` / `tree` (must match the force-included header) |
171+
| `--validate` | 1 | verify vs `reference_gemm` (split-K-aware rtol/atol) |
172+
173+
> The registry `run()` path is a functional dispatch path; its `Perf:` line is a
174+
> cold-but-**non-rotated** measurement, **not** the calibrated apple-to-apple surface.
175+
> Use the `03` driver (`--validate 0`) for Tile-Engine-comparable numbers.
176+
177+
---
178+
179+
## 3. Test (CTest)
180+
181+
The deep-core path is guarded by `test_streamk_registry.py`, which generates, builds,
182+
dispatches, and verifies every `datatype × layout × strategy` against two shapes
183+
(the default plus a small-M/large-K shape that stresses the split-K tolerance). It
184+
**SKIPs** (exit 77) when no GPU or `hipcc` is present.
185+
186+
```bash
187+
# directly:
188+
python3 tests/test_streamk_registry.py --arch gfx942
189+
python3 tests/test_streamk_registry.py --arch gfx942 --datatypes fp16,bf16 --layouts rcr,ccr
190+
191+
# via ctest (from your dispatcher build dir):
192+
ctest -R dispatcher_test_streamk_registry --output-on-failure
193+
```
194+
195+
---
196+
197+
## Verification tolerance (why Stream-K is special)
198+
199+
Stream-K reduces `kbatch` partial products into each output element, so the
200+
accumulation error is larger than a single-pass GEMM. The drivers use the same
201+
split-K-aware tolerance as Tile Engine (`calculate_rtol_atol`): `kbatch` is taken
202+
from the kernel's own tile partitioner, and the tolerance is
203+
`max(per-split threshold, split-K-reduction threshold)`. Using the plain
204+
`get_relative/absolute_threshold(K)` here spuriously FAILs correct atomic results on
205+
small-M/N, large-K shapes.
206+
207+
---
208+
209+
## Known limitations
210+
211+
- **gfx950 (MI350) fp8/bf8 not validated.** On CDNA4 the fp8/bf8 host reference/codec
212+
hits an FNUZ-vs-OCP format mismatch; those combos currently fail verification. fp16
213+
and bf16 are fine on gfx950. Validate/gate before enabling fp8/bf8 there.
214+
- **Tile coverage is narrower than Tile Engine.** The dispatcher emits fewer Stream-K
215+
tiles than TE (e.g. fp16 `rcr` TE=180 vs DISP=73). Numeric+perf parity is validated
216+
per matched tile config, not over the whole TE tile surface. See the coverage note
217+
at the `STREAM_K` variant in `codegen/unified_gemm_codegen.py`.
218+
219+
---
220+
221+
## File map
222+
223+
| Path | Role |
224+
|---|---|
225+
| `codegen/unified_gemm_codegen.py` | generates Stream-K kernels + dispatcher wrappers (`--variants stream_k`) |
226+
| `include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp` | `GeneratedStreamKKernelInstance` (registry/workspace/launch glue) |
227+
| `include/ck_tile/dispatcher/kernel_key.hpp` | registry key carrying `streamk` + `reduction_strategy` |
228+
| `examples/gemm/cpp/03_streamk_gemm_driver.cpp` | standalone driver (direct `launch`, perf surface) |
229+
| `examples/gemm/cpp/04_streamk_registry_driver.cpp` | deep-core driver (Registry → Dispatcher → verify) |
230+
| `tests/test_streamk_registry.py` | CTest `dispatcher_test_streamk_registry` |

‎dispatcher/codegen/arch_filter.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ class OperatorType(Enum):
5050
GEMM = "gemm"
5151
GEMM_PRESHUFFLE = "gemm_preshuffle"
5252
GEMM_MULTI_D = "gemm_multi_d"
53+
GEMM_STREAMK = "gemm_streamk"
5354
CONV_FWD = "conv_fwd"
5455
CONV_BWD_DATA = "conv_bwd_data"
5556
CONV_BWD_WEIGHT = "conv_bwd_weight"
@@ -85,6 +86,20 @@ class OperatorType(Enum):
8586
"tile_n_alignment": 16,
8687
"tile_k_alignment": 8,
8788
},
89+
# NOTE: these are copied from plain GEMM and only gate tile *shape* validity.
90+
# They do NOT express Stream-K's real feasibility requirement -- that a problem
91+
# has enough output tiles to partition K-work across the CUs. That gate is
92+
# runtime (StreamKKernel::IsSupportedArgument / the backend supports() check),
93+
# which lets the dispatcher fall back to a non-Stream-K kernel for too-small
94+
# problems instead of rejecting them at codegen time.
95+
OperatorType.GEMM_STREAMK: {
96+
"min_tile_m": 16,
97+
"min_tile_n": 16,
98+
"min_tile_k": 8,
99+
"tile_m_alignment": 16,
100+
"tile_n_alignment": 16,
101+
"tile_k_alignment": 8,
102+
},
88103
OperatorType.CONV_FWD: {
89104
"min_tile_m": 1, # N dimension can be 1
90105
"min_tile_n": 16, # K (output channels) should be reasonable

0 commit comments

Comments
 (0)