Metal: Optimize q8_0 KV - #25556
Draft
tarruda wants to merge 3 commits into
Draft
Conversation
Add a packed_char4 Q8_0 dequantization helper and use it for the 256-wide Metal vector flash-attention specialization. Loading four signed cache bytes as one packed vector reduces scalar indexing and conversion overhead while leaving other head sizes on the generic helper, since applying the packed path broadly regressed the 64-wide case. Add Qwen 3.5 shaped FLASH_ATTN_EXT coverage for the 256-wide, GQA-16 cache layout, including native and permuted views plus 10K and 20K operation benchmarks. Assisted-by: Codex
Generalize the Metal vector flash-attention kernel to process two adjacent GQA query heads in one threadgroup so they can share each K/V load and dequantization pass. Keep independent query, softmax, output, and reduction state for both heads, and retain the one-head specialization as a fallback. Add a block-parallel Q8_0-to-F16 conversion kernel that dequantizes and transposes K and V into contiguous operation-local scratch. Route eligible one-token, 256-wide, high-GQA attention with at least 1024 visible cache positions through the existing F16 vector attention pipeline using effective F16 buffers and strides. Account for F16 scratch and padding in Metal allocation sizing, insert the producer-to-consumer barrier, and preserve final-chunk padding and partial-result reduction. Extend FLASH_ATTN_EXT correctness and performance coverage for short-context fallback, contiguous and permuted scratch layouts, multiple streams, padding, sinks, ALiBi, logit softcap, and exact Qwen 3.5 397B 10K/20K decode shapes. The persistent KV cache remains Q8_0. Assisted-by: Codex
Move the Q8_0-to-F16 conversion setup ahead of the vector/non-vector flash-attention split so both decode and prompt processing can consume the same temporary K/V representation. Extend materialization eligibility to prompt batches with at least 64 queries while retaining the 256-wide, GQA, and 1024-position gates. Allow the non-vector Metal pipeline helper to select the effective F16 specialization and contiguous F16 stride constants even though the persistent tensors are Q8_0. Pass the scratch buffers and effective strides through final-chunk padding and matrix-oriented attention, and omit quantized threadgroup staging when scratch is active. Add a permuted, unaligned 64-query correctness case and Q8_0/F16 performance cases for 512-query prompts at 10K and 20K depth. Short prompts, small query batches, unsupported shapes, and non-Q8_0 caches continue to use their existing paths. Assisted-by: Codex
github-actions
Bot
added
testing
Contributor
|
Thanks for this — it's a really valuable PR, and you've identified the repeated dequant cost of Q8_0 KV very clearly. Your H=2 GQA approach is quite close to the direction I explored in #23114 (though along a different axis — you're grouping query heads, whereas I was grouping query rows). From the discussion there with @ggerganov , one of the main concerns was the long-term maintainability of hardcoded optimization paths. That's why I've been working on a cross-device, generic approach in #25750 (still a work in progress), and I think the H dimension could potentially fit into that framework later as well. The Q8_0 → F16 materialization also looks like it could be a solid standalone PR. It follows the same idea as what the CUDA backend already does in |
Contributor
Author
I wish I could take credit for this, but my only contribution was prompting GPT.
Feel free to take anything you need out of this PR and submit a new one with the relevant parts. |
|
The performance lift with q8_0 is huge from this PR. I hope the llama.cpp team can prioritize merging it. |
Contributor
Author
|
@anthonyalayo it is pure AI code, and I've not put any effort into splitting or cleaning things up. I opened this with the hope that a metal kernel developer can extract the good parts and submit a proper contribution. |
|
That's fair. I made a fork with the branch so I can utilize it in the meantime. |
This was referenced Aug 19, 2026
ggerganov
added a commit
that referenced
this pull request
Aug 20, 2026
* metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: #25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
ggerganov
added a commit
to ggml-org/ggml
that referenced
this pull request
Aug 21, 2026
…27390) * metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org/llama.cpp#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
ggerganov
added a commit
to ggml-org/whisper.cpp
that referenced
this pull request
Aug 21, 2026
…27390) * metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org/llama.cpp#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
Contributor
Author
|
Trying the same benchmarks I posted on #25555 using the latest master branch, I got much better prompt processing numbers: Token generation still takes a big hit, so I'm leaving this open for now |
fishlikeX
pushed a commit
to fishlikeX/sm70-attn
that referenced
this pull request
Aug 23, 2026
* metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org/llama.cpp#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
fishlikeX
pushed a commit
to fishlikeX/sm70-attn
that referenced
this pull request
Aug 23, 2026
* metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org/llama.cpp#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
therealkenc
pushed a commit
to therealkenc/llama.cpp
that referenced
this pull request
Aug 24, 2026
…rg#27390) * metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
ravel7524
pushed a commit
to ravel7524/llama.cpp
that referenced
this pull request
Aug 30, 2026
…rg#27390) * metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16<block_t, QK, deq_t4x4>: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: ggml-org#25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Suggesting some GPT 5.6 Sol improvements made to Metal KV q8_0
Additional information
More context in this thread: #25555 (comment)
Requirements