HIP: Enable MMA flash attention for RDNA3 with head size 576#19063
Closed
linus-amg wants to merge 1 commit into
Closed
HIP: Enable MMA flash attention for RDNA3 with head size 576#19063linus-amg wants to merge 1 commit into
linus-amg wants to merge 1 commit into
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
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. |
chrisdevchroma
added a commit
to chrisdevchroma/llama.cpp
that referenced
this pull request
May 8, 2026
PR ggml-org#19063 enabled fattn-mma-f16 on RDNA3 for head_size = 576 only. The broader head_size <= 128 dispatch on RDNA3 needs the input tiles to use DATA_LAYOUT_I_MAJOR_MIRRORED to satisfy the static_asserts added by PR ggml-org#22051 in load_ldmatrix(tile<16,8,T,dl>). - fattn-mma-f16.cuh: split mma_tile_sizes for RDNA3 vs RDNA4+MFMA. RDNA3 uses I_MAJOR_MIRRORED for T_A_KQ, T_B_KQ, T_A_VKQ, T_B_VKQ. - mma.cuh: add data_layout templates to load_ldmatrix_trans and to mma(tile<16,8,half2,dl_ab>, ...). Bodies unchanged; the RDNA3 path inside the half2 mma was already coded for halfx16_t which only fits MIRRORED tiles. - mma.cuh: add an RDNA3 get_half2 that uses __shfl_xor_sync at offset 16 to remap C-tile data (column-split across thread subgroups) into a MIRRORED B-tile (each subgroup holds all columns). Parallels the Volta path's offset-2 shuffle. Bench on Strix Halo iGPU (gfx1151), Qwen3.5-9B Q4_K_XL, -fa 1 -ctk f16 -ctv f16 -b 4096 -ub 2048: pp512 @ d=4096: 980 t/s (mainline fattn-tile-f16: 619) pp512 @ d=8192: 911 t/s (mainline: 439) pp512 @ d=16384: 805 t/s (mainline: 281) pp512 @ d=32768: 640 t/s (mainline: ~31) Matches or exceeds the RDNA4 (gfx1200) reference at every depth.
chrisdevchroma
added a commit
to chrisdevchroma/llama.cpp
that referenced
this pull request
May 8, 2026
PR ggml-org#19063 enabled fattn-mma-f16 on RDNA3 for head_size = 576 only. The broader head_size <= 128 dispatch on RDNA3 needs the input tiles to use DATA_LAYOUT_I_MAJOR_MIRRORED to satisfy the static_asserts added by PR ggml-org#22051 in load_ldmatrix(tile<16,8,T,dl>). - fattn-mma-f16.cuh: split mma_tile_sizes for RDNA3 vs RDNA4+MFMA. RDNA3 uses I_MAJOR_MIRRORED for T_A_KQ, T_B_KQ, T_A_VKQ, T_B_VKQ. - mma.cuh: add data_layout templates to load_ldmatrix_trans and to mma(tile<16,8,half2,dl_ab>, ...). Bodies unchanged; the RDNA3 path inside the half2 mma was already coded for halfx16_t which only fits MIRRORED tiles. - mma.cuh: add an RDNA3 get_half2 that uses __shfl_xor_sync at offset 16 to remap C-tile data (column-split across thread subgroups) into a MIRRORED B-tile (each subgroup holds all columns). Parallels the Volta path's offset-2 shuffle.
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
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