mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-11 08:48:45 +00:00
e7798e9560
Add asynchronous XOR shuffle support to the Async GEMM pipeline and the MX GEMM pipeline (#7112) ## Motivation The goal of this work is to apply XOR shuffle (swizzle) to the current `comp_async` GEMM pipeline and the `gemm_mx` pipeline. XOR swizzling has been helpful to avoid LDS bank conflicts, as data are redistributed across LDS banks, such that simultaneous threads accessing different rows land on different LDS banks. ## Technical Details A similar approach to the work in the existing eight-waves pipeline was followed. Currently, XOR swizzle support is available for FP8 and BF8 types. FP4 support is also available for MX GEMM. Should the types not match, or should the async vector width be of an unsupported size, then the pipeline falls through to the previously existing ('unswizzled') path. ## Test Plan Execute `test_ck_tile_gemm_pipeline_comp_async` for the Async GEMM pipeline. Execute `test_ck_tile_mx_gemm_fp8` and `test_ck_tile_mx_gemm_fp4` for the MX GEMM pipeline. ## Test Result The tests passed successfully in the `Alola` cluster with MI350 hardware. ## Submission Checklist - [X] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Fernando Jiménez <fernando.jimenez@streamhpc.com> Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
GEMM with CK Tile
This example demonstrates matrix multiplication (GEMM) using the CK Tile programming model, focusing on tile-based parallelism and modular kernel design.
Algorithm and Math
GEMM computes:
C = A \times B
where A is [M, K], B is [N, K], and C is [M, N].
- BlockTile GEMM: Each Block Tile computes a tile of
Cby loading tiles ofAandB, performing blockwise matrix multiply-accumulation, and writing results back with the epilogue.
Tile Programming Model
- Configuration: The Configuration of how the kernel going to be initialized with Block Tile Dimension, Warps Layout, Warp Tile Dimension, and other improvements.
- Block Tile: Each block tile allocates in the compute unit of AMD GPU grabbing the .
- Pipeline: Modular design allows swapping different memory/computation pipelines (e.g., basic, memory-bound, compute).
- Block GEMM: Block Level implementation on how to coordinate the warps iteration and memory layout in block tile.
- Warp GEMM: Each Warp's GEMM Calculation
- Epilogue: Transferring the Accumulated result from register to global memory.
Features
- Flexible Layouts: Supports row/column-major and custom strides for
A,B,C. - Split K: Split the Block Tile also on K Dimension and add it back after the matrix multiply-accumulation. Have a higher performance when M and N is small and K is large.
- Preshuffled GEMM: In inference task, shuffle the GEMM of B (weight) matrix in the warp layout and bypass the shared memory to do the GEMM calculation. Best performance solution for GEMM.
- Precision: Supports fp16, bf16, fp8, bf8, int4 (for B Matrix).
- Validation: CPU/GPU validation and error tolerance options.
Build & Run
mkdir build && cd build
# you can replace <arch> with the appropriate architecture (for example gfx90a or gfx942) or leave it blank
../script/cmake-ck-dev.sh ../ <arch>
# The basic pipeline method on the gemm calculation
make tile_example_gemm_basic -j`nproc`
# The memory bound pipeline on the gemm calculation
make tile_example_gemm_universal -j`nproc`
# The weight preshuffle pipeline on the gemm calculation
make tile_example_gemm_weight_preshuffle -j`nproc`
# gfx125 only: weight preshuffle TDM pipeline with data cache prefetch controls
make tile_example_gemm_weight_preshuffle_tdm_data_cache_prefetch -j`nproc`
This will result in an executable build/bin/tile_example_gemm_basic & build/bin/tile_example_gemm_universal
example
args:
-m m dimension (default:1024)
-n n dimension (default:2048)
-k k dimension (default:64)
-a_layout Tensor A data layout (default: R)
-b_layout Tensor B data layout (default: C)
-c_layout Tensor C data layout (default: R)
-stride_a Tensor A stride (default:0)
-stride_b Tensor B stride (default:0)
-stride_c Tensor C stride (default:0)
-v 0. No validation, 1. Validation on CPU, 2. Validation on GPU (default:2)
-prec data type. fp16/bf16/fp8/bf8 (default:fp16)
-warmup number of iterations before benchmark the kernel (default:50)
-repeat number of iterations to benchmark the kernel (default:100)
-timer gpu:gpu timer, cpu:cpu timer (default:gpu)
-split_k splitK value (default:1)
-init 0:random, 1:linear, 2:constant(1) (default:0)
-persistent 0:non-persistent, 1:persistent (default:0)
-json 0: No Json, 1: Dump Results in Json format (default:0)
-jsonfile json file name to dump results (default:gemm.json)
Source Structure
- Executables:
gemm_basic.cpp,universal_gemm.cpp(different kinds of GEMM implementation) - Utils:
gemm_utils.hpp(helper functions) - Build:
CMakeLists.txt,run_gemm_example.inc - Scripts:
script/(build and run helpers)
Related CK Tile Examples
- 01_fmha: Fused multi-head attention (FMHA)
- 18_flatmm: Preshuffled GEMM alternative solution
- 16_batched_gemm: Batched GEMM with tiles
For distribution, see include/ck_tile/tile_program/tile_distribution/.