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#9000 (commit 9faa8de)

Browse Source 
feat(ck-tile): add grouped GEMM variant to TE to dispatcher
 bridge (#9000)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> Re-opened from #8130 with a policy-compliant branch name
(`users/muozturk/ck-tile/dispatcher-te-bridge-grouped-gemm`). Supersedes
#8130.

## What this PR does

Routes the **grouped_gemm** variant through the Tile Engine (TE) →
Dispatcher **bridge**: TE only generates configs and benchmarks; the
Dispatcher owns codegen, build, and runtime. This is the grouped
counterpart of the regular-GEMM bridge (#8123/#8479), the fp8/bf8/int8
bridge (#8887), and the Stream-K bridge (#8136).

**This PR now also contains the grouped Dispatcher codegen** that
previously lived in #8075 — that PR has been **closed in favor of this
one** to keep the grouped codegen in a single place (it was otherwise
duplicated across both).

## Why grouped needs special handling

Grouped GEMM is **multi-problem**: one launch runs a *list* of `(M, N,
K)` sub-problems with arrays of A/B/C device pointers.

1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`)
cannot express a list of problems.
2. The generated registry wrapper (`generated_tile_backend.hpp::run()`)
hard-codes the single-problem launch and won't compile against a grouped
`SelectedKernel`.

So the grouped path **bypasses the registry**: a dedicated ctypes lib
calls the generated `SelectedKernel::launch(descs, stream)` directly and
reports the name from the compile-time `KERNEL_NAME` macro.

## Changes

**Codegen (absorbed from #8075)**
- `codegen/arch_filter.py` — `GEMM_GROUPED` operator tile constraints.
- `codegen/unified_gemm_codegen.py` — `GemmVariant.GROUPED`, the grouped
launch generator (DeviceMem internal workspace via `MakeKargs`,
persistent/non-persistent grid), `grouped` in `--variants`.
- `examples/gemm/cpp/02_grouped_gemm_driver.cpp` — standalone,
layout/dtype-generic grouped driver with per-group reference
verification.
- `codegen/README.md` + `examples/gemm/cpp/README.md` — grouped
sections.

**Bridge**
- `bindings/ctypes/grouped_gemm_ctypes_lib.cpp` — multi-problem,
registry-bypass C ABI; per-group device alloc/copy; strides derived from
the compile-time `ALayout/BLayout/CLayout`; warmup/repeat timing matched
to Old-TE (`CK_TILE_BENCH_WARMUP/REPEAT`).
- `python/gemm_utils.py` — `GroupedGemmProblem`/`GroupedGemmResult`,
`GpuGroupedGemmRunner`, `run_grouped`, fp16/bf16/fp8(E4M3 FNUZ)/bf8(E5M2
FNUZ) codecs, output-dtype-aware C buffer.
- `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` +
`run_one_grouped_gemm_kernel.py` — TE driver + worker for the parity
sweep.
- `bindings/ctypes/GROUPED_GEMM_BRIDGE.md` — design README.

## Coverage (= Old-TE grouped runnable set on develop)

| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) |
|---|---|---|---|---|
| rcr / rrr / ccr / crr | ✓ | ✓ | ✓ | ✓ |

C is always row-major. `int8` (rejected by the TE grouped builder) and
`fp32`/`fp64` (no MFMA warp tiles) are excluded on both sides.

## Parity vs Old-TE (MI300X / gfx942)

Apples-to-apples (same warmup=50/repeat=100 both sides, A/B interleaved,
single GPU, both engines rebuilt fresh, stale-`.so` guard, matched
compile flags):

- **Correctness: 64/64 PASS.**
- **Performance: 64/64 within ±15%.**
- The 5 small-shape (1024³ fp8/bf8) rows that initially read >15% were
proven by `rocprof` to be a **measurement-harness artifact** (Old-TE's
JSON `latency(ms)` rounded to 2 decimals → 30–50% TFLOPS swing on ~0.02
ms kernels), **not** a kernel/codegen difference — bridge and Old-TE
launch byte-identical kernels (same grid/VGPR/SGPR, duration ≤3.22%);
full-precision re-measure collapses all 5 to <3%.

## Notes

- Targets `develop`. Depends on #8997 (fp16/bf16 bridge) and #8998
(fp8/bf8/int8 bridge) merging to `develop` first; until then this PR's
diff also shows their content, after which it reduces to the
grouped-only files.
- Supersedes #8075 (closed).
This commit is contained in:
Muhammed Emin Ozturk
2026-07-16 02:55:42 +00:00
committed by assistant-librarian[bot]
parent a6028c883b
commit 6648115aed
12 changed files with 1654 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

335
dispatcher/python/gemm_utils.py
View File

@@ -59,6 +59,88 @@ def _cap(flag: bool) -> str:
return "True" if flag else "False"
# ---------------------------------------------------------------------------
# Dtype codecs: map a bridge dtype token -> numpy dtype for host operands.
#
# fp16 maps to plain numpy; bf16/fp8/bf8 need ml_dtypes. fp8/bf8 use the FNUZ
# encodings (E4M3FNUZ / E5M2FNUZ) that the gfx942 MFMA path expects -- matching
# the regular bridge's fp8/bf8 codec (PR #8887). ml_dtypes is imported lazily so
# the fp16-only path keeps working where ml_dtypes is unavailable.
# ---------------------------------------------------------------------------
# Canonicalize common spellings to a single token.
_DTYPE_ALIASES = {
"fp16": "fp16",
"f16": "fp16",
"half": "fp16",
"float16": "fp16",
"bf16": "bf16",
"bfloat16": "bf16",
"fp8": "fp8",
"fp8_e4m3": "fp8",
"e4m3": "fp8",
"bf8": "bf8",
"fp8_e5m2": "bf8",
"e5m2": "bf8",
}
def numpy_dtype_for(dtype: str):
"""Return the numpy dtype object used for host operands of ``dtype``.
fp16 -> np.float16; bf16/fp8/bf8 require the ``ml_dtypes`` package (imported
lazily) and use FNUZ fp8 encodings for gfx942 parity.
"""
token = _DTYPE_ALIASES.get(str(dtype).lower())
if token is None:
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}")
if token == "fp16":
return np.float16
try:
import ml_dtypes # noqa: WPS433 (lazy: optional dep)
except ImportError as exc: # pragma: no cover - env-dependent
raise RuntimeError(
f"dtype {dtype!r} requires the 'ml_dtypes' package (pip install ml_dtypes)"
) from exc
if token == "bf16":
return np.dtype(ml_dtypes.bfloat16)
if token == "fp8":
return np.dtype(ml_dtypes.float8_e4m3fnuz)
if token == "bf8":
return np.dtype(ml_dtypes.float8_e5m2fnuz)
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") # pragma: no cover
def output_dtype_for(dtype: str) -> str:
"""Return the bridge dtype token of a kernel's OUTPUT for input ``dtype``.
Mirrors ``codegen_common.CommonTypeMappings.get_output_dtype`` (fp8/bf8 ->
fp16, else identity): the generated grouped kernel emits an fp16 ``CDataType``
for fp8/bf8 inputs, so the host C buffer must be sized/typed by the OUTPUT
dtype, not the INPUT dtype. ``codegen_common`` lives on the dispatcher
``codegen`` dir which ctypes_utils already puts on ``sys.path``; import it
lazily so the fp16-only path has no extra dependency.
"""
token = _DTYPE_ALIASES.get(str(dtype).lower())
if token is None:
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}")
try:
from codegen_common import CommonTypeMappings # noqa: WPS433 (lazy)
except ImportError: # pragma: no cover - fall back to the documented mapping
return "fp16" if token in ("fp8", "bf8") else token
return CommonTypeMappings.get_output_dtype(token)
def output_numpy_dtype_for(dtype: str):
"""Numpy dtype of a kernel's OUTPUT buffer for input ``dtype``.
Composition of :func:`output_dtype_for` + :func:`numpy_dtype_for`. For
fp8/bf8 this resolves to ``np.float16`` (2 bytes) because the kernel's
``CDataType`` is fp16; for fp16/bf16 it equals the input dtype.
"""
return numpy_dtype_for(output_dtype_for(dtype))
# ============================================================================
# The shared contract: GemmKernelConfig
# ============================================================================
@@ -151,6 +233,8 @@ class GemmKernelConfig:
name += "_preshuffle"
elif self.variant == "streamk":
name += "_streamk"
elif self.variant == "grouped":
name += "_grouped"
return name
# ------------------------------------------------------------------ #
@@ -264,6 +348,39 @@ class GemmProblem:
return cls(M=int(d["M"]), N=int(d["N"]), K=int(d["K"]))
@dataclass
class GroupedGemmProblem:
"""A grouped GEMM problem: a list of independent (M, N, K) sub-problems
all run by a single grouped kernel launch.
Each group g computes C_g[M_g x N_g] = A_g[M_g x K_g] @ B_g[K_g x N_g].
"""
groups: List[Tuple[int, int, int]]
@classmethod
def uniform(
cls, group_count: int, M: int, N: int, K: int
) -> "GroupedGemmProblem":
"""All groups share the same (M, N, K) shape."""
return cls(groups=[(int(M), int(N), int(K)) for _ in range(int(group_count))])
@property
def group_count(self) -> int:
return len(self.groups)
@property
def flops(self) -> float:
return sum(2.0 * m * n * k for (m, n, k) in self.groups)
def to_dict(self) -> Dict[str, Any]:
return {"groups": [[int(m), int(n), int(k)] for (m, n, k) in self.groups]}
@classmethod
def from_dict(cls, d: Dict[str, Any]) -> "GroupedGemmProblem":
return cls(groups=[(int(m), int(n), int(k)) for (m, n, k) in d["groups"]])
@dataclass
class GemmResult:
output: np.ndarray
@@ -277,6 +394,22 @@ class GemmResult:
return self.status == 0
@dataclass
class GroupedGemmResult:
"""Result of a grouped GEMM launch: one output per group plus aggregate
timing/throughput across the whole batch."""
outputs: List[np.ndarray]
time_ms: float
status: int
tflops: float
kernel_name: str
@property
def success(self) -> bool:
return self.status == 0
# ============================================================================
# ctypes ABI wrapper
# ============================================================================
@@ -294,6 +427,8 @@ class GemmDispatcherLib:
self._path = Path(so_path)
self._lib = ctypes.CDLL(str(self._path))
self._has_indexed = hasattr(self._lib, "dispatcher_get_kernel_name_at")
self._has_grouped = hasattr(self._lib, "dispatcher_run_grouped_gemm")
self._has_single = hasattr(self._lib, "dispatcher_run_gemm")
self._setup_functions()
def _setup_functions(self) -> None:
@@ -316,16 +451,32 @@ class GemmDispatcherLib:
]
lib.dispatcher_get_kernel_name_at.restype = ctypes.c_int
lib.dispatcher_run_gemm.argtypes = [
ctypes.c_void_p, # A (host)
ctypes.c_void_p, # B (host)
ctypes.c_void_p, # C (host)
ctypes.c_int64, # M
ctypes.c_int64, # N
ctypes.c_int64, # K
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_gemm.restype = ctypes.c_int
# Single-problem ABI (regular GEMM .so). Absent on grouped libs.
if self._has_single:
lib.dispatcher_run_gemm.argtypes = [
ctypes.c_void_p, # A (host)
ctypes.c_void_p, # B (host)
ctypes.c_void_p, # C (host)
ctypes.c_int64, # M
ctypes.c_int64, # N
ctypes.c_int64, # K
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_gemm.restype = ctypes.c_int
# Multi-problem ABI (grouped GEMM .so). Absent on regular libs.
if self._has_grouped:
lib.dispatcher_run_grouped_gemm.argtypes = [
ctypes.c_int, # group_count
ctypes.POINTER(ctypes.c_int64), # Ms[]
ctypes.POINTER(ctypes.c_int64), # Ns[]
ctypes.POINTER(ctypes.c_int64), # Ks[]
ctypes.POINTER(ctypes.c_void_p), # A_ptrs[]
ctypes.POINTER(ctypes.c_void_p), # B_ptrs[]
ctypes.POINTER(ctypes.c_void_p), # C_ptrs[]
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_grouped_gemm.restype = ctypes.c_int
lib.dispatcher_cleanup.argtypes = []
lib.dispatcher_cleanup.restype = None
@@ -371,6 +522,52 @@ class GemmDispatcherLib:
)
return status, time_ms.value
def run_grouped(
self,
A_list: List[np.ndarray],
B_list: List[np.ndarray],
C_list: List[np.ndarray],
Ms: List[int],
Ns: List[int],
Ks: List[int],
) -> Tuple[int, float]:
"""Launch the grouped kernel over a batch of (M, N, K) sub-problems.
Each A/B/C entry is a host numpy array already laid out (dtype + row/col
transpose) as the kernel expects for its compile-time layout; the caller
(GpuGroupedGemmRunner) does that per-dtype/per-layout packing. Pointers
are marshalled into ctypes pointer arrays.
"""
if not self._has_grouped:
raise RuntimeError(
f"{self._path} does not expose dispatcher_run_grouped_gemm"
)
g = len(A_list)
c_int64_arr = (ctypes.c_int64 * g)
c_void_arr = (ctypes.c_void_p * g)
ms = c_int64_arr(*[int(m) for m in Ms])
ns = c_int64_arr(*[int(n) for n in Ns])
ks = c_int64_arr(*[int(k) for k in Ks])
a_ptrs = c_void_arr(*[A.ctypes.data_as(ctypes.c_void_p) for A in A_list])
b_ptrs = c_void_arr(*[B.ctypes.data_as(ctypes.c_void_p) for B in B_list])
c_ptrs = c_void_arr(*[C.ctypes.data_as(ctypes.c_void_p) for C in C_list])
time_ms = ctypes.c_float(0.0)
status = self._lib.dispatcher_run_grouped_gemm(
g,
ms,
ns,
ks,
a_ptrs,
b_ptrs,
c_ptrs,
ctypes.byref(time_ms),
)
return status, time_ms.value
def cleanup(self) -> None:
self._lib.dispatcher_cleanup()
@@ -619,6 +816,94 @@ class GpuGemmRunner:
)
class GpuGroupedGemmRunner:
"""High-level runner for the GROUPED variant: construct from a grouped .so
path, call run(A_list, B_list, problem).
Like GpuGemmRunner, the ctypes ABI takes HOST pointers and manages GPU
memory internally (per group), so this runner only marshals the host operand
arrays. The runner is parameterized by ``(dtype, layout)`` (mirroring
``GpuGemmRunner``/``GemmProblem``): the A/B operands are cast to the per-dtype
INPUT numpy codec (fp16/bf16/fp8-E4M3FNUZ/bf8-E5M2FNUZ) and transposed per the
A/B/C layout so the contiguous host buffer matches the layout the kernel was
generated with (the ctypes lib derives strides from the same layouts).
The C/output buffer is sized/typed by the kernel's OUTPUT dtype, not the input
dtype: for fp8/bf8 inputs the generated kernel's ``CDataType`` is fp16, so the
host C buffer is fp16 (2 bytes) even though A/B are 1-byte fp8/bf8. Sizing C by
the input dtype would under-allocate by 2x and the ctypes copy-back would
overrun the host buffer (heap corruption). See :func:`output_numpy_dtype_for`.
"""
def __init__(self, lib_path: Path, dtype: str = "fp16", layout: str = "rcr"):
self.lib = GemmDispatcherLib(lib_path)
if not self.lib.initialize():
raise RuntimeError(
f"Failed to initialize grouped dispatcher .so: {lib_path}"
)
names = self.lib.kernel_names
self._kernel_name = names[0] if names else "unknown"
self._dtype = dtype
# A/B (input) codec vs C (output) codec: they differ for fp8/bf8
# (output is fp16), so keep them distinct to size the C buffer correctly.
self._np_dtype = numpy_dtype_for(dtype)
self._c_np_dtype = output_numpy_dtype_for(dtype)
if len(layout) != 3 or any(ch not in ("r", "c") for ch in layout):
raise ValueError(f"layout must be a 3-char r/c string, got {layout!r}")
self._layout = layout
@property
def kernel_name(self) -> str:
return self._kernel_name
def run(
self,
A_list: List[np.ndarray],
B_list: List[np.ndarray],
problem: GroupedGemmProblem,
) -> GroupedGemmResult:
groups = problem.groups
if len(A_list) != len(groups) or len(B_list) != len(groups):
raise ValueError(
"A_list/B_list length must match the number of groups "
f"({len(A_list)}/{len(B_list)} vs {len(groups)})"
)
Ms = [g[0] for g in groups]
Ns = [g[1] for g in groups]
Ks = [g[2] for g in groups]
la, lb, _lc = self._layout[0], self._layout[1], self._layout[2]
nd = self._np_dtype
c_nd = self._c_np_dtype # OUTPUT dtype (fp16 for fp8/bf8); see __init__.
A_h: List[np.ndarray] = []
B_h: List[np.ndarray] = []
C_h: List[np.ndarray] = []
for A, B, (M, N, _K) in zip(A_list, B_list, groups):
# A logically MxK, B logically KxN, C row-major MxN (CLayout is always
# RowMajor for grouped). Store each operand so its contiguous buffer
# matches its layout: row-major -> as-is, col-major -> transpose.
A_buf = A if la == "r" else A.T
B_buf = B if lb == "r" else B.T
A_h.append(np.ascontiguousarray(A_buf, dtype=nd))
B_h.append(np.ascontiguousarray(B_buf, dtype=nd))
# Size C by the kernel's CDataType (output dtype), NOT the input dtype:
# fp8/bf8 inputs produce fp16 output, so a 1-byte C would be overrun.
C_h.append(np.zeros((M, N), dtype=c_nd))
status, time_ms = self.lib.run_grouped(A_h, B_h, C_h, Ms, Ns, Ks)
tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0
return GroupedGemmResult(
outputs=C_h,
time_ms=time_ms,
status=status,
tflops=tflops,
kernel_name=self._kernel_name,
)
# ============================================================================
# Build API: codegen + hipcc -> .so paths (no GPU)
# ============================================================================
@@ -697,6 +982,18 @@ def _tile_engine_codegen_flags() -> Tuple[str, ...]:
return tuple(flags)
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``.
"""
if config.variant == "grouped":
return "grouped_gemm_ctypes_lib.cpp"
return "gemm_ctypes_lib.cpp"
def _build_compile_jobs(
config: GemmKernelConfig, header: Path
) -> Tuple[Dict[str, Any], Path]:
@@ -705,7 +1002,7 @@ def _build_compile_jobs(
ck_root = root.parent
build_dir = _cu.get_build_dir()
output_dir = _cu.get_generated_kernels_dir()
ctypes_source = root / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp"
ctypes_source = root / "bindings" / "ctypes" / _ctypes_source_name(config)
static_lib = build_dir / "libck_tile_dispatcher.a"
lib_path = build_dir / "examples" / f"lib{config.name}.so"
@@ -786,13 +1083,13 @@ def setup_multiple_gemm_dispatchers(
codegen_script = _cu.get_codegen_path()
output_dir = _cu.get_generated_kernels_dir()
static_lib = _cu.get_build_dir() / "libck_tile_dispatcher.a"
ctypes_source = (
_cu.get_dispatcher_root() / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp"
)
if not static_lib.exists() or not ctypes_source.exists():
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:
raise FileNotFoundError(
"Missing static lib or ctypes source required for compilation:\n"
f" {static_lib}\n {ctypes_source}\n"
f" {static_lib}\n " + "\n ".join(missing) + "\n"
"Build the dispatcher first (cmake + make)."
)
@@ -915,6 +1212,7 @@ def expand_sweep(
arch: str,
dtype: str = "fp16",
layout: str = "rcr",
variant: str = "standard",
) -> List[GemmKernelConfig]:
"""Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list.
@@ -924,8 +1222,8 @@ def expand_sweep(
one GemmKernelConfig. Invalid combinations are dropped via the dispatcher's
own validator, and duplicates (by .name) are collapsed.
The signature is controlled by the `dtype` and `layout` arguments (defaults
to fp16 / rcr).
The operand signature (``dtype``, ``layout``) is applied to every emitted
GemmKernelConfig, so the same sweep expands across any supported dtype/layout.
"""
with open(config_path) as f:
cfg = json.load(f)
@@ -1015,6 +1313,7 @@ def expand_sweep(
pad_k=bool(pk),
persistent=bool(persist),
gfx_arch=arch,
variant=variant,
)
if c.name in seen:
continue