mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-11 17:51:40 +00:00
1c455b1bf5
feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8 (all layouts) (#8998) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Extends the Tile Engine ↔ Dispatcher GEMM **bridge** to the remaining data types TE's plain GEMM has MFMA warp tiles for, beyond the fp16/bf16 surface of #8479: - **fp8** (E4M3) and **bf8** (E5M2) → fp16 output, fp32 accumulate - **int8** → int32 output and accumulate (gfx942) All four A/B layout combinations per dtype (row-major C only, matching #8479). `fp32`/`fp64` are intentionally **excluded** — they appear in TE's dtype-string map but have no MFMA warp tiles in `GEMM_WARP_TILE_SUPPORTED_COMBINATIONS`, so no kernel can be generated/run. **Depends on the fp16/bf16 bridge in #8997** (`users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16`), which carries the bridge infrastructure and is not yet merged. This PR targets `develop`, so until #8997 merges its diff also includes the base bridge changes; please merge #8997 first. ## Changes - **Codegen** (`codegen_common.py`, `unified_gemm_codegen.py`): add `int32` to the dtype maps; `get_output_dtype` int8→int32; new `get_acc_dtype` (int8→int32, else fp32); derive `AccDataType`/`CDataType`, the `GEMM_KEY_DTYPE_{C,ACC}` macros, and the registry `dtype_c`/`dtype_acc` from the dtype instead of hard-coding `float`/`fp32`. - **Host harness** (`gemm_utils.py`): fp8/bf8 **FNUZ** (gfx942) uint8 codecs — exact decode (matches device `fp8_t`/`bf8_t`), nearest-representable saturating encode (same pattern as the existing bf16 helper); `GpuGemmRunner.run` encodes A/B and sizes the C buffer per dtype; `expand_sweep` sets `dtype_c`/`dtype_acc`. - **Tests**: `test_gemm_utils.py` adds CPU-only fp8/bf8 codec + output-dtype tests (all green); `test_gemm_parity.py` adds fp8/bf8/int8 cases with dtype-aware inputs/references/tolerances (int8 is bit-exact), GPU-gated like the existing cases. ## Verification done - `test_gemm_utils.py` + `test_codegen_common.py`: **54 passed** (CPU). - Codegen smoke: fp8/int8/fp16 each generate 1 kernel + 1 wrapper, 0 failed; emitted `ADataType/CDataType/AccDataType` and `GEMM_KEY_*` macros are correct (int8→int32_t acc/C; fp8→fp16_t C). - `test_gemm_parity.py` collects 60 cases and skips cleanly without a GPU. - The 16 unrelated failures in `test_examples_integration` / `test_grouped_conv_codegen` / `test_library_caching` are **pre-existing** (verified identical on the base branch; they require a built dispatcher `.a` / GPU). ## Test plan - [x] Merge #8997 (fp16/bf16 bridge), then this reduces to just the fp8/bf8/int8 delta on `develop`. - [x] On an MI300X (gfx942) node: run `python3 tests/test_gemm_parity.py` and confirm fp8/bf8/int8 parity; tune the fp8/bf8 tolerances if needed (current values are first-cut headroom). - [x] FNUZ vs OCP: the fp8/bf8 host codec targets the gfx942 FNUZ format; validate / extend for gfx950 (OCP) before enabling there.
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)