Skip to content

ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts #25956

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 5 commits into from
Jul 31, 2026
Merged

ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts #25956

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2906002
improve fa of quantized kv cache
yomaytk Jul 19, 2026
f98ca4e
Fix some bugs and some comments.
yomaytk Jul 22, 2026
94aec92
fix v type check and some comments
yomaytk Jul 29, 2026
bddd57a
Fix build error caused by rebasing
yomaytk Jul 29, 2026
4cb36d2
editorconfig checking pass
yomaytk Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 45 additions & 30 deletions ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -591,7 +591,8 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
ggml_type dst_type;
uint32_t head_dim_qk;
uint32_t head_dim_v;
bool kv_direct;
bool k_direct;
bool v_direct;
bool kv_overlap;
bool has_mask;
bool has_sinks;
Expand All @@ -600,8 +601,9 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
bool operator==(const ggml_webgpu_flash_attn_common_pipeline_key & other) const {
return q_type == other.q_type && k_type == other.k_type && v_type == other.v_type &&
dst_type == other.dst_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v &&
kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask &&
has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap;
k_direct == other.k_direct && v_direct == other.v_direct && kv_overlap == other.kv_overlap &&
has_mask == other.has_mask && has_sinks == other.has_sinks &&
uses_logit_softcap == other.uses_logit_softcap;
}
};

Expand All @@ -613,7 +615,8 @@ inline void ggml_webgpu_flash_attn_hash_common_pipeline_key(size_t &
ggml_webgpu_hash_combine(seed, key.dst_type);
ggml_webgpu_hash_combine(seed, key.head_dim_qk);
ggml_webgpu_hash_combine(seed, key.head_dim_v);
ggml_webgpu_hash_combine(seed, key.kv_direct);
ggml_webgpu_hash_combine(seed, key.k_direct);
ggml_webgpu_hash_combine(seed, key.v_direct);
ggml_webgpu_hash_combine(seed, key.kv_overlap);
ggml_webgpu_hash_combine(seed, key.has_mask);
ggml_webgpu_hash_combine(seed, key.has_sinks);
Expand Down Expand Up @@ -687,12 +690,13 @@ inline bool ggml_webgpu_flash_attn_float_vec4_aligned(const ggml_tensor * K,
ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
}

inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q,
const ggml_tensor * K,
const ggml_tensor * V,
uint32_t kv_direct_align) {
return K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && (Q->ne[0] % kv_direct_align == 0) &&
(K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
inline bool ggml_webgpu_flash_attn_k_direct(const ggml_tensor * Q, const ggml_tensor * K, uint32_t kv_direct_align) {
return (K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_Q4_0) &&
(Q->ne[0] % kv_direct_align == 0) && (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
}

inline bool ggml_webgpu_flash_attn_v_direct(const ggml_tensor * Q, const ggml_tensor * V, uint32_t kv_direct_align) {
return ggml_webgpu_flash_attn_k_direct(Q, V, kv_direct_align);
}

inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key(
Expand All @@ -706,10 +710,11 @@ inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_co
key.dst_type = context.dst->type;
key.head_dim_qk = (uint32_t) context.src0->ne[0];
key.head_dim_v = (uint32_t) context.src2->ne[0];
key.kv_direct = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align);
key.kv_overlap = kv_overlap;
key.has_mask = context.src3 != nullptr;
key.has_sinks = context.src4 != nullptr;
key.k_direct = ggml_webgpu_flash_attn_k_direct(context.src0, context.src1, kv_direct_align);
key.v_direct = ggml_webgpu_flash_attn_v_direct(context.src0, context.src2, kv_direct_align);
key.kv_overlap = kv_overlap;
key.has_mask = context.src3 != nullptr;
key.has_sinks = context.src4 != nullptr;
key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
return key;
}
Expand Down Expand Up @@ -794,9 +799,13 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(
defines.push_back("LOGIT_SOFTCAP");
variant += "_lgsc";
}
if (key.kv_direct) {
defines.push_back("KV_DIRECT");
variant += "_kvdirect";
if (key.k_direct) {
defines.push_back("K_DIRECT");
variant += "_k_direct";
}
if (key.v_direct) {
defines.push_back("V_DIRECT");
variant += "_v_direct";
}
if (key.kv_overlap) {
defines.push_back("KV_OVERLAP");
Expand All @@ -815,6 +824,12 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(

if (ggml_is_quantized(key.k_type) || ggml_is_quantized(key.v_type)) {
defines.push_back("U32_DEQUANT_HELPERS");
if (ggml_is_quantized(key.k_type)) {
defines.push_back("LOADERS_QUANTIZED_K");
}
if (ggml_is_quantized(key.v_type)) {
defines.push_back("LOADERS_QUANTIZED_V");
}
}

return defines;
Expand Down Expand Up @@ -2792,12 +2807,14 @@ class ggml_webgpu_shader_lib {
ggml_webgpu_flash_attn_pipeline_key key = {};
key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(
context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap);
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.v_type == GGML_TYPE_F16;
key.use_sg_matrix = decisions.use_sg_matrix;

const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u,
key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask,
key.common.k_direct || key.common.v_direct);
GGML_ASSERT(max_kv_tile > 0);

decisions.kv_tile = decisions.use_sg_matrix ?
Expand All @@ -2809,7 +2826,7 @@ class ggml_webgpu_shader_lib {
std::min(context.max_wg_size, std::max(GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE,
GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE * context.max_subgroup_size));

if (key.common.kv_direct) {
if (key.common.k_direct || key.common.v_direct) {
decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD);
while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) {
decisions.kv_tile -= decisions.use_sg_matrix ? context.sg_mat_n : context.min_subgroup_size;
Expand Down Expand Up @@ -2856,9 +2873,9 @@ class ggml_webgpu_shader_lib {
}

ggml_webgpu_flash_attn_vec_decisions decisions = {};
decisions.kv_tile =
ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
decisions.kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
key.common.head_dim_v, key.common.has_mask,
key.common.k_direct || key.common.v_direct);
decisions.wg_size = context.max_subgroup_size;

std::string variant = "flash_attn_vec";
Expand All @@ -2870,12 +2887,10 @@ class ggml_webgpu_shader_lib {
variant += "_mask_blk";
}

uint32_t d_split = context.min_subgroup_size;
if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16) {
const uint32_t D = key.common.head_dim_qk | key.common.head_dim_v;
const uint32_t D_lsb = D & (~(D - 1u));
d_split = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
}
uint32_t d_split = context.min_subgroup_size;
const uint32_t D = key.common.head_dim_qk | key.common.head_dim_v;
const uint32_t D_lsb = D & (~(D - 1u));
d_split = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));

defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
variant += "_dsplit" + std::to_string(d_split);
Expand Down
10 changes: 6 additions & 4 deletions ggml/src/ggml-webgpu/ggml-webgpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3839,7 +3839,8 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
const auto & capabilities = ctx->webgpu_global_ctx->capabilities;
if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V)) {
const bool kv_direct =
ggml_webgpu_flash_attn_kv_direct(Q, K, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
ggml_webgpu_flash_attn_k_direct(Q, K, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH) ||
ggml_webgpu_flash_attn_v_direct(Q, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
const uint32_t kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(
capabilities.limits.maxComputeWorkgroupStorageSize, (uint32_t) Q->ne[0], (uint32_t) V->ne[0],
mask != nullptr, kv_direct);
Expand Down Expand Up @@ -4448,9 +4449,10 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
const uint32_t q_tile =
use_subgroup_matrix ? capabilities.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE;
const uint32_t kv_granularity = use_subgroup_matrix ? capabilities.sg_mat_n : 1u;
const bool kv_direct = use_subgroup_matrix ?
ggml_webgpu_flash_attn_kv_direct(src0, src1, src2, capabilities.sg_mat_k) :
false;
const bool kv_direct = use_subgroup_matrix ?
ggml_webgpu_flash_attn_k_direct(src0, src1, capabilities.sg_mat_k) ||
ggml_webgpu_flash_attn_v_direct(src0, src2, capabilities.sg_mat_k) :
false;
const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
capabilities.limits.maxComputeWorkgroupStorageSize, q_tile, kv_granularity, (uint32_t) src0->ne[0],
(uint32_t) src2->ne[0], op->src[3] != nullptr, kv_direct);
Expand Down
51 changes: 48 additions & 3 deletions ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,12 @@ fn get_byte_i32(value: u32, index: u32) -> i32 {
#endif

#ifdef U32_DEQUANT_HELPERS

fn f16_from_u16(bits: u32) -> f16 {
let packed = unpack2x16float(bits);
return f16(packed[0]);
}

#ifdef DECLARE_BYTE_LOADERS_SRC
fn load_u16_at_src(byte_offset: u32) -> u32 {
let word = src[byte_offset / 4u];
Expand Down Expand Up @@ -36,7 +42,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
let d_bits = (word >> shift) & 0xFFFFu;
return unpack2x16float(d_bits)[0];
}
#endif
#endif // DECLARE_BYTE_LOADERS_SRC

#ifdef DECLARE_BYTE_LOADERS_SRC0
fn load_u16_at_src0(byte_offset: u32) -> u32 {
Expand Down Expand Up @@ -72,8 +78,47 @@ fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
let d_bits = (word >> shift) & 0xFFFFu;
return unpack2x16float(d_bits)[0];
}
#endif
#endif
#endif // DECLARE_BYTE_LOADERS_SRC0

#ifdef LOADERS_QUANTIZED_K
fn load_k_u16_at(byte_offset: u32) -> u32 {
let word = K[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}

fn load_k_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = K[word_idx];
if (shift == 0u) {
return lo;
}
let hi = K[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif // LOADERS_QUANTIZED_K

#ifdef LOADERS_QUANTIZED_V
fn load_v_u16_at(byte_offset: u32) -> u32 {
let word = V[byte_offset / 4u];
let shift = (byte_offset & 2u) * 8u;
return (word >> shift) & 0xFFFFu;
}

fn load_v_u32_at(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 3u) * 8u;
let lo = V[word_idx];
if (shift == 0u) {
return lo;
}
let hi = V[word_idx + 1u];
return (lo >> shift) | (hi << (32u - shift));
}
#endif // LOADERS_QUANTIZED_V

#endif // U32_DEQUANT_HELPERS



Expand Down
27 changes: 13 additions & 14 deletions ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ const FLOAT_MIN: f32 = -1.0e9;
// The number of Q rows processed per workgroup
var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;

#ifndef KV_DIRECT
#if !defined(K_DIRECT) || !defined(V_DIRECT)
const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
// we can reuse the same shmem for K and V since we only need one at a time
var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
Expand Down Expand Up @@ -183,13 +183,12 @@ fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u3
return (*buf)[scalar_index >> 2u];
}

#ifndef KV_DIRECT
#if !defined(K_DIRECT) || !defined(V_DIRECT)
#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)
#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
let k_row = elem_idx / HEAD_DIM_QK;
Expand All @@ -204,7 +203,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
}
#endif

#if !defined(V_Q4_0) && !defined(V_Q8_0)
#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
let v_row = elem_idx / HEAD_DIM_V;
Expand Down Expand Up @@ -296,7 +295,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}

// load k tile into shared memory
#ifndef KV_DIRECT
#ifndef K_DIRECT
load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
#endif

Expand All @@ -306,7 +305,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
// TODO: this loop seems to be the current largest bottleneck
// this bracket exists to scope the lifetime of variables, reducing register pressure
{
#ifdef KV_DIRECT
#ifdef K_DIRECT
let k_block_row = kv_tile + subgroup_id * SG_MAT_N;
var k_global_offset = k_head_offset + k_block_row * params.stride_k1;
#else
Expand All @@ -318,7 +317,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

var q_cur = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, 0u, false, HEAD_DIM_QK);

#ifdef KV_DIRECT
#ifdef K_DIRECT
var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + 0u, true, params.stride_k1);
#else
var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + 0u, true, HEAD_DIM_QK);
Expand All @@ -328,7 +327,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
for (; t + 1u < HEAD_DIM_QK / SG_MAT_K; t += 2u) {
let h0 = t * SG_MAT_K;
var q0 = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h0, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h0, true, params.stride_k1);
#else
var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h0, true, HEAD_DIM_QK);
Expand All @@ -339,7 +338,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

let h1 = (t + 1u) * SG_MAT_K;
var q1g = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h1, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h1, true, params.stride_k1);
#else
var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h1, true, HEAD_DIM_QK);
Expand All @@ -353,7 +352,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
if (t < HEAD_DIM_QK / SG_MAT_K) {
let h = t * SG_MAT_K;
var qn = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h, false, HEAD_DIM_QK);
#ifdef KV_DIRECT
#ifdef K_DIRECT
var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h, true, params.stride_k1);
#else
var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h, true, HEAD_DIM_QK);
Expand All @@ -365,7 +364,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

acc = subgroupMatrixMultiplyAccumulate(q_cur, k_cur, acc);

#ifdef KV_DIRECT
#ifdef K_DIRECT
k_global_offset += num_subgroups * SG_MAT_N * params.stride_k1;
#else
k_block_offset += num_subgroups * SG_MAT_N * HEAD_DIM_QK;
Expand Down Expand Up @@ -436,7 +435,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
}

// load v tile into shared memory
#ifndef KV_DIRECT
#ifndef V_DIRECT
load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
#endif

Expand Down Expand Up @@ -464,7 +463,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
);

// load V submatrix from global or shared memory
#ifdef KV_DIRECT
#ifdef V_DIRECT
let v_block_row = kv_tile + kv_block * SG_MAT_N;
let v_global_offset = v_head_offset + v_block_row * params.stride_v1 + head_dim_block;
var v_sg_mat: subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K> = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(
Expand Down
Loading
Loading