vulkan: add Intel Xe flash attention optimization kernels (2/3, Xe-LPG Plus/Xe2/Xe3)

[fish-jiang]

Overview

Co-authors: @jxia4intel, @sliu39

PR 2/3 of the Intel Xe optimization series — see #24408 (mega PR, draft) for the full feature set.

Target platforms: Xe-LPG Plus, Xe2, Xe3

This PR adds Intel Xe-specific flash attention optimization kernels for both ARLH iGPU (Xe1, UMA, coopmat1) and Xe2/Xe3. Dependency: builds on top of #24404 (Xe-LPG Plus coopmat1 enable). Independent of #24407 (GEMM+CW).

Flash Attention (Intel Xe)

• New Vulkan shaders: single-phase prefill (flash_attn_hdim64/96/128) and two-phase split prefill/decode variants
• Pipelines keyed by (head_dim, gqa_ratio) for runtime dispatch across various GQA ratios without combinatorial pipeline proliferation
• Supports non-power-of-two GQA ratios via subgroup splitting (qk_groups)
• Intel Xe1 (integrated GPU, UMA, cooperative matrix) and Xe2 paths with separate warptile tuning
• Two-phase decode splits softmax reduction across subgroups; shared QK state copy (fa_copy_qstate) between prefill phases

Performance (Panther Lake B390 + Windows OS)

BEFORE:
C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64>llama-bench.exe -p 8192 -n 0 -r 3 -fa 1 --delay 10 -ngl 99 -m C:\Users\dungeon\Desktop\models\Qwen3.5-35B-A3B-Q4_K_M\Qwen3.5-35B-A3B-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gpt-oss-20b-Q4_K_M\gpt-oss-20b-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gemma-4-26B-A4B-it-UD-Q4_K_M\gemma-4-26B-A4B-it-UD-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3-0.6B.Q4_K_M\Qwen3-0.6B.Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3.5-4B-Q4_K_M\Qwen3.5-4B-Q4_K_M.gguf
load_backend: loaded RPC backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-rpc.dll
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = Intel(R) Arc(TM) B390 GPU (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
load_backend: loaded Vulkan backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-vulkan.dll
load_backend: loaded CPU backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-cpu-alderlake.dll
| model                          |       size |     params | backend    | ngl |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q4_K - Medium |  20.49 GiB |    34.66 B | Vulkan     |  99 |   1 |          pp8192 |        397.28 ± 1.83 |
| gpt-oss 20B Q4_K - Medium      |  10.81 GiB |    20.91 B | Vulkan     |  99 |   1 |          pp8192 |        557.55 ± 2.70 |
| gemma4 26B.A4B Q4_K - Medium   |  15.70 GiB |    25.23 B | Vulkan     |  99 |   1 |          pp8192 |        341.19 ± 1.37 |
| qwen3 0.6B Q4_K - Medium       | 456.11 MiB |   751.63 M | Vulkan     |  99 |   1 |          pp8192 |        802.22 ± 6.66 |
| qwen35 4B Q4_K - Medium        |   2.54 GiB |     4.21 B | Vulkan     |  99 |   1 |          pp8192 |       723.48 ± 11.85 |

build: 3571fa543 (9490)

C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64>

AFTER:
C:\upsteaming_build\subPR3_FA\Release>llama-bench.exe -p 8192 -n 0 -r 3 -fa 1 --delay 10 -ngl 99 -m C:\Users\dungeon\Desktop\models\Qwen3.5-35B-A3B-Q4_K_M\Qwen3.5-35B-A3B-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gpt-oss-20b-Q4_K_M\gpt-oss-20b-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gemma-4-26B-A4B-it-UD-Q4_K_M\gemma-4-26B-A4B-it-UD-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3-0.6B.Q4_K_M\Qwen3-0.6B.Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3.5-4B-Q4_K_M\Qwen3.5-4B-Q4_K_M.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = Intel(R) Arc(TM) B390 GPU (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q4_K - Medium |  20.49 GiB |    34.66 B | Vulkan     |  99 |   1 |          pp8192 |       551.12 ± 30.16 |
| gpt-oss 20B Q4_K - Medium      |  10.81 GiB |    20.91 B | Vulkan     |  99 |   1 |          pp8192 |        819.09 ± 7.08 |
| gemma4 26B.A4B Q4_K - Medium   |  15.70 GiB |    25.23 B | Vulkan     |  99 |   1 |          pp8192 |        624.12 ± 2.05 |
| qwen3 0.6B Q4_K - Medium       | 456.11 MiB |   751.63 M | Vulkan     |  99 |   1 |          pp8192 |      3395.03 ± 29.09 |
| qwen35 4B Q4_K - Medium        |   2.54 GiB |     4.21 B | Vulkan     |  99 |   1 |          pp8192 |      1031.70 ± 44.36 |

build: 472f80478 (9492)

C:\upsteaming_build\subPR3_FA\Release>

BEFORE:
C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64>llama-bench.exe -p 0 -n 128 -d 8192 -r 3 -fa 1 --delay 10 -ngl 99 -m C:\Users\dungeon\Desktop\models\Qwen3.5-35B-A3B-Q4_K_M\Qwen3.5-35B-A3B-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gpt-oss-20b-Q4_K_M\gpt-oss-20b-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gemma-4-26B-A4B-it-UD-Q4_K_M\gemma-4-26B-A4B-it-UD-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3-0.6B.Q4_K_M\Qwen3-0.6B.Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3.5-4B-Q4_K_M\Qwen3.5-4B-Q4_K_M.gguf
load_backend: loaded RPC backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-rpc.dll
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = Intel(R) Arc(TM) B390 GPU (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
load_backend: loaded Vulkan backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-vulkan.dll
load_backend: loaded CPU backend from C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64\ggml-cpu-alderlake.dll
| model                          |       size |     params | backend    | ngl |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q4_K - Medium |  20.49 GiB |    34.66 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         25.72 ± 0.49 |
| gpt-oss 20B Q4_K - Medium      |  10.81 GiB |    20.91 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         25.27 ± 0.06 |
| gemma4 26B.A4B Q4_K - Medium   |  15.70 GiB |    25.23 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         20.72 ± 0.11 |
| qwen3 0.6B Q4_K - Medium       | 456.11 MiB |   751.63 M | Vulkan     |  99 |   1 |   tg128 @ d8192 |         56.30 ± 0.06 |
| qwen35 4B Q4_K - Medium        |   2.54 GiB |     4.21 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         26.47 ± 0.52 |

build: 3571fa543 (9490)

C:\Users\dungeon\Downloads\llama-b9490-bin-win-vulkan-x64>

AFTER:
C:\upsteaming_build\subPR3_FA\Release>llama-bench.exe -p 0 -n 128 -d 8192 -r 3 -fa 1 --delay 10 -ngl 99 -m C:\Users\dungeon\Desktop\models\Qwen3.5-35B-A3B-Q4_K_M\Qwen3.5-35B-A3B-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gpt-oss-20b-Q4_K_M\gpt-oss-20b-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\gemma-4-26B-A4B-it-UD-Q4_K_M\gemma-4-26B-A4B-it-UD-Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3-0.6B.Q4_K_M\Qwen3-0.6B.Q4_K_M.gguf,C:\Users\dungeon\Desktop\models\Qwen3.5-4B-Q4_K_M\Qwen3.5-4B-Q4_K_M.gguf
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = Intel(R) Arc(TM) B390 GPU (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q4_K - Medium |  20.49 GiB |    34.66 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         27.44 ± 0.19 |
| gpt-oss 20B Q4_K - Medium      |  10.81 GiB |    20.91 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         31.86 ± 0.21 |
| gemma4 26B.A4B Q4_K - Medium   |  15.70 GiB |    25.23 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         22.20 ± 0.11 |
| qwen3 0.6B Q4_K - Medium       | 456.11 MiB |   751.63 M | Vulkan     |  99 |   1 |   tg128 @ d8192 |         60.08 ± 0.15 |
| qwen35 4B Q4_K - Medium        |   2.54 GiB |     4.21 B | Vulkan     |  99 |   1 |   tg128 @ d8192 |         27.79 ± 0.60 |

build: 472f80478 (9492)

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, used claude code, then lots of manual review/tweaking.

[ggml-gh-bot]

Hi @fish-jiang, thanks for your contribution!

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

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 3 open PRs.

• AI-generated content: This project does not accept PRs, descriptions or commit messages that are fully or predominantly AI-generated. If you have used AI to assist you in writing code, please make sure to disclose that explicitly.

• Large PR: Large changes require prior discussion (e.g. an issue or RFC) and maintainers may not be able to review this PR as-is. Consider splitting it into smaller, focused PRs.

---

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

[0cc4m]

This is way too much complexity and code for one vendor.

[virajwad]

Hi @0cc4m I understand your point, it is large (2300+ lines of code). We will circle back and re-look at the code, but any high level feedback you could please offer us how we could improve it?

Intel perf w/ the current flash attention shaders is very unoptimized, especially in long context scenarios. We tried to write custom FA shaders that significantly help our platforms and users, across lots of models, and gives improvement in both prefill and decode stages (MegaPR has FA=ON perf graphs). Any feedback you could offer that could help this PR go in would be very beneficial!

[jeffbolznv]

A few questions from my perspective:

• What is the bottleneck in the current FA shaders and why can't they be tweaked to avoid it?
• What is the purpose/goal of the split-phase FA shaders? What is the size of the temp buffer in practice?
• I'm skeptical of a need for hardcoded head dimensions. Why can't it be parameterized?

[0cc4m]

Yes, basically what Jeff said. Vulkan is meant to be generic, code should be shared between vendors where possible. That way everyone benefits from optimizations and other improvements. Of course, that is not always possible, and device-specific tuning is still necessary. But 2000 lines of additional code specifically for one vendor is excessive and would make maintenance of the backend much harder.

I know Intel's architecture is special and differs in significant ways from AMD and Nvidia, but I also know that many of the issues we are dealing with are down to unoptimized/immature drivers (Linux ANV especially), and that is the main thing that needs to improve, in my opinion.

The other problem is that it has been pretty hard for me to figure out how to write shaders in a way that works well on Intel, especially in regards to subgroups and subgroup operators. I've made great progress, but it's still not up there. I do appreciate that you want to help out with that. I hope we can find a less intrusive way to improve the Vulkan performance for Intel GPUs.

[sliu39]

A few questions from my perspective:

• What is the bottleneck in the current FA shaders and why can't they be tweaked to avoid it?
• What is the purpose/goal of the split-phase FA shaders? What is the size of the temp buffer in practice?
• I'm skeptical of a need for hardcoded head dimensions. Why can't it be parameterized?

@jeffbolznv
Thanks for the review questions.
Overall speaking, the Intel GPU arch diff is why intensive change on FA kernel, with this context, for each question:
• As Intel GPUs have specific GRF/SLM message/Thread/WG arch, if reuse the existing FA kernel, major changes will be needed beyond tweaking, including dispatch/memory access pattern, SLM layout, variable location (SLM to GRF), output rescaling/GEMM PV flow. Potential impact on other vendors.
• The temp buffer size <= (130MB/16k context). Since roofline of FA kernels are counted as 2xGEMM, the purpose for 2 phase design:
For prefill with large head_dim (like 256), the 2xGEMM all calculation bound, so are more suitable and flexible with GEMM style design. For single-phase FA, due to GRF/SLM size limitation, either GRF spill or extra reduce with GEMM(Q^K), both have extra overheads.
For decode, the 2xGEMM all memory bound, can hit better BW utilization with different dispatch pattern. Single-phase FA, with high qk_ratio (4/8), faces dilemma between duplicated KV cache access or too few work groups for dispatch, both result in low memory BW utilization.
• The head_dim 64/96/128 kernels have significant difference with memory access pattern/SLM layout behavior/matrix P reduce logic, unifying them will introduce conditional branches in main loop, the performance trade off may need further evaluation.

[jeffbolznv]

Does this fundamentally come down to Intel having fewer registers per shader core available? In my experience that has been a common problem with trying to choose FA tile size.

Do I understand correctly that the split prefill is really not flash attention, it's just softmax(Q*K) spilled to memory and then the second multiply in a separate dispatch, and this whole thing is chunked to put a bound on the temporary buffer size?

One of the bigger problems having a totally separate path like this is that it won't be easily testable on other devices, which makes it very difficult to maintain.

[sliu39]

Does this fundamentally come down to Intel having fewer registers per shader core available? In my experience that has been a common problem with trying to choose FA tile size.

Exactly, smaller GRF bytes per warp is one of the key diffs, and there are more others including smaller subgroup size and larger warp count per workgroup, SLM reshape for message reduce, efficient WG dispatch/memory access pattern.

Do I understand correctly that the split prefill is really not flash attention, it's just softmax(Q*K) spilled to memory and then the second multiply in a separate dispatch, and this whole thing is chunked to put a bound on the temporary buffer size?

One of the bigger problems having a totally separate path like this is that it won't be easily testable on other devices, which makes it very difficult to maintain.

Agree with the point, split prefill does not strictly comply with FA definition, but softmax operation is still fused as epilog/preprocessing in phase1/phase2. It’s a tradeoff in large head dim case to avoid quad reduce for Q^K or GRF spill with all-in-one FA kernel. For maintenance, other vendors with smaller subgroup/GRF size may also use it or the idea of it, there may be dedicate effort in case of major changes like Vulkan FA op definition upgrade, elsewise would be covered in test-backend-ops and E2E model test.

[0cc4m]

But the most obvious first attempt to make flash attention run better on Intel Battlemage would be to start using coopmat. We currently require 16x16x16 coopmat shape support for Flash Attention, which IIRC Intel GPUs don't offer. They have 8x8x16? Have you tried adapting the existing coopmat FA shader to support this?

[sliu39]

But the most obvious first attempt to make flash attention run better on Intel Battlemage would be to start using coopmat. We currently require 16x16x16 coopmat shape support for Flash Attention, which IIRC Intel GPUs don't offer. They have 8x8x16? Have you tried adapting the existing coopmat FA shader to support this?

Yes, we initially tried to add CM1 with existing FA kernel flow, but due to fundamental GPU arch diffs mentioned in above discussion, we were unable to get consistent prefill win over FA off path unless making significant changes in kernel flow. To avoid impacting existing path, those changes were put in separate path.

[0cc4m]

We can investigate that direction if it's really the only way for Intel, but different shaders for decode, prefill and even specific head sizes would make maintenance too hard. Would it be possible to push it into 2 shaders, one for each stage?

[sliu39]

We can investigate that direction if it's really the only way for Intel, but different shaders for decode, prefill and even specific head sizes would make maintenance too hard. Would it be possible to push it into 2 shaders, one for each stage?

@0cc4m

Thanks for the suggestion, tried experiments of prefill/decode unified 2-phases attention solution, the prefill performance seems to be OK, but decode performance dropped, and need an dedicated decode path to close the gap.
Also want to clarify, the previous suggestion of 2 shaders implementation, does that mean unified 2-phases attention for prefill and decode, or 2 single phase attention shaders, one for prefill and one for decode?

[0cc4m]

I mean I want to keep the number of shaders to maintain minimal. Currently we have 3 for Flash Attention and that works reasonably well across most devices. The fewer the better.

[virajwad]

Hi @0cc4m We've experimented with a consolidated 2 shader solution, however we are seeing significant regressions. We wanted to share the data with you. We looked at the perf in two aspects on an Arrow Lake (ARL-H) Xe2 Intel platform, covering the same models as our mega-pr:

1. Compare 2 shader solution vs llama.cpp master as baseline
2. Compare 2 shader solution vs our mega pr (answer q: can we maintain most of our perf gains when consolidating?)

Prefill

Aspect 1 - We see good perf increases in FA=ON case for 8K input tokens, no FA=ON regressions.

Testing aspect 2, we are able to maintain most prefill gains from our optimizations. Some models gain perf and some lose compared to our mega pr, but average is 1.03x

Decode

Aspect 1 - Significant perf loss across every model we tested, so 2 shader solution causes our decode to be much worse than current master.

Testing aspect 2, the decode perf on average is 0.38x compared to our mega pr optimizations in flash attention. Here are the ratios:

We have a hunch that if we increase to 3 shader solution, we may be able to solve the decode regression (decode will be split-phase), but we need time to experiment and measure once more. @0cc4m could we please get your early thoughts on increasing the shader count? We still want to respect the ask to consolidate this PR a lot, but we are having difficulties porting these optimizations on 2 shaders.

[virajwad]

Hi @0cc4m just checking on the above! Thank you

[0cc4m]

In general, the fewer shaders the better. If you need 3 then you need 3, but someone will have to maintain them. Are they coopmat-only? My second concern is that I still have no access to any Battlemage hardware, so I cannot test anything that doesn't run on my A770.

[bbharti]

In general, the fewer shaders the better. If you need 3 then you need 3, but someone will have to maintain them. Are they coopmat-only? My second concern is that I still have no access to any Battlemage hardware, so I cannot test anything that doesn't run on my A770.
Hi @0cc4m we should be able to maintain the changes. Let me suggest a name(s) by tomorrow. Regarding hardware, definitely you should have better machines. I will try to secure one, please let me know where to ship.

[0cc4m]

If you can help out with hardware that would be good, you can send me a mail with details.