Skip to content

ggml-webgpu: makes the flash attn vec path subgroup-aware #23040

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
reeselevine merged 2 commits into from
May 14, 2026
Merged

ggml-webgpu: makes the flash attn vec path subgroup-aware #23040

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4c0f629
ggml-webgpu: makes the flash attn vec path compile and size its split…
ArberSephirotheca May 14, 2026
d53c0c1
ggml-webgpu: remove the extra max_wg_size >= max_subgroup_size guard.…
ArberSephirotheca May 14, 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
7 changes: 5 additions & 2 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 @@ -770,7 +770,10 @@ inline ggml_webgpu_flash_attn_decisions ggml_webgpu_flash_attn_get_decisions(
(v_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u);
const bool kv_vec_type_supported =
K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0;
const bool use_vec = context.supports_subgroups && (context.src0->ne[1] < 20) && (context.src0->ne[0] % 32 == 0) &&
// Compile with enough invocations to cover the largest reported subgroup.
const bool vec_subgroup_supported = context.max_wg_size >= context.max_subgroup_size;
Comment thread
ArberSephirotheca marked this conversation as resolved.
Outdated
const bool use_vec = context.supports_subgroups && vec_subgroup_supported && (context.src0->ne[1] < 20) &&
(context.src0->ne[0] % 32 == 0) &&

@reeselevine reeselevine May 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the % 32 for? Especially if we’re trying to be subgroup size agnostic.

@reeselevine reeselevine May 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment everywhere 32 is still used

@ArberSephirotheca ArberSephirotheca May 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea we can remove other places where 32 is used.
But for this specific code, I think 32 is not related to subgroup size, but rather the QK head dim size? For f16 we want the head dim to be divisible by 4, and for q4/q8 path, it should be divisible by ggml_blck_size(type), which is 32. Just update the logic.

@reeselevine reeselevine May 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah thanks for the clarification, the code is clearer now too. It seems like llama.cpp does this check up front as well: https://github.com/ggml-org/llama.cpp/blob/master/src/llama-context.cpp#L3275

But I think it's probably good to have the check here too, since ggml is technically independent from llama.cpp.

(context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) &&
kv_vec_type_supported && (K->type != GGML_TYPE_F16 || f16_vec4_aligned) &&
(context.src2->type == K->type);
Expand Down Expand Up @@ -808,7 +811,7 @@ inline ggml_webgpu_flash_attn_decisions ggml_webgpu_flash_attn_get_decisions(
decisions.q_tile = 1u;
decisions.kv_tile = std::max(8u, std::min(32u, max_kv_tile));
decisions.kv_tile = (decisions.kv_tile / 8u) * 8u;
decisions.wg_size = std::max(1u, std::min<uint32_t>(32u, context.max_subgroup_size));
decisions.wg_size = context.max_subgroup_size;
if (decisions.kv_direct) {
decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD);
while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) {
Expand Down
13 changes: 8 additions & 5 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 @@ -1832,7 +1832,7 @@ static webgpu_encoded_op ggml_webgpu_flash_attn(webgpu_context & ctx,
uint32_t blk_nblk1 = 0;
uint32_t blk_batch_count = 0;

const uint32_t vec_nwg_cap = std::max(1u, std::min<uint32_t>(32u, ctx->global_ctx->capabilities.max_subgroup_size));
const uint32_t vec_nwg_cap = std::min<uint32_t>(32u, ctx->global_ctx->capabilities.min_subgroup_size);
uint32_t nwg = 1u;
const uint64_t kv_span = (uint64_t) std::max(1u, decisions->kv_tile);
while ((2u * nwg * kv_span) < (uint64_t) K->ne[1] && nwg < vec_nwg_cap) {
Expand Down Expand Up @@ -1953,8 +1953,11 @@ static webgpu_encoded_op ggml_webgpu_flash_attn(webgpu_context & ctx,
std::vector<uint32_t> reduce_params;
std::vector<wgpu::BindGroupEntry> reduce_entries;
if (use_vec_reduce) {
const uint32_t reduce_wg_size = std::max(
32u, std::min<uint32_t>(nwg * 32u, ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup));
const uint32_t reduce_sg_size = ctx->global_ctx->capabilities.max_subgroup_size;
const uint32_t reduce_wg_size =
std::max(reduce_sg_size, (uint32_t) std::min<uint64_t>(
(uint64_t) nwg * reduce_sg_size,
ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup));
ggml_webgpu_shader_lib_context reduce_shader_ctx = shader_lib_ctx;
reduce_shader_ctx.max_wg_size = reduce_wg_size;
reduce_pipeline = ctx->shader_lib->get_flash_attn_vec_reduce_pipeline(reduce_shader_ctx);
Expand Down Expand Up @@ -3542,8 +3545,8 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) {
const uint32_t kv_tile = decisions.kv_tile;

const uint32_t vec_nwg_cap = std::max(
1u, std::min<uint32_t>(32u, ctx->webgpu_global_ctx->capabilities.max_subgroup_size));
const uint32_t vec_nwg_cap =
std::min<uint32_t>(32u, ctx->webgpu_global_ctx->capabilities.min_subgroup_size);
uint32_t nwg = 1u;
const uint64_t kv_span = (uint64_t) std::max(1u, kv_tile);
while ((2u * nwg * kv_span) < (uint64_t) K->ne[1] && nwg < vec_nwg_cap) {
Expand Down
Loading