b10532

github-actions released this 21 Aug 00:42

70aff25

metal : dequantize quantized KV to F16 before flash attention (#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: #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

Website:

https://llama.app

Attestations:

https://github.com/ggml-org/llama.cpp/attestations/42016035

macOS/iOS:

macOS Apple Silicon (arm64)
macOS Apple Silicon (arm64, KleidiAI enabled) DISABLED
macOS Intel (x64)
iOS XCFramework

Linux:

Android:

Android arm64 (CPU)

Windows:

openEuler:

DISABLED
openEuler x86 (310p)
openEuler x86 (910b, ACL Graph)
openEuler aarch64 (310p)
openEuler aarch64 (910b, ACL Graph)

UI:

UI

Assets 28