ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts - #25956
Conversation
github-actions
Bot
added
ggml
fa637fb to
750096b
Compare
| @@ -270,7 +269,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>, | |||
| local_scores[slot] = FLOAT_MIN; | |||
| } | |||
|
|
|||
| #ifndef KV_DIRECT | |||
| // The tile path stages K/V in shared memory so each tile can be reused across | |||
| // Q_TILE query rows. It therefore does not use the direct path. | |||
There was a problem hiding this comment.
I added this comment since the current implementation allows the direct path only when using subgroup matrices.
| #if defined(K_DIRECT) || defined(V_DIRECT) | ||
| // Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value, | ||
| // so caching it is more efficient, even on the direct path. | ||
| var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>; |
There was a problem hiding this comment.
I confirmed that we can improve performance by sharing only the scale factor of the quantized K/V, in the test of the PR description.
reeselevine
left a comment
reeselevine
left a comment
There was a problem hiding this comment.
thanks for the improvements, sorry for the slow review. Just a couple minor questions I had.
| key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct; | ||
| key.use_sg_matrix = decisions.use_sg_matrix; | ||
| key.common.k_direct &= decisions.use_sg_matrix && key.common.k_type == GGML_TYPE_F16; | ||
| key.common.v_direct &= decisions.use_sg_matrix && key.common.k_type == GGML_TYPE_F16; |
There was a problem hiding this comment.
should this be key.common.v_type?
There was a problem hiding this comment.
also this calculation for whether to use direct path seems to be in contradiction to the new ggml_webgpu_flash_attn_k/v_direct helpers?
There was a problem hiding this comment.
should this be key.common.v_type?
Thanks, yeah this should be v_type.
also this calculation for whether to use direct path seems to be in contradiction to the new ggml_webgpu_flash_attn_k/v_direct helpers?
New helpers ggml_webgpu_flash_attn_k/v_direct expands the supported type from F16 to F16/Q8_0/Q4_0, but the current subgroup matrices FA doesn’t have the direct path for quantized KV cache. So I added this additional F16 check in this condition.
But we can remove this if we implement the direct path for subgroup matrices FA with quantized KV cache, so may I follow in later?
It's Ok, thanks for the review in busy time :) |
750096b to
94aec92
Compare
yomaytk
added
the
merge ready
…ggml-org#25956) * improve fa of quantized kv cache * Fix some bugs and some comments. * fix v type check and some comments * Fix build error caused by rebasing * editorconfig checking pass
…ggml-org#25956) * improve fa of quantized kv cache * Fix some bugs and some comments. * fix v type check and some comments * Fix build error caused by rebasing * editorconfig checking pass
…ggml-org#25956) * improve fa of quantized kv cache * Fix some bugs and some comments. * fix v type check and some comments * Fix build error caused by rebasing * editorconfig checking pass
3 participants
Overview
This PR mainly addresses two things:
flash_attn_vecto the quantized KV cache, which improves TG performance as shown in the following table. It significantly improves performance for the quantized KV cache, and the performance is slightly higher than with f16 KV cache.-ctkand-ctv, and ggml-webgpu already handles the K and V types separately.It also includes some refactoring of the WGSL code.
Performance
machine: M5 Max, model: Qwen3.5-35B-A3B-Q4_K_M.gguf
Requirements