kanshan/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-09 12:40:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

ggml-webgpu: tune subgroup split (d_split) in flash_attn_vec (#25418)

Browse Source
This commit is contained in:
Masashi Yoshimura
2026-07-09 08:34:19 +09:00
committed by GitHub
parent 92366df30d
commit 32e41fa5b4
2 changed files with 16 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+9 -16
View File
@@ -2821,23 +2821,16 @@ class ggml_webgpu_shader_lib {
variant.resize(variant.size() - (sizeof("_mask") - 1));
variant += "_mask_blk";
}
uint32_t vec_ne = 1u;
if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16 &&
key.common.head_dim_qk == key.common.head_dim_v) {
switch (key.common.head_dim_qk) {
case 64:
case 192:
case 576:
vec_ne = 2u;
break;
case 96:
vec_ne = 4u;
break;
default:
break;
}
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));
}
defines.push_back(std::string("VEC_NE=") + std::to_string(vec_ne) + "u");
defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
variant += "_dsplit" + std::to_string(d_split);
auto pipeline_decisions = std::make_shared<ggml_webgpu_flash_attn_vec_decisions>(decisions);
webgpu_pipeline pipeline =
ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
+7 -7
View File
@@ -39,9 +39,6 @@ enable subgroups;
#define KV_GRANULARITY 8
#define KV_TILE 16
#define WG_SIZE 64
#ifndef VEC_NE
#define VEC_NE 4u
#endif
#define KV_BLOCKS (KV_TILE / KV_GRANULARITY)
@@ -367,11 +364,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
// accumulate q block * k block into registers across the entire KV tile
if (!skip_tile) {
let num_of_threads = subgroup_size / VEC_NE;
let num_of_threads:u32 = D_SPLIT;
let tx = sg_inv_id % num_of_threads;
let ty = sg_inv_id / num_of_threads;
if (subgroup_id == 0u && q_row_start < params.seq_len_q) {
for (var kv_base : u32 = 0u; kv_base < KV_TILE; kv_base += VEC_NE) {
for (var kv_base : u32 = 0u; kv_base < KV_TILE; kv_base += subgroup_size / D_SPLIT) {
let kv_idx = kv_base + ty;
var partial_sum: f32 = 0.0;
let kv_valid = kv_idx < KV_TILE && (kv_tile + kv_idx) < params.seq_len_kv;
@@ -486,15 +483,18 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
if (!skip_tile) {
// we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
// we want to compute O += P * V across the full KV tile
let ne_threads : u32 = VEC_NE;
let ne_threads : u32 = subgroup_size / D_SPLIT;
let nl_threads = max(1u, subgroup_size / ne_threads);
let tx_pv = sg_inv_id % nl_threads;
let ty_pv = sg_inv_id / nl_threads;
if (subgroup_id == 0u && q_row_start < params.seq_len_q) {
for (var vec_col = tx_pv; vec_col < (HEAD_DIM_V / 4u); vec_col += nl_threads) {
var lo = vec4<f32>(0.0, 0.0, 0.0, 0.0);
for (var cc = 0u; cc < KV_TILE / ne_threads; cc += 1u) {
for (var cc = 0u; cc * ne_threads < KV_TILE; cc += 1u) {
let kv_idx = cc * ne_threads + ty_pv;
if (kv_idx >= KV_TILE) {
continue;
}
let v_row = kv_tile + kv_idx;
if (v_row >= params.seq_len_kv) {
continue;