HIP: Enable MMA flash attention for RDNA3 with head size 576#19063
Closed
linus-amg wants to merge 1 commit intoggml-org:masterfrom
Closed
HIP: Enable MMA flash attention for RDNA3 with head size 576#19063linus-amg wants to merge 1 commit intoggml-org:masterfrom
linus-amg wants to merge 1 commit intoggml-org:masterfrom
Conversation
This enables MMA-based flash attention on RDNA3 GPUs (gfx1100/1101/1102) for models with head size 576, such as GLM-4.7-Flash and other MLA (Multi-head Latent Attention) models. Previously, flash attention with head size 576 only worked on CUDA (via PR ggml-org#18953) and RDNA4. RDNA3 users had to disable flash attention, resulting in ~3x slower inference. Changes: - fattn.cu: Route RDNA3 + head size 576 to MMA kernel (was RDNA4-only) - fattn-mma-f16.cuh: Enable AMD WMMA for all RDNA3/RDNA4, allow DKQ==576 - mma.cuh: Add RDNA3 to make_identity_mat(), add f16->f16 WMMA intrinsic Tested on AMD RX 7900 XTX (gfx1100) with GLM-4.7-Flash-REAP-23B: - FA off: ~77 t/s - FA on (before, broken): ~27 t/s - FA on (after fix): ~83 t/s
github-actions
bot
added
Nvidia GPU
3 tasks
Author
|
Closing this PR - the RDNA3 f16→f16 WMMA implementation produces incorrect output due to unpacked output format incompatibility with the tile structure. RDNA3 works correctly with tile-based flash attention instead of MMA. May revisit with a proper fix in the future. |
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.
Summary
This PR enables MMA-based flash attention on RDNA3 GPUs (gfx1100/1101/1102) for models with head size 576, such as GLM-4.7-Flash and other MLA (Multi-head Latent Attention) models.
Previously, flash attention with head size 576 only worked on CUDA (via #18953) and RDNA4. RDNA3 users had to disable flash attention, resulting in ~3x slower inference.
Changes
fattn.cu: Route RDNA3 + head size 576 to MMA kernel (was RDNA4-only)fattn-mma-f16.cuh:DKQ == 576in AMD path (was limited to ≤128)mma.cuh:make_identity_mat()Performance
Tested on AMD RX 7900 XTX (gfx1100) with GLM-4.7-Flash-REAP-23B-A3B:
Testing
-DGGML_HIP=ON -DGGML_HIP_ROCWMMA_FATTN=ON -DGPU_TARGETS="gfx1100"Related