mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 17:19:12 +00:00
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).
CK Tile Dispatcher Examples
Comprehensive examples for GEMM and Grouped Convolution operations with GPU execution.
Quick Start
Step 1: Build
cd /path/to/composable_kernel/dispatcher
mkdir -p build && cd build
cmake .. \
-DCMAKE_PREFIX_PATH=/opt/rocm \
-DCMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \
-DCMAKE_BUILD_TYPE=Release \
-DGPU_TARGETS="gfx942" \
-DBUILD_DISPATCHER_EXAMPLES=ON
# Build everything (C++ examples + Python libraries)
make -j$(nproc)
# Or build ONLY Python libraries (faster)
make python_libs -j$(nproc)
Step 2: Run C++ Examples
cd build/examples
# GEMM
./gemm_01_basic
./gemm_02_multi_size
./gemm_03_benchmark_validation
./gemm_04_heuristics
./gemm_05_json_export
./gemm_06_multi_registry
Step 3: Run Python Examples
cd /path/to/composable_kernel/dispatcher
# GEMM
python3 examples/gemm/python/01_basic_gemm.py
python3 examples/gemm/python/04_validation.py
python3 examples/gemm/python/07_stress_test.py
python3 examples/gemm/python/08_heuristics.py
Directory Structure
examples/
|---- gemm/
| |---- cpp/ # 7 C++ GEMM examples
| +---- python/ # 11 Python GEMM examples
|
|---- grouped_conv/
| |---- cpp/ # 7 C++ Grouped Conv examples
| +---- python/ # 6 Python Grouped Conv examples
|
|---- fmha/
| |---- cpp/ # 35 C++ FMHA examples (all variants)
| +---- python/ # 38 Python FMHA examples (JIT-compiled)
|
+---- README.md
GEMM Examples
C++ Examples
| # | Example | Description |
|---|---|---|
| 01 | gemm_01_basic |
Basic GEMM with declarative API, autofill, autocorrect |
| 02 | gemm_02_multi_size |
Wildcard expansion for multiple configurations |
| 03 | gemm_03_benchmark_validation |
Performance benchmarking with CPU/GPU validation |
| 04 | gemm_04_heuristics |
Heuristic-based kernel selection |
| 05 | gemm_05_json_export |
Registry JSON export for external tools |
| 06 | gemm_06_multi_registry |
Multiple registries with named kernel sets |
Details: gemm/cpp/README.md
Python Examples
| # | Example | Description |
|---|---|---|
| 01 | 01_basic_gemm.py |
Basic GEMM with multi-kernel support |
| 02 | 02_batch_gemm.py |
Batched GEMM operations |
| 03 | 03_benchmark.py |
Performance benchmarking |
| 04 | 04_validation.py |
CPU reference validation |
| 05 | 05_numpy_integration.py |
NumPy array integration |
| 06 | 06_json_export.py |
Registry JSON export |
| 07 | 07_stress_test.py |
Multi-kernel stress testing (48 configs) |
| 08 | 08_heuristics.py |
Heuristic-based kernel selection (24 configs) |
| 09 | 09_multi_registry.py |
Multiple registries |
| 10 | 10_advanced_benchmark.py |
Advanced benchmark with full control |
| 11 | 11_json_import.py |
Import kernels from JSON |
Details: gemm/python/README.md
Key Features
Declarative Kernel API
Both C++ and Python examples use a declarative approach:
C++ (DECL_KERNEL_SET macro):
DECL_KERNEL_SET(my_kernels,
.add(
Signature().dtype("fp16").layout("rcr"),
Algorithm().tile(256, 256, 32).wave(2, 2, 1).warp(32, 32, 16)
.pipeline("compv4").scheduler("intrawave"),
"gfx942"
)
);
Python (KernelConfig):
config = KernelConfig(
tile_m=256, tile_n=256, tile_k=32,
wave_m=2, wave_n=2, wave_k=1,
warp_tile_m=32, warp_tile_n=32, warp_tile_k=16,
pipeline="compv4", scheduler="intrawave"
)
Autofill and Autocorrect
The build system automatically:
- Autofills missing parameters with sensible defaults
- Autocorrects invalid parameters based on architecture constraints
- Expands wildcards (
*,-1,ANY_INT) to all valid configurations
Architecture Filtering
Kernel configurations are validated against GPU architecture constraints:
- Tile divisibility requirements
- Warp tile constraints
- Pipeline compatibility
Invalid configurations are automatically pruned during code generation.
Validation Examples
C++ Validation
./gemm_03_benchmark_validation --verify 1 # GEMM with CPU reference
./gemm_03_benchmark_validation --verify 2 # GEMM with GPU reference
Python Validation
python3 examples/gemm/python/04_validation.py
python3 examples/gemm/python/07_stress_test.py # Multi-kernel validation
Troubleshooting
Python: Library not found
# Run from dispatcher directory
cd /path/to/composable_kernel/dispatcher
python3 examples/gemm/python/01_basic_gemm.py
C++: Executables not found
# Build with examples enabled
cmake .. -DBUILD_DISPATCHER_EXAMPLES=ON
make -j$(nproc)
# Run from build/examples
cd build/examples
./gemm_01_basic
GPU not detected
rocminfo | grep "Name:"
# Should show: gfx942, gfx90a, etc.
Grouped Convolution
Grouped convolution support has been re-introduced with a unified infrastructure shared with GEMM.
Infrastructure
The grouped convolution code generation, utilities, and build scripts are available:
| Component | Location |
|---|---|
| C++ Headers | include/ck_tile/dispatcher/grouped_conv_*.hpp |
| Python Codegen | codegen/unified_grouped_conv_codegen.py |
| Python Utils | python/grouped_conv_utils.py |
| Build Script | scripts/compile_grouped_conv_examples.py |
Building Grouped Conv Kernels
# Generate grouped conv kernels
python3 codegen/unified_grouped_conv_codegen.py \
--output-dir build/generated_kernels \
--datatype fp16 --variant forward --ndim-spatial 2
# Compile a grouped conv example
python3 scripts/compile_grouped_conv_examples.py my_grouped_conv_example.cpp
See the main README for more details.