Repository: ggml-org/llama.cpp
Pull request: Vulkan: Tune Flash Attention for MoE on AMD GPUs

See https://github.com/ggml-org/llama.cpp/issues/17715
Supercedes https://github.com/ggml-org/llama.cpp/pull/18033

This is based on some guesswork on my part, since I'm not that familiar with the Vulkan Flash Attention code yet. Let me know if this makes sense. Performance looks alright.

<details>
<summary>AMD Radeon 8060S</summary

With Coopmat:

| model                          | size      | params  | ngl | fa | test          | t/s (before)    | t/s (after)     | diff  |
|--------------------------------|-----------|---------|-----|----|---------------|-----------------|-----------------|-------|
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512         | 488.62 ± 5.65   | 487.46 ± 4.62   | -0.2% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128         | 43.67 ± 0.17    | 43.29 ± 0.18    | -0.9% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 266.87 ± 2.25   | 269.40 ± 0.56   | +0.9% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 33.05 ± 0.10    | 32.99 ± 0.13    | -0.2% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | pp512         | 627.26 ± 2.75   | 631.85 ± 20.90  | +0.7% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | tg128         | 45.44 ± 0.24    | 45.47 ± 0.05    | +0.1% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 306.56 ± 1.99   | 309.79 ± 2.08   | +1.1% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 33.93 ± 0.15    | 34.50 ± 0.14    | +1.7% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512         | 1261.92 ± 70.65 | 1322.13 ± 30.72 | +4.8% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128         | 117.06 ± 0.26   | 117.71 ± 0.34   | +0.6% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512 @ d8192 | 1102.37 ± 24.31 | 1136.17 ± 9.01  | +3.1% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128 @ d8192 | 110.35 ± 1.02   | 111.08 ± 0.94   | +0.7% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512         | 611.04 ± 1.81   | 633.09 ± 14.23  | +3.6% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128         | 94.54 ± 1.12    | 101.33 ± 0.67   | +7.2% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512 @ d8192 | 277.39 ± 1.15   | 284.11 ± 4.44   | +2.4% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128 @ d8192 | 62.02 ± 0.21    | 62.65 ± 0.18    | +1.0% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512         | 956.45 ± 43.81  | 990.88 ± 22.63  | +3.6% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128         | 74.06 ± 0.52    | 75.84 ± 0.15    | +2.4% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512 @ d8192 | 635.69 ± 7.64   | 667.39 ± 4.65   | +5.0% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128 @ d8192 | 64.86 ± 0.82    | 66.59 ± 0.25    | +2.7% |

Without coopmat:

| model                          | size      | params  | ngl | fa | test          | t/s (before)    | t/s (after)     | diff  |
|--------------------------------|-----------|---------|-----|----|---------------|-----------------|-----------------|-------|
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512         | 743.42 ± 17.17  | 812.31 ± 7.25   | +9.3% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128         | 43.37 ± 0.13    | 44.12 ± 0.19    | +1.7% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 166.39 ± 0.82   | 167.35 ± 0.25   | +0.6% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 32.51 ± 0.13    | 32.79 ± 0.04    | +0.9% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | pp512         | 813.98 ± 16.18  | 868.31 ± 13.72  | +6.7% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | tg128         | 45.00 ± 0.04    | 46.18 ± 0.11    | +2.6% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 169.09 ± 1.41   | 170.72 ± 0.31   | +1.0% |
| llama 8B Q4_0                  | 4.33 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 33.14 ± 0.07    | 33.72 ± 0.09    | +1.8% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512         | 1666.99 ± 75.09 | 1707.99 ± 45.57 | +2.5% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128         | 116.60 ± 0.34   | 117.58 ± 1.01   | +0.8% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512 @ d8192 | 1329.74 ± 6.17  | 1344.71 ± 20.99 | +1.1% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128 @ d8192 | 110.72 ± 0.87   | 110.19 ± 0.45   | -0.5% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512         | 548.95 ± 3.79   | 584.59 ± 13.05  | +6.5% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128         | 93.77 ± 0.50    | 101.50 ± 0.36   | +8.2% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512 @ d8192 | 172.44 ± 0.44   | 170.28 ± 0.59   | -1.3% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128 @ d8192 | 61.79 ± 0.29    | 62.44 ± 0.26    | +1.1% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512         | 1153.83 ± 43.07 | 1166.10 ± 30.29 | +1.1% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128         | 72.02 ± 0.34    | 75.37 ± 0.44    | +4.7% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512 @ d8192 | 671.43 ± 4.57   | 679.51 ± 3.69   | +1.2% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128 @ d8192 | 67.25 ± 0.64    | 68.46 ± 0.34    | +1.8% |

</details>

<details>
<summary>AMD Radeon Pro VII</summary>

| model                          | size      | params  | ngl | fa | test          | t/s (before)   | t/s (after)    | diff   |
|--------------------------------|-----------|---------|-----|----|---------------|----------------|----------------|--------|
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512         | 680.94 ± 1.48  | 665.84 ± 0.73  | -2.2%  |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128         | 76.77 ± 0.23   | 80.15 ± 0.27   | +4.4%  |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 168.38 ± 0.23  | 167.81 ± 0.31  | -0.3%  |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 51.26 ± 0.07   | 50.91 ± 0.01   | -0.7%  |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | pp512         | 184.39 ± 0.23  | 181.04 ± 0.49  | -1.8%  |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | tg128         | 24.62 ± 0.03   | 24.95 ± 0.02   | +1.3%  |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | pp512 @ d8192 | 31.45 ± 0.19   | 31.41 ± 0.21   | -0.1%  |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | tg128 @ d8192 | 11.58 ± 0.00   | 11.55 ± 0.01   | -0.3%  |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512         | 1364.56 ± 2.24 | 1311.74 ± 2.57 | -3.9%  |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128         | 116.31 ± 0.19  | 117.18 ± 0.07  | +0.7%  |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512 @ d8192 | 1077.54 ± 5.86 | 1046.76 ± 6.43 | -2.9%  |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128 @ d8192 | 110.14 ± 0.23  | 109.77 ± 0.16  | -0.3%  |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512         | 590.86 ± 3.50  | 596.91 ± 2.34  | +1.0%  |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128         | 84.25 ± 0.09   | 96.28 ± 0.07   | +14.3% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512 @ d8192 | 162.30 ± 0.18  | 161.83 ± 0.23  | -0.3%  |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128 @ d8192 | 57.34 ± 0.09   | 57.31 ± 0.09   | -0.1%  |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512         | 1171.81 ± 7.10 | 1191.57 ± 4.13 | +1.7%  |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128         | 131.52 ± 0.13  | 139.36 ± 0.11  | +6.0%  |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512 @ d8192 | 558.73 ± 1.25  | 563.02 ± 1.66  | +0.8%  |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128 @ d8192 | 116.72 ± 0.22  | 119.10 ± 0.23  | +2.0%  |

</details>

<details>
<summary>RTX 3090</summary>

| model                          | size      | params  | ngl | fa | test          | t/s (before)     | t/s (after)      | diff  |
|--------------------------------|-----------|---------|-----|----|---------------|------------------|------------------|-------|
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512         | 4863.26 ± 78.80  | 4807.87 ± 19.28  | -1.1% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128         | 129.86 ± 1.15    | 127.44 ± 0.69    | -1.9% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | pp512 @ d8192 | 2821.90 ± 83.73  | 2776.25 ± 84.02  | -1.6% |
| llama 8B Q4_K - Small          | 4.36 GiB  | 8.03 B  | 99  | 1  | tg128 @ d8192 | 108.18 ± 0.05    | 107.21 ± 0.25    | -0.9% |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | pp512         | 1656.30 ± 7.07   | 1623.96 ± 3.64   | -2.0% |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | tg128         | 43.32 ± 0.13     | 42.82 ± 0.07     | -1.2% |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | pp512 @ d8192 | 1274.57 ± 14.86  | 1260.11 ± 13.05  | -1.1% |
| llama 13B Q5_K - Small         | 15.18 GiB | 23.57 B | 99  | 1  | tg128 @ d8192 | 39.42 ± 0.05     | 38.97 ± 0.04     | -1.1% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512         | 4538.43 ± 24.11  | 4525.94 ± 16.30  | -0.3% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128         | 223.54 ± 0.46    | 221.74 ± 0.70    | -0.8% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | pp512 @ d8192 | 3360.39 ± 92.82  | 3391.57 ± 117.26 | +0.9% |
| granitehybrid 1B Q4_K - Small  | 3.75 GiB  | 6.94 B  | 99  | 1  | tg128 @ d8192 | 206.29 ± 1.05    | 203.94 ± 0.89    | -1.1% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512         | 2929.01 ± 6.71   | 2925.94 ± 33.82  | -0.1% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128         | 184.45 ± 0.72    | 183.88 ± 0.88    | -0.3% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | pp512 @ d8192 | 1885.31 ± 18.20  | 1869.33 ± 14.37  | -0.8% |
| qwen3moe 30B.A3B Q2_K - Medium | 10.48 GiB | 30.53 B | 99  | 1  | tg128 @ d8192 | 143.63 ± 0.32    | 143.19 ± 0.34    | -0.3% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512         | 4070.55 ± 17.57  | 4067.58 ± 41.42  | -0.1% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128         | 183.32 ± 0.73    | 182.58 ± 0.53    | -0.4% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | pp512 @ d8192 | 2926.57 ± 121.41 | 2898.52 ± 119.52 | -1.0% |
| gpt-oss 20B Q8_0               | 11.27 GiB | 20.91 B | 99  | 1  | tg128 @ d8192 | 157.61 ± 0.86    | 157.41 ± 1.05    | -0.1% |

</details>

Changed files:
filename: ggml/src/ggml-vulkan/ggml-vulkan.cpp
patch: @@ -379,18 +379,18 @@ enum FaCodePath {
 };
 
 struct vk_fa_pipeline_state {
-    vk_fa_pipeline_state(uint32_t HSK, uint32_t HSV, bool small_rows, FaCodePath path, bool aligned, bool f32acc)
-        : HSK(HSK), HSV(HSV), small_rows(small_rows), path(path), aligned(aligned), f32acc(f32acc) {}
+    vk_fa_pipeline_state(uint32_t HSK, uint32_t HSV, bool small_rows, bool small_cache, FaCodePath path, bool aligned, bool f32acc)
+        : HSK(HSK), HSV(HSV), small_rows(small_rows), small_cache(small_cache), path(path), aligned(aligned), f32acc(f32acc) {}
 
     uint32_t HSK, HSV;
-    bool small_rows;
+    bool small_rows, small_cache;
     FaCodePath path;
     bool aligned;
     bool f32acc;
 
     bool operator<(const vk_fa_pipeline_state &b) const {
-        return std::tie(HSK, HSV, small_rows, path, aligned, f32acc) <
-               std::tie(b.HSK, b.HSV, b.small_rows, b.path, b.aligned, b.f32acc);
+        return std::tie(HSK, HSV, small_rows, small_cache, path, aligned, f32acc) <
+               std::tie(b.HSK, b.HSV, b.small_rows, b.small_cache, b.path, b.aligned, b.f32acc);
     }
 };
 
@@ -2582,10 +2582,10 @@ static void ggml_vk_wait_events(vk_context& ctx, std::vector<vk::Event>&& events
 static constexpr uint32_t flash_attention_num_small_rows = 32;
 static constexpr uint32_t scalar_flash_attention_num_small_rows = 1;
 
-static uint32_t get_fa_scalar_num_large_rows(uint32_t hsk, uint32_t hsv) {
+static uint32_t get_fa_scalar_num_large_rows(uint32_t hsk, uint32_t hsv, bool small_cache) {
     if (hsv >= 192) {
         return 2;
-    } else if ((hsv | hsk) & 8) {
+    } else if ((hsv | hsk) & 8 || small_cache) {
         return 4;
     } else {
         return 8;
@@ -2607,9 +2607,8 @@ static uint32_t get_fa_num_small_rows(FaCodePath path) {
     }
 }
 
-static std::array<uint32_t, 2> fa_rows_cols(FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows) {
+static std::array<uint32_t, 2> fa_rows_cols(FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows, bool small_cache) {
     GGML_UNUSED(clamp);
-    GGML_UNUSED(hsv);
 
     if (path == FA_SCALAR) {
         if (small_rows) {
@@ -2618,9 +2617,9 @@ static std::array<uint32_t, 2> fa_rows_cols(FaCodePath path, uint32_t hsk, uint3
             if ((hsv | hsk) & 8) {
                 // HSV/HSK not being a multiple of 16 makes D_split smaller, which makes cols_per_iter
                 // larger, and Bc needs to be >= cols_per_thread. 64 is large enough, 32 is not.
-                return {get_fa_scalar_num_large_rows(hsk, hsv), 64};
+                return {get_fa_scalar_num_large_rows(hsk, hsv, small_cache), 64};
             } else {
-                return {get_fa_scalar_num_large_rows(hsk, hsv), 32};
+                return {get_fa_scalar_num_large_rows(hsk, hsv, small_cache), 32};
             }
         }
     }
@@ -2649,8 +2648,8 @@ static std::array<uint32_t, 2> fa_rows_cols(FaCodePath path, uint32_t hsk, uint3
     return {64, 64};
 }
 
-static uint32_t fa_align(FaCodePath path, uint32_t hsk, uint32_t hsv, ggml_type type, bool small_rows) {
-    return fa_rows_cols(path, hsk, hsv, 0, type, small_rows)[1];
+static uint32_t fa_align(FaCodePath path, uint32_t hsk, uint32_t hsv, ggml_type type, bool small_rows, bool small_cache) {
+    return fa_rows_cols(path, hsk, hsv, 0, type, small_rows, small_cache)[1];
 }
 
 static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector<uint32_t>& warptile, bool mul_mat_id, ggml_type src0_type) {
@@ -2992,11 +2991,11 @@ static void ggml_vk_load_shaders(vk_device& device) {
                                        align, disable_robustness, require_full_subgroups, required_subgroup_size);
     };
 
-    auto const &fa_wg_denoms = [&](FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows) -> std::array<uint32_t, 3> {
-        return {fa_rows_cols(path, hsk, hsv, clamp, type, small_rows)[0], 1, 1};
+    auto const &fa_wg_denoms = [&](FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows, bool small_cache) -> std::array<uint32_t, 3> {
+        return {fa_rows_cols(path, hsk, hsv, clamp, type, small_rows, small_cache)[0], 1, 1};
     };
 
-    auto const &fa_spec_constants = [&](FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows) -> std::vector<uint32_t> {
+    auto const &fa_spec_constants = [&](FaCodePath path, uint32_t hsk, uint32_t hsv, uint32_t clamp, ggml_type type, bool small_rows, bool small_cache) -> std::vector<uint32_t> {
         // For large number of rows, 128 invocations seems to work best.
         // For small number of rows (e.g. N==1), 256 works better. But matrix granularity for 256 is 32, so we
         // can't use 256 for D==80.
@@ -3006,7 +3005,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
         uint32_t wg_size = (path == FA_SCALAR || path == FA_COOPMAT1)
                             ? scalar_flash_attention_workgroup_size
                             : ((small_rows && (D % 32) == 0) ? 256 : 128);
-        auto rows_cols = fa_rows_cols(path, hsk, hsv, clamp, type, small_rows);
+        auto rows_cols = fa_rows_cols(path, hsk, hsv, clamp, type, small_rows, small_cache);
 
         // D_split can't be larger than a subgroup because we use subgroupShuffle to reduce it.
         // D_split can't be larger than the LSB of D divided by 4 due to vectorization in the shader.
@@ -3021,21 +3020,22 @@ static void ggml_vk_load_shaders(vk_device& device) {
             uint32_t HSK = fa.first.HSK; \
             uint32_t HSV = fa.first.HSV; \
             bool small_rows = fa.first.small_rows; \
+            bool small_cache = fa.first.small_cache; \
             FaCodePath path = fa.first.path; \
             bool aligned = fa.first.aligned; \
             bool f32acc = fa.first.f32acc; \
             if (path == FAPATH) { \
                 if (aligned) { \
                     if (f32acc) { \
-                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows), fa_align(FAPATH,HSK,HSV,TYPE,small_rows), true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
+                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
                     } else { \
-                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows), fa_align(FAPATH,HSK,HSV,TYPE,small_rows), true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
+                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
                     } \
                 } else { \
                     if (f32acc) { \
-                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f32acc"         #NAMELC, flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows), 1,                                        true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
+                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f32acc"         #NAMELC, flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ##            SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1,                                        true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
                     } else { \
-                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f16acc"         #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows), 1,                                        true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
+                        ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f16acc"         #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len,  flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data,  "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1,                                        true, true, (FAPATH==FA_COOPMAT1 ? 32 : 0));     \
                     } \
                 } \
             } \
@@ -8008,11 +8008,11 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx
     }
 }
 
-static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const uint32_t hsk, uint32_t hsv) {
+static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const uint32_t hsk, uint32_t hsv, bool small_cache) {
     // Needs to be kept up to date on shader changes
     GGML_UNUSED(hsv);
     const uint32_t wg_size = scalar_flash_attention_workgroup_size;
-    const uint32_t Br = get_fa_scalar_num_large_rows(hsk, hsv);
+    const uint32_t Br = get_fa_scalar_num_large_rows(hsk, hsv, small_cache);
     const uint32_t Bc = scalar_flash_attention_Bc;
 
     const uint32_t tmpsh = wg_size * sizeof(float);
@@ -8136,14 +8136,16 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
     uint32_t workgroups_y = (uint32_t)neq2;
     uint32_t workgroups_z = (uint32_t)neq3;
 
+    const bool small_cache = nek1 < 1024;
+
     // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
     // For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
     uint32_t max_gqa;
     switch (path) {
     case FA_SCALAR:
     case FA_COOPMAT1:
         // We may switch from coopmat1 to scalar, so use the scalar limit for both
-        max_gqa = get_fa_scalar_num_large_rows(HSK, HSV);
+        max_gqa = get_fa_scalar_num_large_rows(HSK, HSV, small_cache);
         break;
     case FA_COOPMAT2:
         max_gqa = get_fa_num_small_rows(FA_COOPMAT2);
@@ -8177,7 +8179,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
 
     // with large hsk/hsv, scalar path may need to use small_rows to fit in shared memory
     if (path == FA_SCALAR &&
-        !ggml_vk_flash_attn_scalar_shmem_support(ctx->device, HSK, HSV)) {
+        !ggml_vk_flash_attn_scalar_shmem_support(ctx->device, HSK, HSV, small_cache)) {
         small_rows = true;
     }
 
@@ -8193,7 +8195,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
         v_stride /= 4;
     }
 
-    uint32_t alignment = fa_align(path, HSK, HSV, k->type, small_rows);
+    uint32_t alignment = fa_align(path, HSK, HSV, k->type, small_rows, small_cache);
     bool aligned = (KV % alignment) == 0 &&
                    // the "aligned" shader variant will forcibly align strides, for performance
                    (q_stride & 7) == 0 && (k_stride & 7) == 0 && (v_stride & 7) == 0;
@@ -8205,7 +8207,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
 
     bool f32acc = path == FA_SCALAR || dst->op_params[3] == GGML_PREC_F32;
 
-    vk_fa_pipeline_state fa_pipeline_state(HSK, HSV, small_rows, path, aligned, f32acc);
+    vk_fa_pipeline_state fa_pipeline_state(HSK, HSV, small_rows, small_cache, path, aligned, f32acc);
 
     vk_pipeline pipeline = nullptr;