Skip to content

HIP: Enable MMA flash attention for RDNA3 with head size 576#19063

Closed
linus-amg wants to merge 1 commit into
ggml-org:masterfrom
linus-amg:hip-rdna3-mma-fattn-576
Closed

HIP: Enable MMA flash attention for RDNA3 with head size 576#19063
linus-amg wants to merge 1 commit into
ggml-org:masterfrom
linus-amg:hip-rdna3-mma-fattn-576

Conversation

@linus-amg

Copy link
Copy Markdown

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

  1. fattn.cu: Route RDNA3 + head size 576 to MMA kernel (was RDNA4-only)
  2. fattn-mma-f16.cuh:
    • Enable AMD WMMA guards for all RDNA3/RDNA4 (was RDNA4-only)
    • Allow DKQ == 576 in AMD path (was limited to ≤128)
  3. mma.cuh:
    • Add RDNA3 to make_identity_mat()
    • Add RDNA3 f16→f16 WMMA intrinsic with correct 4-argument signature

Performance

Tested on AMD RX 7900 XTX (gfx1100) with GLM-4.7-Flash-REAP-23B-A3B:

Configuration Generation Speed
FA off (before) ~77 t/s
FA on (before - broken) ~27 t/s
FA on (after fix) ~83 t/s

Testing

  • Builds successfully with -DGGML_HIP=ON -DGGML_HIP_ROCWMMA_FATTN=ON -DGPU_TARGETS="gfx1100"
  • GLM-4.7-Flash-REAP inference works with flash attention enabled
  • No regressions on standard head sizes (64, 128)

Related

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 github-actions Bot added Nvidia GPU Issues specific to Nvidia GPUs ggml changes relating to the ggml tensor library for machine learning labels Jan 24, 2026
@linus-amg

Copy link
Copy Markdown
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.

@linus-amg linus-amg closed this Jan 24, 2026
@linus-amg linus-amg deleted the hip-rdna3-mma-fattn-576 branch January 24, 2026 14:23
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ggml changes relating to the ggml tensor library for machine learning Nvidia GPU Issues specific to Nvidia GPUs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant