ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 17:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
6648115aed5183e9906ec8fc0b368635c5fb83a4
composable_kernel/dispatcher/python
History
Muhammed Emin Ozturk 6648115aed [rocm-libraries] ROCm/rocm-libraries#9000 (commit 9faa8de) 
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).
2026-07-16 02:55:42 +00:00
..
CMakeLists.txt
[rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb)
2026-04-09 10:38:33 -07:00
ctypes_utils.py
[rocm-libraries] ROCm/rocm-libraries#8997 (commit 6e9bfd9)
2026-07-07 01:15:38 +00:00
dispatcher_common.py
[rocm-libraries] ROCm/rocm-libraries#5260 (commit a1834d2)
2026-05-17 00:29:40 -07:00
fmha_utils.py
[rocm-libraries] ROCm/rocm-libraries#5260 (commit a1834d2)
2026-05-17 00:29:40 -07:00
gemm_utils.py
[rocm-libraries] ROCm/rocm-libraries#9000 (commit 9faa8de)
2026-07-16 02:55:42 +00:00
grouped_conv_utils.py
[rocm-libraries] ROCm/rocm-libraries#6327 (commit 1e7a12e)
2026-05-08 13:47:13 -07:00
pytest.ini
Adding dispatcher architecture (#3300)
2026-01-22 09:34:33 -08:00
README.md
[rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb)
2026-04-09 10:38:33 -07:00
requirements.txt
[rocm-libraries] ROCm/rocm-libraries#5676 (commit 1d18339)
2026-04-01 19:25:55 -07:00

README.md

CK Tile Dispatcher Python Utilities

This directory contains Python utilities used by the dispatcher examples.

Contents

Shared Utilities (used by both GEMM and Grouped Conv)

  • dispatcher_common.py - Shared dispatcher infrastructure
    • Path helpers (get_dispatcher_root, get_build_dir, etc.)
    • ValidationResultBase - Structured validation feedback
    • validate_wave_config, validate_warp_tile_config, validate_trait_combo
    • auto_correct_wave, auto_correct_trait - Auto-correction helpers
    • Colors - Cross-platform ANSI color support
    • print_phase, print_success, print_error, print_info - Phased output
    • cleanup_generated_kernels - Cleanup helper

GEMM Utilities

  • ctypes_utils.py - Core ctypes utilities for GEMM Python examples
    • KernelConfig - Kernel configuration dataclass
    • setup_gemm_dispatcher() - Setup dispatcher with auto-correction
    • cleanup_gemm() - Cleanup dispatcher resources
    • GemmRunner - GPU execution helper
    • Auto-correction and validation utilities

Grouped Convolution Utilities

  • grouped_conv_utils.py - Utilities for grouped convolution
    • GroupedConvValidationResult - Validation result (extends ValidationResultBase)
    • validate_grouped_conv_config - Validate a grouped conv config
    • auto_correct_grouped_conv_config - Auto-correct invalid configs
    • get_grouped_conv_default_config - Get default config for a variant
    • GroupedConvDataType - Data type enum (FP16, BF16, FP32, FP8, BF8, INT8)
    • format_grouped_conv_summary - Human-readable config summary

Usage

GEMM Examples

The GEMM Python examples in dispatcher/examples/gemm/python/ import:

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))

from ctypes_utils import (
    KernelConfig,
    setup_gemm_dispatcher,
    cleanup_gemm,
    GemmRunner,
)

Grouped Conv Usage

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))

from grouped_conv_utils import (
    validate_grouped_conv_config,
    auto_correct_grouped_conv_config,
    get_grouped_conv_default_config,
    GroupedConvDataType,
)

# Get a default config
config = get_grouped_conv_default_config(variant="forward", arch="gfx942")

# Validate
result = validate_grouped_conv_config(config)
print(f"Valid: {result.is_valid}")

Requirements

  • Python 3.8+
  • NumPy
  • HIP runtime (for GPU execution)
Reference in New Issue View Git Blame Copy Permalink