[MelonUsk1337]

Summary

Staging K/V through shared memory in the scalar flash-attention path (get_fa_tuning_params_scalar) was gated to NVIDIA when it was introduced in #19625. On AMD RDNA the same staging is a large win, including for head sizes >= 256 which the NVIDIA gate excludes. This PR enables it for AMD non-GCN devices in the scalar path only; GCN and the coopmat1 path are left unchanged.

Ref: #25207

Measurements

AMD BC-250 (gfx1013, RDNA2, RADV / Mesa 25.2.6, Debian 13), b10727 stock vs. the same build with this change.
gemma-4-26B-A4B-it-UD-IQ3_S (MoE; hsk/hsv = 256 on SWA layers, 512 on global layers), FA on, llama-bench -p 2048:

depth	stock pp2048	patched pp2048	delta
d0	baseline		+12.3 %
d4096	baseline		+33.7 %
d8192	baseline		+50.0 %

Production use since 2026-09-01 on a 5-node RPC cluster (dirigent + 4 workers over 2.5 GbE) with Qwen3.5-122B-A10B IQ3_S: prefill 96 -> 162 t/s (+68 %), generation 19.0 -> 29.5 t/s (+55 %) on a 16k-token prompt; a daily 20k-token batch job went from 803 s to ~440 s. No correctness or stability regressions observed over ~6 days of continuous use across five gfx1013 devices.

AMD V620 (RDNA2), reported by @draetheus in #25207 with the same change: qwen35 27B Q6_K, +29.8 % at d8192, +62.4 % at d16384.

Notes

• The NVIDIA gate keeps its hsk < 256 && hsv < 256 condition unchanged. The AMD RDNA gate deliberately has no head-size condition, because the measured gains above are on head sizes 256 and 512.
• Not tested on RDNA3/RDNA4 (coopmat paths); the coopmat1 function is intentionally untouched. Happy to gate narrower (e.g. !device->coopmat_support) if maintainers prefer.
• GCN untested, therefore excluded.

[ggml-gh-bot]

Hi @MelonUsk1337, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[KevinFleischer]

Hi @MelonUsk1337, nice results! I'm on AMD Vulcan and would love this PR! The bot flagged this PR because the description doesn't follow the PR template. To save you some time, here's your description rearranged into it. You'd just need to edit the PR's first post, paste this in, and fill in the Requirements section (only you can confirm those):

Overview

Enables K/V shared-memory staging (shmem_staging) in the scalar flash-attention path (get_fa_tuning_params_scalar) for AMD RDNA devices. When this was added in #19625, it was only turned on for NVIDIA.

• Adds an amd_rdna check (AMD vendor, excludes GCN) and turns on shmem_staging for it
• The NVIDIA condition stays the same (hsk < 256 && hsv < 256)
• The AMD RDNA condition has no head-size limit, because the biggest gains were at hsk/hsv = 256 and 512
• The coopmat1 path and GCN are unchanged

This gives a large prompt-processing speedup on RDNA2, and the gain grows with context depth.

Additional information

Related: #25207, #19625

AMD BC-250 (gfx1013, RDNA2, RADV / Mesa 25.2.6, Debian 13)
gemma-4-26B-A4B-it-UD-IQ3_S (MoE; hsk/hsv = 256 on SWA layers, 512 on global layers), pp2048:

depth	delta
d0	+12.3 %
d4096	+33.7 %
d8192	+50.0 %

AMD V620 (RDNA2), qwen35 27B Q6_K: +29.8 % at d8192, +62.4 % at d16384.

Real-world use: running since 2026-09-01 on a 5-node RPC cluster (Qwen3.5-122B-A10B IQ3_S). Prefill went from 96 to 162 t/s (+68 %) and generation from 19.0 to 29.5 t/s (+55 %).

Not tested on RDNA3/RDNA4 or GCN. Test results on those from anyone with the hardware would be welcome.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: <YES / NO – if yes, describe how AI was used>

[KevinFleischer]

Additionally: Here is a quick CR that claude did. Maybe have a look especially on the OTHERS topic that you have touched:

The idea makes sense. Shared-memory staging of K/V was an NVIDIA-only tuning choice in #19625. Turning it on for AMD is a small, low-risk tuning change, and the gains on two RDNA2 cards are big. But there are a few things a reviewer would probably raise:

• The check covers more than RDNA. != AMD_GCN also matches OTHER, which the code falls back to when it can't tell which AMD architecture a card is. So the name amd_rdna is misleading. It turns staging on for RDNA3 and RDNA4, which the PR says weren't tested, and possibly for GCN cards the detection didn't recognize. A tighter check (only AMD_RDNA1/2/3, or only RDNA2 for now) or some RDNA3/4 benchmarks would help.
• The head-size limit is removed for AMD. On NVIDIA, staging was limited to head sizes under 256, apparently because of shared-memory limits. The BC-250 run shows head size 512 works on RDNA2. I couldn't check whether the code's shared-memory limit check accounts for staging on every AMD card.
• Test coverage is narrow. It was tested on two RDNA2 cards, apparently with the Mesa RADV driver on Linux. Other AMD Vulkan drivers (AMDVLK, Windows) weren't tested.
• Style nit. The five-line code comment with benchmark numbers belongs in the PR description rather than the source.

[Clarck1]

Tested this on RDNA2 (RX 6900 XT). The speedup is real, but test-backend-ops finds two crashes with the patch that it doesn't hit without it.

Setup

• AMD Radeon RX 6900 XT (RADV NAVI21) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 65536 | int dot: 1 | matrix cores: none
• Mesa 25.2.8 (RADV), kernel 7.0.0-38, Ubuntu 24.04
• llama.cpp b11617, and b11617 with this PR's diff applied on top (applies cleanly; the hunk lands at line ~1247)
• Model: a Qwen3.6-35B-A3B derivative (qwen35moe, hybrid Gated DeltaNet + full attention), IQ3_S, 10 full-attention layers with head size 256/256 and GQA 8

Speed: llama-bench -ngl 99 -fa 1 -ctk q8_0 -ctv q8_0 -b 2048 -ub 2048 -p 2048 -n 128 -d 0,8192,16384,32768 -r 3, builds interleaved A-B-B-A, values in t/s.

test	b11617	b11617 + #28507	change
pp2048 @ d0	2024	2012	-0.6%
pp2048 @ d8192	1480	1649	+11.4%
pp2048 @ d16384	1186	1393	+17.5%
pp2048 @ d32768	846	1066	+25.9%
tg128 @ d0	102.4	102.5	+0.2%
tg128 @ d32768	91.6	91.5	-0.2%

Generation is unchanged at every depth (d8192 and d16384 too).

Output: llama-perplexity KL divergence, 3 chunks of 8192, patched build against unpatched logits:

• mean KLD 0.000000
• max 0.000061
• same top p 99.992%

Crash 1: SIGSEGV with head size 72

test-backend-ops -o FLASH_ATTN_EXT -b Vulkan0:

• without the patch: 5323/5323 OK, 2/2 backends passed;
• with the patch: segfault every time (reproduced twice, once with -p "hsk=72"). The last case printed before the crash is:

FLASH_ATTN_EXT(hsk=72,hsv=72,nh=4,nr23=[1,1],kv=512,nb=75,mask=1,sinks=1,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,2,1,3],kv_view=1,v_is_view_of_k=0,n_kv_max=0): OK

Head size 72 matters in practice: the Qwen3-VL / Qwen3.5 vision encoder (clip.vision.embedding_length=1152, head_count=16) uses it. A real image through llama-mtmd-cli -fa on did not crash for me, but the test does.

Crash 2: GPU hang with head size 320/256

To avoid crash 1 I gated the staging to AMD non-GCN only when hsk >= 256 && hsv >= 256. With that, the full FLASH_ATTN_EXT run got further (4868 OK, 0 FAIL) and then hung the GPU:

FLASH_ATTN_EXT(hsk=320,hsv=256,nh=4,nr23=[4,1],kv=512,nb=75,mask=0,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3],kv_view=1,v_is_view_of_k=0,n_kv_max=0): OK
terminate called after throwing an instance of 'vk::DeviceLostError'
  what():  vk::Queue::submit: ErrorDeviceLost

Kernel log:

amdgpu 0000:04:00.0: ring comp_1.1.0 timeout, signaled seq=233400, emitted seq=233401
amdgpu 0000:04:00.0:  Process test-backend-op pid 977278 thread test-backend-op pid 977278
amdgpu 0000:04:00.0: Starting comp_1.1.0 ring reset
amdgpu 0000:04:00.0: Ring comp_1.1.0 reset succeeded
amdgpu 0000:04:00.0: [drm] device wedged, but recovered through reset

Since the hang happened with my narrower gate, the cases that fail are reachable with the original diff too (it enables staging for every head size on RDNA). The 256/256 cases pass: -p "hsk=256" gives 151/151 OK.

So on RDNA2 the gain is there for head size 256. But enabling staging for every head size breaks at least 72/72 (segfault) and 320/256 (device lost), which suggests the shared memory sizing doesn't fit for some head sizes. Happy to rerun test-backend-ops or a bench on this card if you push a new revision.

Assisted-by: Claude Opus 5.5 (ran the tests and drafted this comment; the numbers are from my machine)