vulkan : dequant q8_0 KV once in coopmat1

[Nathanw1014]

Overview

Removing redundant coopmat1 FA dequantization of KV at prefill. Coopmat1 currently Dequants 32 times (once per workgroup), Reorganises the F16 KV to per-head-contiguous in scratch to enable faster memory-bound read by FA.

Closes #25491

Additional information

Scratch size grows with KV cache size, the trade off for inflight memory-usage vs token throughput: (~268 MB @128k).

Patched vs unpatched greedy output (temp 0, seed 1, ~15.5k-token prompt) is byte-identical, no change in results, just speed.

Unpatched:

model	size	params	backend	ngl	type_k	type_v	fa	test	t/s
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	pp512 @ d32768	200.17 ± 3.27
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	tg32 @ d32768	36.71 ± 0.23
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	pp512 @ d65536	99.08 ± 0.50
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	tg32 @ d65536	26.27 ± 0.01

Patched:

model	size	params	backend	ngl	type_k	type_v	fa	test	t/s
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	pp512 @ d32768	281.99 ± 3.16
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	tg32 @ d32768	36.73 ± 0.06
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	pp512 @ d65536	166.17 ± 2.07
qwen3moe 30B.A3B Q6_K	24.53 GiB	30.53 B	Vulkan	999	q8_0	q8_0	1	tg32 @ d65536	27.50 ± 0.11

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes : Claude (Opus 4.8)
  Assisted with benchmarking, analysis, and review; design + implementation directed by me

[thomas9120]

There’s a complimentary issue here that may be worth looking at: #25207 (comment)

[jeffbolznv]

Doesn't this defeat the purpose of using the quantized cache? Seems like you'll end up with worse memory usage than just using f16 in the first place.

I was curious so I measured the perf on my system:

in coopmat1 mode:

before:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 0 -p 512 -r 10 --prio 1 -m c:\models\Qwen3.6-35B-A3B-UD-Q3_K_M.gguf -ctk q8_0 -ctv q8_0 -d 32768,65536
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl | type_k | type_v |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -----: | -----: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d32768 |      2618.84 ± 45.06 |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d65536 |     1476.18 ± 290.20 |

after:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 0 -p 512 -r 10 --prio 1 -m c:\models\Qwen3.6-35B-A3B-UD-Q3_K_M.gguf -ctk q8_0 -ctv q8_0 -d 32768,65536
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat
| model                          |       size |     params | backend    | ngl | type_k | type_v |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -----: | -----: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d32768 |     3222.40 ± 188.52 |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d65536 |      2206.10 ± 92.85 |

f16:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 0 -p 512 -r 10 --prio 1 -m c:\models\Qwen3.6-35B-A3B-UD-Q3_K_M.gguf -d 32768,65536
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | 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 Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   1 |  pp512 @ d32768 |      3374.56 ± 74.28 |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   1 |  pp512 @ d65536 |      2339.73 ± 13.84 |

in coopmat2 mode:

q8_0:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 0 -p 512 -r 10 --prio 1 -m c:\models\Qwen3.6-35B-A3B-UD-Q3_K_M.gguf -ctk q8_0 -ctv q8_0 -d 32768,65536
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2v
| model                          |       size |     params | backend    | ngl | type_k | type_v |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -----: | -----: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d32768 |     5828.22 ± 297.92 |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   q8_0 |   q8_0 |   1 |  pp512 @ d65536 |     4420.88 ± 348.16 |

f16:

Z:\github\jeffbolznv\llama.cpp\build\bin\RelWithDebInfo>llama-bench.exe -fa 1 -n 0 -p 512 -r 10 --prio 1 -m c:\models\Qwen3.6-35B-A3B-UD-Q3_K_M.gguf -d 32768,65536
ggml_vulkan: Found 1 Vulkan devices:
ggml_vulkan: 0 = NVIDIA GeForce RTX 5090 (NVIDIA) | uma: 0 | fp16: 1 | bf16: 1 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: NV_coopmat2v
| model                          |       size |     params | backend    | ngl |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --: | --------------: | -------------------: |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   1 |  pp512 @ d32768 |      6255.83 ± 68.85 |
| qwen35moe 35B.A3B Q3_K - Medium |  15.45 GiB |    34.66 B | Vulkan     |  -1 |   1 |  pp512 @ d65536 |      4943.92 ± 47.37 |

In coopmat2 mode, there's only about a 10% penalty for using q8_0 directly. So I think there's room for improvement in the coopmat1 path without needing to use additional memory.

[Nathanw1014]

Thanks Jeff,

It only dequantizes one f16 layer into scratch at a time, so it keeps the full q8_0 benefit (~6 GB + 256 MB scratch vs ~12 GB for f16 @128k). The one f16 layer in scratch (256 MB @128k) is extra memory, but ~4% on top of the quantized cache. That ~4% (in this PR's current state) is the tradeoff for up to ~2× prefill throughput.

The benefit grows with context depth, and grows further on lower-bandwidth unified-memory hardware. As I understand it, AMD/RADV can't use coopmat2, so coopmat1 seems like the only prefill path.

The current implementation reuses the existing prealloc_x buffer, minimal code change, a memory-for-throughput tradeoff.

A few directions I had in mind as follow-ups (open to your view on any):

• Bounding the scratch. Either, release/shrink it after prefill, or a chunked/streamed scratch to cap the peak.
• The other quant KV types — q4_0, q4_1, q5_0, q5_1, iq4_nl, via the same block -> contiguous mapping.

Are the current tradeoffs ok for this PR or would you suggest the scratch memory usage is addressed first?

[jeffbolznv]

Thanks for the explanation, maybe this is more viable than I thought. Can you explain why the transpose is needed? I would have expected the quantized and dequantized KV cache to have matching layout, and AFAIK the normal layout for the KV cache is fine.

[Nathanw1014]

Removing the redundant per-workgroup dequant means materializing the KV as an f16 scratch which is nearly double q8_0's bytes. Read strided (without transposing), that costs more than the dequant it saved: pp512 @128k regressed to 30.8, below the 45.8 q8_0 baseline. But that write has to happen anyway, so we use it to lay the scratch out for an efficient FA read, which brings it to 93 t/s, ~2× the baseline.

The tensor strides show why. For the f16 K tensor ne = [128, 256, 4] (HS, KV, n_head_kv), the byte strides are nb = [2, 1024, 256]. Stepping one KV position for a fixed head jumps nb[1] = HS × n_head_kv × 2 = 128 × 4 × 2 = 1024B, where per-head-contiguous would be HS × 2 = 256B. The transpose lays the scratch per-head-contiguous, so the FA reads sequentially.

The same reorder may help f16 too, it's strided the same way. Although for f16 it'd be a pure extra copy, with no dequant to ride on.

[jeffbolznv]

In either case the innermost dimension is contiguous and whole cache lines. Why does it matter whether the rows are permuted?

[Nathanw1014]

Per-row it's identical, the difference is across rows. One head's rows are 1 KB apart natively vs back-to-back in the scratch, so each workgroup's KV walk is 256-on/768-off instead of a sequential stream. The gaps aren't wasted, the other heads' workgroups want those bytes, but exploiting that takes cache and bandwidth to spare. If I'm reading your f16 number right, the 5090 has both, the stride cost there is negligible. The 8060S has neither, so the stride is paid in full.

Dequanting once means the FA loop reads f16, roughly double q8_0's bytes. Higher-bandwidth cards absorb that even with the rows unpermuted. On lower-bandwidth hardware like Strix Halo the doubled strided reads cost more than the decode they save at depth. The transpose is what lets dequant-once work everywhere. It hands the cost back, on a write we're already making.

To isolate the transpose I benched dequant-once with a matching-layout (strided) scratch against this PR. Same pass, same scratch, only the layout differs:

depth	dequant-once strided	dequant-once coalesced (PR)	transpose gain
d16384	327.6 ±1.5	439.0 ±2.3	1.34×
d32768	183.6 ±1.0	281.0 ±1.2	1.53×
d131072	30.8	93.0	3.02×

Full coverage at 16k/32k, r=5, same session (base = 683f0c7):

config	d16384	d32768
base f16	344.5 ±1.6	181.1 ±1.5
base q8_0, in-shader dequant	316.7 ±2.6	200.8 ±2.0
dequant once, strided	327.6 ±1.5	183.6 ±1.0
dequant once, coalesced (PR)	439.0 ±2.3	281.0 ±1.2

Without the transpose, dequant-once vs base goes +3% @16k, −9% @32k, −33% @128k. My read of that progression: the doubled strided traffic costs nothing until the cache stops absorbing it, and the transpose gain (1.34× -> 1.53× -> 3.02×) is that cost handed back.

ggml_vulkan: 0 = Radeon 8060S Graphics (RADV GFX1151) (radv) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 64 | shared memory: 65536 | int dot: 1 | matrix cores: KHR_coopmat

Ryzen AI Max+ 395 (Strix Halo), 64 GB unified LPDDR5X-8000 shared with the CPU. Qwen3 30B A3B Q6_K, -fa 1 -n 0 -p 512 -r 5 -ctk q8_0 -ctv q8_0.

Native strided f16 still edges out the coalesced path in your numbers, so I'd expect strided ≈ coalesced on higher bandwidth cards. Only if it's of interest: the no-transpose variant is a host-only two-hunk change. I can post the diff here or push a branch to my fork.

[jeffbolznv]

It shouldn't matter if it's skipping 768B if it's still fetching whole cache lines. Isn't the cache line size 128B?

[Nathanw1014]

Yes, 128B lines, fully used in both layouts, so the cost sits below the line level.

The transpose wasn't theory-first, the version without it regressed, and that's what sent me measuring.

This box is UMA so I could test where: the CPU shares the DRAM and memory controller. Sequential sustains 73 GB/s. The same data in 256B chunks back-to-back: 70, so chunking is free. The same 256B chunks at 1KB stride: 22. Same chunks, same fully-used lines, only the gaps differ. A 3.2× tax, consistent with the 3.02× the transpose recovers on the GPU at 128k — and it scales with the gap: 512B per 1KB lands at 30. I don't know exactly what's charging for it down there, my guess would be row-buffer locality.