Add INT8 support for LDS transpose load#2214
Open
stefankoncarevic wants to merge 17 commits intolds-transpose-load-fp8from
Open
Add INT8 support for LDS transpose load#2214stefankoncarevic wants to merge 17 commits intolds-transpose-load-fp8from
stefankoncarevic wants to merge 17 commits intolds-transpose-load-fp8from
Conversation
This commit adds LDS transpose load optimization support for the attention kernel's GEMM operations. GEMM0 (K × Q): - Added decideLDSTransposeForOperands() call for GEMM0 - K matrix (operand A): Can use LDS transpose when directToLDS is enabled - Q matrix (operand B): Can use LDS transpose only when NOT prefetched to registers. Added bLoadsFromLDS parameter to decideLDSTransposeForOperands() to handle the prefetch case correctly. GEMM1 (V × P): - Added decideLDSTransposeForOperands() call for GEMM1 - V matrix (operand A): Can use LDS transpose when directToLDS is enabled - P matrix (operand B): Never uses LDS transpose since it comes from registers (softmax output), not from global memory. directToLDS is always false for P. API changes: - Extended decideLDSTransposeForOperands() with optional bLoadsFromLDS parameter (default=true). When false, operand B is immediately marked as NOT USABLE for LDS transpose, regardless of other constraints. BlockwiseMatrixParamsAttr changes: - Replaced mnPerXdl with accelDDim and accelKDim parameters. The MFMA instruction geometry (e.g., 16x16, 16x32, 32x16) requires both the D dimension (M or N) and K dimension to correctly configure the LDS transpose load. mnPerXdl only captured one dimension, which was insufficient for determining the correct hardware transpose behavior and offset calculations.
This commit extends the LDS transpose load optimization to support workgroups with 8 waves (blockSize=512) and 16 waves (blockSize=1024). Previously, the optimization was limited to 1-4 waves only. This restriction has been lifted to enable LDS transpose load for larger workgroup sizes commonly used in high-performance GEMM configurations. Changes: - Extended numWaves limit from 4 to 16 in decideLDSTransposeForOperands() - Added wave grid layout computation for 8 waves: - 2×4, 4×2 (preferred balanced layouts) - 1×8, 8×1 (fallback layouts) - Added wave grid layout computation for 16 waves: - 4×4 (preferred balanced layout) - 2×8, 8×2 (semi-balanced layouts) - 1×16, 16×1 (fallback layouts) Updated tests: - lds_transpose_attributes_toblockwise.mlir: Changed CHECK-NOT to CHECK for 8 and 16 wave tests, confirming LDS transpose is now enabled for these configurations - PrLdsTransposeLoad.toml: Added e2e test cases for 8-wave (4×2, 1×8) and 16-wave (8×2, 1×16) grid configurations
When one operand uses regular load and the other uses LDS transpose load, the regular load must use a compatible K-access pattern. The new formula is only applied when: - useLdsTransposeLoad is true (hybrid scenario) - kVec >= kBase (enough elements to decompose) This ensures correct data alignment between regular and transpose loads for MFMA operations, and prevents assertion failures when kpack < kBase. Changes: - Add useLdsTransposeLoad parameter to wrapLDSBufferForLoad - Implement hybrid K-access formula with blk_d/blk_k split - Pass LDS transpose state from BlockwiseGemmToThreadwise - Update tests in PrLdsTransposeLoad.toml
- Move kBase variable declaration earlier in wrapLDSBufferForLoad - Remove duplicate MfmaInsnAttr declaration, reuse existing mfmaAttr - Add negative test for LDS transpose load on unsupported arch (gfx942)
- Add GEMM1 LDS transpose tests (V transpose + P prefetch) to nightly - Create PrLdsTransposeLoadAttention.toml with 14 quick PR tests
Implement ds_read_tr8_b64 offset formulas for FP8/BF8 MFMA (16x32, 32x16). Enable mixed fp8/bf8 type combinations for GEMM operations on gfx950.
- Add INT8 (i8) support in LdsTransposeLoad.cpp for ds_read_tr8_b64 - Support mfma_i32_16x16x32_i8, mfma_i32_16x16x64_i8, mfma_i32_32x32x16_i8, mfma_i32_32x32x32_i8 - Add INT8 16x64 and 32x32 MFMA geometries with double-rate K coverage - Handle kpack=1 case for INT8 MFMAs with kBase=16 in AccelEmitter.cpp - Add validation for INT8 MFMA geometries in RockDialect.cpp - Add e2e tests for INT8 LDS transpose in GEMM and Attention
stefankoncarevic
requested review from
dhernandez0,
djramic,
justinrosner,
pabloantoniom and
umangyadav
stefankoncarevic
force-pushed
the
lds-transpose-load-fp8
branch
3 times, most recently
from
f3176a8 to
a75ab7a
Compare
Disable LDS transpose load for INT8 convolutions when N=1600 (40x40 spatial output) and K<=M or K>2*M. This fixes two significant performance regressions: - 1x64x40x40 K=64: -62.87% regression - 1x384x40x40 K=128: -43.79% regression The heuristic has no impact on GEMM INT8 (no problems with N=1600) and does not affect any CONV INT8 improvements.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Extends LDS transpose load optimization to support INT8 data types for GEMM and Attention kernels on gfx950. This enables hardware-accelerated transposed loads (ds_read_tr8_b64) for all INT8 MFMAs (16x16x32, 16x16x64, 32x32x16, 32x32x32), improving performance for INT8 quantized inference.
Technical Details
Test Plan
Added MLIR unit tests
Added E2E tests
All tests verified on gfx950 hardware with numerical correctness validation
Test Result
Submission Checklist