Muhammed Emin Ozturk 1c455b1bf5 [rocm-libraries] ROCm/rocm-libraries#8998 (commit 5501ef1) ...

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.

2026-07-09 23:58:00 +00:00

12 KiB

Raw Blame History

View Raw