ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts (#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
This commit is contained in:
Masashi Yoshimura
2026-07-31 09:08:40 +03:00
committed by GitHub
parent 5f55650a78
commit 8f4646a63e
7 changed files with 264 additions and 139 deletions
@@ -153,7 +153,6 @@ var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
#define QUANT_SHMEM kv_shmem
#define QUANT_OUT_TYPE f16
#include "quant_inner_loops.tmpl"
#include "flash_attn_quant_staging.tmpl"
#if !defined(K_Q4_0) && !defined(K_Q8_0)
@@ -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.
#ifndef K_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#endif
@@ -333,7 +334,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
workgroupBarrier();
#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.
#ifndef V_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif