Stream-K GEMM (Dispatcher Deep-Core Path)

July 15, 2026 · View on GitHub

Stream-K is a single GEMM that splits the K dimension across compute units (CUs) and reduces the partial results, instead of giving each CU a whole output tile. It keeps every CU busy on shapes where a classic data-parallel tiling would leave some idle (tall-skinny / large-K problems), at the cost of a reduction step.

This document explains how to generate, build, run, and test the Stream-K kernels through the CK Tile dispatcher.

Validated platform: AMD Instinct MI300X (gfx942). See Known limitations for gfx950 (MI350) status.

Why Stream-K needs its own path

A plain GEMM rides Dispatcher::run(A, B, C, problem). Stream-K cannot use that signature unchanged: it needs a reduction workspace and a reduction strategy, so its host args type (ck_tile::StreamKHostArgs) is ABI-incompatible with the regular GemmHostArgs. The deep-core path makes Stream-K ride the registry anyway:

codegen (unified_gemm_codegen.py)
   -> generated Stream-K kernel + dispatcher wrapper
   -> Registry::register_kernel(GeneratedStreamKKernelInstance)
   -> Dispatcher::select_kernel(Problem.streamk + reduction_strategy)
   -> GeneratedStreamKKernelInstance::run()  (Dispatcher owns the workspace)
   -> SelectedKernel::launch(StreamKHostArgs, cfg, workspace)

Reduction strategies

The reduction strategy is a compile-time property, so each strategy is a distinct kernel. The registry holds all three side by side and the dispatcher selects by Problem::reduction_strategy:

Strategy	Workspace	Identifier suffix	Notes
`atomic`	none	`_streamk`	partials accumulate directly into C via atomics
`linear`	yes	`_streamk_linear`	partials reduced through a device workspace, in order
`tree`	yes	`_streamk_tree`	tree reduction through a device workspace

Supported datatypes / layouts

Datatypes: fp16, bf16, fp8, bf8. (fp32/fp64 have no MFMA warp tiles; int8 Stream-K is out of scope for this path.)
Layouts: rcr, rrr, ccr, crr — A/B in either order, C is row-major (the atomic C-reset relies on it).

Prerequisites

A full ROCm toolchain with HIP headers (hip/hip_runtime.h) and hipcc. Bare SLURM compute nodes on the cluster often ship an incomplete ROCm, so build inside the CK ROCm container, e.g.:

# on a GPU node (pyxis/enroot), mounting your home:
srun --jobid=<JOBID> --overlap \
  --container-image=/cluster/images/ck/ck_rocm7.1.1_therock_<date>.sqsh \
  --container-mounts=$HOME:$HOME \
  bash -lc '<commands below>'

All commands below are run from the dispatcher root (projects/composablekernel/dispatcher).

1. Generate a Stream-K kernel

The codegen emits all three reduction-strategy headers from one tile config:

python3 codegen/unified_gemm_codegen.py \
    --datatype fp16 --layout rcr \
    --gpu-target gfx942 \
    --variants stream_k \
    --tile-config-json '{
      "tile_config":  {"tile_m":[128],"tile_n":[128],"tile_k":[64],
                       "warp_m":[2],"warp_n":[2],"warp_k":[1],
                       "warp_tile_m":[32],"warp_tile_n":[32],"warp_tile_k":[16],
                       "block_size":[256]},
      "trait_config": {"pipeline":["compv3"],"epilogue":["cshuffle"],"scheduler":["intrawave"],
                       "pad_m":[false],"pad_n":[false],"pad_k":[false],"persistent":[false]},
      "streamk_config": {"reduction_strategy":["atomic","linear","tree"]}
    }' \
    --output-dir ./gen_fp16_rcr

This produces, per strategy, a header named:

gemm_<dtype>_<layout>_compv3_cshuffle_intrawave_<padM>_<padN>_<padK>_<persistent>_<TILE>_<variant>.hpp
# variant ∈ { streamk, streamk_linear, streamk_tree }

Each header force-includes into the global namespace: SelectedKernel, ADataType/BDataType/CDataType/AccDataType, ALayout/BLayout/CLayout, KERNEL_NAME.

Omit --tile-config-json to generate the full arch-filtered tile set instead of a single config. Use --show-arch-info to print what a target GPU supports.

2a. Run via the standalone driver (`03_streamk_gemm_driver.cpp`)

Calls SelectedKernel::launch() directly (bypasses the dispatcher). Use this for apple-to-apple performance measurement against Tile Engine.

HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp

hipcc -std=c++17 --offload-arch=gfx942 -O3 \
  -DCK_TILE_SINGLE_KERNEL_INCLUDE \
  -I ../include -I gen_fp16_rcr \
  -include "$HDR" \
  examples/gemm/cpp/03_streamk_gemm_driver.cpp -o streamk_gemm_driver

# performance (cold cache, TE-matched defaults):
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 0
# correctness (single cold shot so C matches the reference):
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 1

Option	Default	Meaning
`--m/--n/--k`	3840/4096/2048	GEMM dims
`--warmup`	50	warmup iterations (timing)
`--repeat`	100	timed iterations
`--validate`	1	verify vs `reference_gemm`; forces 1 cold shot, no rotation
`--timer`	1	use the GPU timer
`--flush_cache`	1	flush L2 each iter (cold measurement, like Tile Engine)
`--rotating_count`	1000	rotating input copies to defeat cache (Tile Engine default)

Methodology: leaving the cache warm over-reports TFlops and is the entire source of spurious "dispatcher vs Tile Engine" perf gaps. Always measure perf with the cold-cache defaults (--validate 0); run correctness separately (--validate 1).

2b. Run via the registry/dispatcher (`04_streamk_registry_driver.cpp`)

Exercises the full deep-core path: registers the kernel, lets the dispatcher select it by Problem::reduction_strategy, runs it (dispatcher owns the workspace), and verifies vs the reference with a split-K-aware tolerance.

HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp

# core objects (once, no force-include):
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/dispatcher.cpp -o dispatcher.o
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/registry.cpp   -o registry.o

# driver (force-include one strategy's header):
hipcc -std=c++17 --offload-arch=gfx942 -O3 \
  -DCK_TILE_SINGLE_KERNEL_INCLUDE -DGFX_ARCH='"gfx942"' \
  -I ../include -I include -I gen_fp16_rcr -include "$HDR" \
  -c examples/gemm/cpp/04_streamk_registry_driver.cpp -o drv04.o
hipcc --offload-arch=gfx942 drv04.o dispatcher.o registry.o -o streamk_registry_driver

./streamk_registry_driver --m 3840 --n 4096 --k 2048 --strategy atomic --validate 1

Option	Default	Meaning
`--m/--n/--k`	3840/4096/2048	GEMM dims
`--strategy`	atomic	`atomic` / `linear` / `tree` (must match the force-included header)
`--validate`	1	verify vs `reference_gemm` (split-K-aware rtol/atol)

The registry run() path is a functional dispatch path; its Perf: line is a cold-but-non-rotated measurement, not the calibrated apple-to-apple surface. Use the 03 driver (--validate 0) for Tile-Engine-comparable numbers.

3. Test (CTest)

The deep-core path is guarded by test_streamk_registry.py, which generates, builds, dispatches, and verifies every datatype × layout × strategy against two shapes (the default plus a small-M/large-K shape that stresses the split-K tolerance). It SKIPs (exit 77) when no GPU or hipcc is present.

# directly:
python3 tests/test_streamk_registry.py --arch gfx942
python3 tests/test_streamk_registry.py --arch gfx942 --datatypes fp16,bf16 --layouts rcr,ccr

# via ctest (from your dispatcher build dir):
ctest -R dispatcher_test_streamk_registry --output-on-failure

Verification tolerance (why Stream-K is special)

Stream-K reduces kbatch partial products into each output element, so the accumulation error is larger than a single-pass GEMM. The drivers use the same split-K-aware tolerance as Tile Engine (calculate_rtol_atol): kbatch is taken from the kernel's own tile partitioner, and the tolerance is max(per-split threshold, split-K-reduction threshold). Using the plain get_relative/absolute_threshold(K) here spuriously FAILs correct atomic results on small-M/N, large-K shapes.

Known limitations

gfx950 (MI350) fp8/bf8 not validated. On CDNA4 the fp8/bf8 host reference/codec hits an FNUZ-vs-OCP format mismatch; those combos currently fail verification. fp16 and bf16 are fine on gfx950. Validate/gate before enabling fp8/bf8 there.
Tile coverage is narrower than Tile Engine. The dispatcher emits fewer Stream-K tiles than TE (e.g. fp16 rcr TE=180 vs DISP=73). Numeric+perf parity is validated per matched tile config, not over the whole TE tile surface. See the coverage note at the STREAM_K variant in codegen/unified_gemm_codegen.py.

File map

Path	Role
`codegen/unified_gemm_codegen.py`	generates Stream-K kernels + dispatcher wrappers (`--variants stream_k`)
`include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp`	`GeneratedStreamKKernelInstance` (registry/workspace/launch glue)
`include/ck_tile/dispatcher/kernel_key.hpp`	registry key carrying `streamk` + `reduction_strategy`
`examples/gemm/cpp/03_streamk_gemm_driver.cpp`	standalone driver (direct `launch`, perf surface)
`examples/gemm/cpp/04_streamk_registry_driver.cpp`	deep-core driver (Registry → Dispatcher → verify)
`tests/test_streamk_registry.py`	CTest `dispatcher_test_streamk_registry`