AMD PR Set 3 of 6 — Flash Attention Optimization RDNA3.5

[Geramy]

AMD PR Set 3 of 6 — Review and Summarization Flash Attention Optimization RDNA3.5

Johannes and llama.cpp team, below is a group of related patches from AMD intended to optimize a set of high-priority models on the ROCm backend. I've grouped these together — from a larger set of patches — to hopefully make your job easier in reviewing them, and I've provided a short summary for each. These PR numbers (#40, #41, etc.) refer to the AMD team's downstream fork at AMD-Ecosystem/llama.cpp, where this work landed first; the goal is to bring it upstream in reviewable batches.

This third group is the smallest: two patches covering flash-attention with head dim D=256 on RDNA3.5, one for each build variant (rocWMMA on and off). Both are RDNA3.5 gated and change kernel selection/tile config.

3. Flash-Attention RDNA3.5 D=256 Optimizations

1. #40 — ggml-cuda: use mma-f16 flash-attn for D=256 on RDNA3.5 — roberteg16
   • The rocWMMA-on path. Routes D=256 GQA to the mma-f16 kernel, which packs ncols2 query heads into a full 16-wide tile instead of leaving the wmma-f16 tile mostly empty. +13.7% pp1024.
2. #41 — ggml-cuda: tune D=256 tile flash-attn config for RDNA3.5 — roberteg16
   • The rocWMMA-off path. Overrides just the D=256 / ncols=32 tile config with a smaller K tile and higher occupancy (nbatch_K 128 -> 64, occupancy 3 -> 4). Small effect on this model; helps models with a 256-wide head dim more directly.

Note: Happy to do the front-line review on these conventions, obvious issues, and benchmark coverage.

I have spoke to @am17an please feel free to start pushing PRs into llama.cpp and reference this discussion, also as you tag this discussion I'll update the PR links to the llama.cpp PRs as they come in for accuracy.

Discussions:
AMD PR Set 1 of 6 — Tooling and Profiling modifications and additions
AMD PR Set 2 of 6 — MMQ and MoE - MatMul Optimizations
AMD PR Set 3 of 6 — Flash Attention Optimization RDNA3.5
AMD PR Set 4 of 6 Matvec and Fusion Decode Optimizations
AMD PR Set 5 of 6 — GEMM Prefill Weight Optimizations
AMD PR Set 6 of 6 - Layout and Concurrency Optimizations

[Geramy]

A small update in general for all discussions, a Tailscale device has been provided to am17an for testing with RDNA3.5 and 4.
AMD also has offered to provide devices for CI pipeline here,

[jimw567](https://github.com/jimw567)
[1 hour ago](https://github.com/ggml-org/llama.cpp/discussions/26377#discussioncomment-17856007)
@Geramy and @am17an Regarding hardware, I have just secured two boards (strix halo and strix point) to be used as GHA runners in the public domain. Can you please let me know who I can work with to add them up to the list of runners in upstream llama.cpp?

I myself can provide another machine through Tailscale with RDNA3.5 and 4 but the 4 is over TB5 in an external razor enclosure and would include a .pem and ssh access with root.

[am17an]

For this set of PRs some important information - rocWMMA has been removed #26046