mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-21 05:19:20 +00:00
86591de476
[CK] [CK_Tile] Add FMHA scaffolding to CK kernel dispatcher (#5260) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation The CK Tile dispatcher currently supports GEMM and Grouped Convolution but has no support for Fused Multi-Head Attention (FMHA). The example/ck_tile/01_fmha folder contains a comprehensive FMHA implementation with forward, backward, split-KV, paged-KV, append-KV, and batch-prefill kernels across multiple GPU architectures — but there is no unified dispatch layer for it. This PR ports the FMHA stack into the dispatcher, following the same architectural patterns established by GEMM and Grouped Convolution, enabling runtime kernel selection, JIT compilation from Python, and a declarative C++ example flow. Autotuning heuristics to follow. ## Technical Details This PR adds FMHA scaffolding to the CK dispatcher framework, mirroring GEMM's layered architecture. Seven new C++ runtime headers provide type definitions (coexisting with upstream headers via __has_include, requiring zero modifications to example/ck_tile/01_fmha/), a problem builder with 18+ setters, Signature + Algorithm kernel key matching, a virtual kernel instance, a DECL_FMHA_KERNEL_SET macro with wildcard support and named tile/wave/warp setters, arch-aware registry with JSON export, and a dispatcher with seqtune-aware selection, configurable timing, and multi-stage execution plans for split-KV (two-stage) and backward (three-stage). The codegen pipeline is driven by a fmha_arch_specs.json capturing per-arch tile tables and pipeline constraints for five architectures (gfx90a/942/950/1100/1201), migrated from hardcoded logic in 01_fmha/codegen/, with supporting modules for C++ symbol mappings, validation rules, and named receipt profiles (ck_default, flash, pytorch, aiter, fp32, fp8). Python integration (fmha_utils.py) mirrors the C++ layer with JIT compilation, parallel multi-kernel builds, HIP memory management via ctypes, tolerance-based validation, and a NumPy CPU reference with GQA support. Twenty-seven C++ and thirty-two Python examples cover the full feature surface — forward, split-KV, masks, bias, dropout, GQA, backward, append-KV, batch prefill, fp8, logits soft cap, sink tokens, and parameter sweeps — all JIT-compiled on the fly. ## Test Plan Seven test files cover the runtime types, codegen, and end-to-end correctness. C++ unit tests validate the problem builder, dispatcher planning (single-stage for forward/paged-KV/append-KV; multi-stage for split-KV and backward), registry operations, and the kernel-set declaration macro. Python unit tests verify codegen emission, profile filtering, and 15 validation rules for masks, hdim constraints, and pipeline requirements. GPU execution validation in 01_basic_fmha --validate reports zero errors across 65,536 elements with max absolute error of 7.29e-05. A gold-standard parity suite (test_fmha_parity.py) runs 14 configurations through both the upstream tile_example_fmha_fwd and the dispatcher, comparing exit codes to confirm behavioral parity — all 14 match. ## Test Result The C++ smoke test builds and passes all 9 compiled examples, and a Python JIT sweep (29_sweep_seqlen.py) passes 7/7 configurations reaching up to 375 TFLOPS at seqlen 2048. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
71 lines
7.2 KiB
Markdown
71 lines
7.2 KiB
Markdown
**CK Tile operation support by data type, layout, and GPU target:**
|
|
|
|
| Op | CK Tile Kernel | fp16 | fp8 | bf16 | bf8 | int8 | fp4 | fp6 | rcr | rrr | ccr | crr | 90a | 942 | 950 | 1201 |
|
|
|:--:|----------------|:----:|:---:|:----:|:---:|:----:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:----:|
|
|
| GEMM | gemm_universal [1][2]<br>engine: gemm_universal/<br>example: 03_gemm/ | ✅ | ✅ | ✅ | ✅ | ❌ | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
|
| GEMM | gemm_multi_d [3]<br>engine: gemm_multi_d/<br>example: 19_gemm_multi_d/ | ✅ | | | | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
|
| GEMM | gemm_preshuffle [4]<br>engine: gemm_preshuffle/ | ✅ | ✅ | ✅ | ✅ | | | | ✅ | | | | ✅ | ✅ | ✅ | ❌ |
|
|
| GEMM | streamk_gemm [5][6][7]<br>engine: gemm_streamk/<br>example: 40_streamk_gemm/ | ✅ | ✅ | ❌ | ❌ | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
|
|
| GEMM | batched_gemm<br>example: 16_batched_gemm/ | ❌ | | | | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | batched_contraction<br>example: 41_batched_contraction/ | ❌ | | | | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | block_scale_gemm<br>example: 38_block_scale_gemm/ | | ❌ | ❌ | ❌ | | ❌ | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | flatmm<br>example: 18_flatmm/ | ❌ | ❌ | ❌ | ❌ | | ❌ | ❌ | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | gemm_multi_abd<br>example: 22_gemm_multi_abd/ | ❌ | | | | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | grouped_gemm<br>example: 17_grouped_gemm/ | ❌ | ❌ | ❌ | | | | | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
|
| GEMM | grouped_gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Reduce | multi_reduce2d [8]<br>engine: reduce/<br>example: 05_reduce/ | ✅ | | ❌ | | | | | | | | | ❌ | ✅ | ✅ | ❌ |
|
|
| Reduce | reduce2d<br>example: 05_reduce/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Attention | fmha<br>engine: fmha/<br>example: 01_fmha/ | ✅ | ✅ | ✅ | ❌ | | | | | | | | ✅ | ✅ | ✅ | ❌ |
|
|
| Attention | sparse_attn<br>example: 50_sparse_attn/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Activation | softmax | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Activation | topk_softmax<br>example: 09_topk_softmax/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Conv | grouped_conv<br>example: 20_grouped_convolution/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Data Move | batched_transpose<br>example: 35_batched_transpose/ | ❌ | ❌ | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Data Move | image_to_column<br>example: 04_img2col/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Data Move | permute<br>example: 06_permute/ | ❌ | ❌ | | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Elementwise | elementwise<br>example: 21_elementwise/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| MoE | fused_moe<br>example: 15_fused_moe/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Norm | add_rmsnorm2d_rdquant<br>example: 11_add_rmsnorm2d_rdquant/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Norm | layernorm2d<br>example: 02_layernorm2d/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Norm | norm_reduce | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Norm | rmsnorm2d<br>example: 10_rmsnorm2d/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Pooling | pooling<br>example: 36_pooling/ | ❌ | | | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
| Quant | smoothquant<br>example: 12_smoothquant/ | ❌ | | ❌ | | | | | | | | | ❌ | ❌ | ❌ | ❌ |
|
|
|
|
**Legend:**
|
|
- **CK Tile Kernel column:** First line is the kernel name. Lines prefixed with "engine:" show the tile engine directory under `ops/`. Lines prefixed with "example:" show the CK Tile example directory under `example/ck_tile/`.
|
|
- **Green cell** (✅): CK Tile implementation exists **and** the tile engine supports it.
|
|
- **Red cell** (❌): CK Tile implementation exists **but** the tile engine does **not** support it.
|
|
- **Grey cell** (blank): No CK Tile implementation exists for this combination.
|
|
|
|
**Notes:**
|
|
- All CK Tile GEMM and reduce kernels are architecturally generic (no compile-time GPU guards). The gfx filtering in the tile engine is a validation/testing scope decision, not a code limitation.
|
|
- [1] **gemm_universal:** CMake defaults to `fp8;fp16`. Building bf16/bf8 requires `-DGEMM_UNIVERSAL_DATATYPE="fp16;fp8;bf16;bf8"`.
|
|
- [2] **gemm_universal:** CK Tile supports int8 GEMM (with int32 output) but the tile engine has no int8 configuration.
|
|
- [3] **gemm_multi_d:** CK Tile kernel is type-generic but example and tile engine are fp16-only. Adding other types requires new tile engine configurations.
|
|
- [4] **gemm_preshuffle:** Only supports rcr layout (A=row, B=column, C=row) due to the pre-shuffle data format requirement.
|
|
- [5] **streamk_gemm:** CK Tile supports bf16 and bf8 for streamk, but the tile engine has no default tile configs for them.
|
|
- [6] **streamk_gemm:** Builder and default configs support all 4 layouts, but CMake defaults to `rcr` only. Building others requires `-DGEMM_STREAMK_LAYOUT="rcr;rrr;ccr;crr"`.
|
|
- [7] **streamk_gemm:** CK Tile kernels have no arch-specific guards; tile engine filtering is pending validation for gfx950/gfx1201.
|
|
- [8] **multi_reduce2d:** CK Tile's reduce example supports bf16 input but the tile engine only configures fp16. The reduce kernel adapts to wave32/wave64 at runtime via `is_wave32()`.
|
|
- Reduce operations do not use matrix layouts.
|
|
|
|
**Layout codes:** Each layout code specifies the memory layout of tensors A, B, and C as row-major (r) or column-major (c). For example, `rcr` means A is row-major, B is column-major, and C is row-major. For gemm_multi_d, the instance builder uses 4-character codes (e.g., `rcrr`) where the 4th character specifies the D tensor layout; in the table above, the 3-character A/B/C portion is shown since the D layout is always row-major (r) for all supported configurations.
|
|
|
|
**Data type mapping:** The column labels (fp16, fp8, bf16, bf8, int8, fp4, fp6) refer to the input configuration label passed to the tile engine or CK Tile example. Each label determines the actual types used for the source tensors (A, B), accumulator, and output tensor (C). For 16-bit and 8-bit float types, A and B use the label type, the accumulator is fp32, and the output type C matches the input type for fp16 and bf16 but is promoted to fp16 for fp8 and bf8 since 8-bit precision is insufficient for output storage. int8 uses int32 for both accumulation and output. fp4 is a mixed-precision weight type where B is fp4 and A uses the activation type (fp16 or bf16). fp6 is used by the microscaling (MX) flatmm pipeline where both A and B are fp6 with fp32 accumulation and fp32 output.
|
|
|
|
**Data type mapping per config label:**
|
|
|
|
| Config Label | A (source) | B (source) | Acc | C (output) |
|
|
|:------------:|:----------:|:----------:|:---:|:----------:|
|
|
| fp16 | fp16 | fp16 | fp32 | fp16 |
|
|
| bf16 | bf16 | bf16 | fp32 | bf16 |
|
|
| int8 | int8 | int8 | int32 | int32 |
|
|
| fp8 | fp8 | fp8 | fp32 | fp16 |
|
|
| bf8 | bf8 | bf8 | fp32 | fp16 |
|
|
| fp6 | fp6 | fp6 | fp32 | fp32 |
|
|
| fp4 | fp16 or bf16 | fp4 | fp32 | fp16 or bf16 |
|
|
|
|
For gemm_multi_d, the D tensors (D0, D1) use the same type as the config label (fp16).
|