mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-18 01:28:27 +00:00
58c1bcbe2e
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).
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 feedbackvalidate_wave_config,validate_warp_tile_config,validate_trait_comboauto_correct_wave,auto_correct_trait- Auto-correction helpersColors- Cross-platform ANSI color supportprint_phase,print_success,print_error,print_info- Phased outputcleanup_generated_kernels- Cleanup helper
- Path helpers (
GEMM Utilities
ctypes_utils.py- Core ctypes utilities for GEMM Python examplesKernelConfig- Kernel configuration dataclasssetup_gemm_dispatcher()- Setup dispatcher with auto-correctioncleanup_gemm()- Cleanup dispatcher resourcesGemmRunner- GPU execution helper- Auto-correction and validation utilities
Grouped Convolution Utilities
grouped_conv_utils.py- Utilities for grouped convolutionGroupedConvValidationResult- Validation result (extendsValidationResultBase)validate_grouped_conv_config- Validate a grouped conv configauto_correct_grouped_conv_config- Auto-correct invalid configsget_grouped_conv_default_config- Get default config for a variantGroupedConvDataType- 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)