ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#9028 (commit 2d6a3d6)

Browse Source 
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).
This commit is contained in:
Muhammed Emin Ozturk
2026-07-17 04:08:05 +00:00
committed by assistant-librarian[bot]
parent c401998b69
commit 58c1bcbe2e
9 changed files with 1696 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
dispatcher/python/ctypes_utils.py
View File

@@ -1389,7 +1389,7 @@ class KernelConfig:
gfx_arch: str = "gfx942"
# GEMM variant (affects arch filter validation)
# "standard", "preshuffle", or "multi_d"
# "standard", "preshuffle", "multi_d", or "stream_k"
variant: str = "standard"
@property

77
dispatcher/python/gemm_utils.py
View File

@@ -187,6 +187,9 @@ class GemmKernelConfig:
gfx_arch: str = "gfx942"
variant: str = "standard"
# Stream-K reduction strategy: "atomic" (default), "linear", or "tree".
# Only meaningful when variant == "stream_k".
reduction_strategy: str = "atomic"
# ------------------------------------------------------------------ #
# Derived string fragments
@@ -231,8 +234,12 @@ class GemmKernelConfig:
)
if self.variant == "preshuffle":
name += "_preshuffle"
elif self.variant == "streamk":
elif self.variant == "stream_k":
name += "_streamk"
# Atomic keeps the bare "_streamk" suffix (original parity); linear
# and tree are disambiguated, matching KernelNaming.generate.
if self.reduction_strategy != "atomic":
name += f"_{self.reduction_strategy}"
elif self.variant == "grouped":
name += "_grouped"
return name
@@ -247,7 +254,7 @@ class GemmKernelConfig:
triple ``warp_*`` and the MFMA triple ``warp_tile_*``. We translate
from dispatcher semantics here so the mapping cannot drift.
"""
return {
cfg = {
"tile_config": {
"tile_m": [self.tile_m],
"tile_n": [self.tile_n],
@@ -271,6 +278,11 @@ class GemmKernelConfig:
"persistent": [self.persistent],
},
}
# Pin the single reduction strategy so stream-K codegen emits exactly this
# kernel (the generator otherwise expands all strategies in its default).
if self.variant == "stream_k":
cfg["streamk_config"] = {"reduction_strategy": [self.reduction_strategy]}
return cfg
def to_dict(self) -> Dict[str, Any]:
return {
@@ -985,10 +997,16 @@ def _tile_engine_codegen_flags() -> Tuple[str, ...]:
def _ctypes_source_name(config: GemmKernelConfig) -> str:
"""Pick the ctypes ABI source for a config's variant.
The grouped kernel has a multi-problem launch signature that the
single-problem ``gemm_ctypes_lib.cpp`` cannot express, so grouped configs
compile against the dedicated ``grouped_gemm_ctypes_lib.cpp``.
Variants whose launch ABI differs from the single-problem
``dispatcher_run_gemm`` path need their own lib:
* stream_k keeps the single-problem C ABI (single A/B/C, M/N/K) but its
lib builds a ``StreamKHostArgs`` and calls ``SelectedKernel::launch``
directly instead of routing through the registry.
* grouped has a multi-problem launch signature the single-problem
``gemm_ctypes_lib.cpp`` cannot express.
"""
if config.variant == "stream_k":
return "streamk_gemm_ctypes_lib.cpp"
if config.variant == "grouped":
return "grouped_gemm_ctypes_lib.cpp"
return "gemm_ctypes_lib.cpp"
@@ -1007,6 +1025,29 @@ def _build_compile_jobs(
lib_path = build_dir / "examples" / f"lib{config.name}.so"
obj_file = lib_path.with_suffix(".o")
# The Stream-K path skips the cmake build that would normally create this
# directory, so ensure it exists before hipcc writes the object/.so here.
lib_path.parent.mkdir(parents=True, exist_ok=True)
# Per-variant AMDGPU codegen flags. The regular path matches Tile Engine's
# gemm_universal build via _tile_engine_codegen_flags(). Stream-K must instead
# match TE's gemm_streamk build EXACTLY for a fair A/B: -enable-post-misched=0
# is applied unconditionally (not persistent-gated) and it does NOT use
# -enable-noalias-to-md-conversion=0.
is_streamk = getattr(config, "variant", "") == "stream_k"
variant_flags = (
[
"-std=c++20",
"-fno-offload-uniform-block",
"-mllvm", "--lsr-drop-solution=1",
"-mllvm", "-enable-post-misched=0",
"-mllvm", "-amdgpu-early-inline-all=true",
"-mllvm", "-amdgpu-function-calls=false",
"--offload-compress",
]
if is_streamk
else list(_tile_engine_codegen_flags())
)
compile_cmd = [
_resolve_hipcc(),
@@ -1022,13 +1063,13 @@ def _build_compile_jobs(
"-D__HIP_PLATFORM_AMD__",
f"--offload-arch={config.gfx_arch}",
f'-DGFX_ARCH="{config.gfx_arch}"',
# Match Tile Engine's AMDGPU codegen flags exactly (see
# Match Tile Engine's AMDGPU codegen flags exactly (see variant_flags /
# _tile_engine_codegen_flags). Without them the kernel is compiled with
# different inlining/register allocation, which changes occupancy;
# persistent kernels size their grid by occupancy
# (UniversalGemmKernel::MaxOccupancyGridSize = #CUs x occupancy), so a
# mismatch shows up as large perf gaps vs Tile Engine on persistent tiles.
*_tile_engine_codegen_flags(),
*variant_flags,
"-Wno-undefined-func-template",
"-Wno-float-equal",
str(ctypes_source),
@@ -1042,7 +1083,10 @@ def _build_compile_jobs(
f"--offload-arch={config.gfx_arch}",
"--hip-link",
str(obj_file),
str(static_lib),
# The Stream-K ctypes lib launches the force-included kernel directly and
# references no registry/dispatcher symbols, so its .so does not need the
# dispatcher static lib. The regular path still links it.
*([] if is_streamk else [str(static_lib)]),
"-o",
str(lib_path),
]
@@ -1086,7 +1130,11 @@ def setup_multiple_gemm_dispatchers(
ctypes_dir = _cu.get_dispatcher_root() / "bindings" / "ctypes"
needed_sources = {ctypes_dir / _ctypes_source_name(c) for c in configs}
missing = [str(p) for p in needed_sources if not p.exists()]
if not static_lib.exists() or missing:
# Stream-K .so links only the force-included kernel (no registry/dispatcher
# symbols), so it does not need the dispatcher static lib; only the regular
# path requires it.
streamk_build = all(c.variant == "stream_k" for c in configs)
if (not streamk_build and not static_lib.exists()) or missing:
raise FileNotFoundError(
"Missing static lib or ctypes source required for compilation:\n"
f" {static_lib}\n " + "\n ".join(missing) + "\n"
@@ -1249,6 +1297,14 @@ def expand_sweep(
pad_ks = _expand_values(tr.get("pad_k"), [False])
persistents = _expand_values(tr.get("persistent"), [False])
# Stream-K only: sweep reduction strategies (atomic/linear/tree). Other
# variants keep a single dummy value so the product is unaffected.
if variant == "stream_k":
sk = cfg.get("streamk_config", {})
reductions = _expand_values(sk.get("reduction_strategy"), ["atomic"])
else:
reductions = ["atomic"]
la, lb, lc = layout[0], layout[1], layout[2]
configs: List[GemmKernelConfig] = []
@@ -1270,6 +1326,7 @@ def expand_sweep(
pn,
pk,
persist,
red,
) in itertools.product(
tile_ms,
tile_ns,
@@ -1287,6 +1344,7 @@ def expand_sweep(
pad_ns,
pad_ks,
persistents,
reductions,
):
c = GemmKernelConfig(
dtype_a=dtype,
@@ -1314,6 +1372,7 @@ def expand_sweep(
persistent=bool(persist),
gfx_arch=arch,
variant=variant,
reduction_strategy=red,
)
if c.name in seen:
continue