metal : remove the in-kernel dequantization from the FA kernels - #27398
Draft
ggerganov wants to merge 2 commits into
Draft
metal : remove the in-kernel dequantization from the FA kernels#27398ggerganov wants to merge 2 commits into
ggerganov wants to merge 2 commits into
Conversation
github-actions
Bot
added
ggml
Resolve the TODO in test_flash_attn_ext: the branch that creates V as a sub-view of K (MLA-based models) was hardcoded for the 576/512 head shapes. Add a v_is_view_of_k test case parameter (default false) and select the sub-view branch on it; the existing 576/512 (DeepSeek MLA) cases now pass it explicitly, so the test coverage is unchanged. Also add more V-is-sub-view-of-K cases: the 320/256 (Mistral4 MLA) and 192/128 head shapes, and full views with equal head sizes (128/128 F16, 64/64 q8_0). Assisted-by: pi:llama.cpp/Qwen3.8-27B
ggerganov
force-pushed
the
gg/tbo-v-is-view-of-k
branch
from
August 20, 2026 10:54
9126f3a to
242dbc8
Compare
the quantized KV cache is dequantized to F16 before flash attention, so the in-kernel dequantization paths (quantized K/V branches, dequant template parameters, k4x4_t/v4x4_t thread types and the quantized kernel instantiations) are dead code now this removes: - the quantized K and V branches from kernel_flash_attn_ext_impl - the quantized K and V branches from kernel_flash_attn_ext_vec - the dequant template parameters from both FA kernels - 125 quantized kernel instantiations (non-vec and vec) - the is_q shared memory term from the host FATTN_SMEM calculation the vec bf16 kernels used the f16 thread types with a bfloat4 device type to dequantize in-kernel; they now use a new FA_TYPES_BF macro with bfloat thread types and read the bfloat device memory directly Assisted-by: pi:llama.cpp/Qwen3.8-27B
ggerganov
force-pushed
the
gg/metal-fa-remove-quant-kv
branch
from
August 20, 2026 10:58
06337b6 to
1134adc
Compare
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
1 participant
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
Follow-up to #27390. Now that quantized KV is dequantized to F16 before flash attention, the in-kernel dequantization paths in
kernel_flash_attn_ext_impl/kernel_flash_attn_ext_vecare dead code. This removes:kd4x4_t/vd4x4_t,nl_k/nl_v,deq_k/deq_vand the_t4variants)k4x4_t/v4x4_tthread types from theFA_TYPESmacrossk/svthreadgroup scratch and theis_qshared-memory term on the host sideThe vec bf16 kernels previously differed from the f16 ones only in the in-kernel dequant device type (
bfloat4dequantized tohalfthreads); they now use a newFA_TYPES_BFmacro withbfloatthread types and read the bfloat device memory directly.Additional info
The
dequantize_*functions are kept - they are still used by the mul_mm kernels and by thekernel_flash_attn_ext_kv_f16dequant kernel from #27390.Requirements