mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-11 17:51:40 +00:00
7a710c3542
feat(ck): Extend and optimize Quant Gemm Kernel for Aiter a8w8 (#8423) ## Motivation JIRA ID : ROCM-20362 JIRA ID : ROCM-26313 The main purpose of this PR is to allow using CK and CK Tile GEMM for Aiter a8w8 (row-col quantization) and improve its performance. ## Technical Details ### Multiple D for Aiter a8w8 with bias * Support multiple D (bias) in Quant GEMM kernel * Extend CShuffleEpilogue * support int8 -> int32 WarpGemms with row-col quantization * allow shuffling in fp32 before applying multiple D to prevent precision loss ### Large tensors support * Support large tensors in the Quant GEMM kernel by offsetting pointers of matrices A, D and C. This feature can be used when M is large, N and K are relatively small and layout is RCR, it's currently enabled only row-col quantization. * Allow broadcasting of D column vectors in the old CK's `DeviceGemmMultiD_Xdl_CShuffle_V3` with large tensors, this case is used to implement row-col quant scales in Aiter. ### Optimization and workarounds * Use literal 0 as scales for unscaled 16x16x128 and 32x32x64 mfma: llvm uses v_mfma for `__builtin_amdgcn_mfma_scale_f32_..._f8f6f4` instead of v_mfma_scale only if scales are literal 0 values. These instruction don't need loading scales and save vector registers. See https://github.com/ROCm/llvm-project/blob/therock-7.13/llvm/lib/Target/AMDGPU/SIInstrInfo.td#L317-L327 * Add workaround for inefficient buffer_load to lds on 7.2 The 3rd argument of buffer_load_dwordx4 is a scalar register. But the compiler generates a waterfall loop as if lanes can have a different value, even though the original values comes from as scalar register. * Use buffer_store_dwordx4 to store 8 bf16 values in epilogue instead of two buffer_store_dwordx2 * Optimize eight waves pipeline: * Improve instruction scheduling * Remove unneeded barriers * Use nontemporal store/load for C and D matrices in the Quant GEMM kernel (they are used once per block but may consume cache that can be better used for matrices A and B) * Use more efficient padding in epilogue with CTransposed ## Test Plan A new test is added for multiple D Quant Gemm (`TestCkTileGemmRowColQuantMultiD/*.RowColQuantMultiDTest`): ``` ninja test_tile_gemm_quant_rowcol && bin/test_tile_gemm_quant_rowcol ``` Testing the large tensor support is not feasible with the current testing infrastructure because of reference calculations on CPU - it takes several minutes to run a single test case. Such cases are tested manually in Aiter. ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.