Georgi Gerganov
70aff25250
* 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: https://github.com/ggml-org/llama.cpp/pull/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
30 lines
1.5 KiB
Markdown
30 lines
1.5 KiB
Markdown
You are a coding agent. Here are some very important rules that you must follow:
|
|
|
|
General:
|
|
- Be very precise and concise when writing code, comments, explanations, etc.
|
|
- If an inline comment exceeds 2 lines, replace it with: `// note: TODO LATER`
|
|
- PR and commit titles format: `<module> : <title>`. Lookup recents for examples
|
|
- Don't try to build or run the code unless you are explicitly asked to do so
|
|
- Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources
|
|
|
|
Coding:
|
|
- When in doubt, always refer to the CONTRIBUTING.md file of the project
|
|
- In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments
|
|
- When referencing issues or PRs in comments, use the format:
|
|
- C/C++ code: `// ref: <url>`
|
|
- Other (CMake, etc.): `# ref: <url>`
|
|
|
|
Pull requests (PRs):
|
|
- New branch names are prefixed with "gg/"
|
|
- Before opening a pull request, ask the user to confirm the description
|
|
- When creating a pull request, look for the repository's PR template and follow it
|
|
- For the AI usage disclosure section, write "YES. pi:llama.cpp/[MODEL]"
|
|
- Ask the user to tell you what model was used and write it in place of [MODEL]
|
|
- Always create the pull requests in draft mode
|
|
|
|
Commits:
|
|
- On every commit that you make, include a "Assisted-by: pi:llama.cpp/[MODEL]" tag
|
|
- Do not explicitly set the git author in commits - rely on the default git config
|
|
- Always use `--no-gpg-sign` when committing
|
|
- Never `git push` without explicit confirmation from the user
|