Skip to content

[ggml-webgpu] Implement 2D workgroups for scale, binary, and unary ops #24044

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 22 commits into from
Jun 8, 2026
Merged

[ggml-webgpu] Implement 2D workgroups for scale, binary, and unary ops #24044

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
df60497
Merge
nikhilJain17 Jan 30, 2026
5ae7583
Merge
nikhilJain17 Feb 2, 2026
70238e0
merge
nikhilJain17 Feb 20, 2026
9f9d064
Merge with master
nikhilJain17 Feb 24, 2026
e93b40d
Merge
nikhilJain17 Mar 4, 2026
52b381b
Merge
nikhilJain17 Mar 5, 2026
5904ae4
merge with master
nikhilJain17 Mar 10, 2026
ae076c0
merge with master
nikhilJain17 Mar 14, 2026
de413da
Merge
nikhilJain17 Mar 18, 2026
2c32cdc
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 Mar 19, 2026
19438c1
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 Mar 30, 2026
cf82d32
Only run webgpu CI on my fork
nikhilJain17 Mar 30, 2026
76662da
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 Apr 16, 2026
e92f013
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 Apr 18, 2026
931d5d1
Add webgpu only workflow
nikhilJain17 Apr 18, 2026
69bdd8d
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 May 11, 2026
bb6ed31
Merge branch 'master' of https://github.com/nikhilJain17/llama.cpp
nikhilJain17 May 20, 2026
33763d7
fix
nikhilJain17 Jun 2, 2026
d4b5e5b
Implement 2d workgroups for more operations
nikhilJain17 Jun 2, 2026
8c263fd
fix
nikhilJain17 Jun 2, 2026
3610e6b
Fix type
nikhilJain17 Jun 2, 2026
771a0e3
Move back to global_invocation_id
nikhilJain17 Jun 3, 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
32 changes: 20 additions & 12 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 @@ -621,10 +621,11 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx,
uint32_t value,
size_t offset,
size_t size) {
std::vector<uint32_t> params = { (uint32_t) offset, (uint32_t) size, value };
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry(0, buf, 0, buf.GetSize()) };
size_t bytes_per_wg = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup * ctx->capabilities.memset_bytes_per_thread;
uint32_t wg_x = CEIL_DIV(size + 3, bytes_per_wg);
std::vector<uint32_t> params = { (uint32_t) offset, (uint32_t) size, value };
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry(0, buf, 0, buf.GetSize()) };
size_t bytes_per_wg =
ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup * ctx->capabilities.memset_bytes_per_thread;
uint32_t wg_x = CEIL_DIV(size + 3, bytes_per_wg);

ctx->queue.WriteBuffer(ctx->memset_params_buf, 0, params.data(), params.size() * sizeof(uint32_t));

Expand Down Expand Up @@ -1362,7 +1363,7 @@ static webgpu_encoded_op ggml_webgpu_get_rows(webgpu_context & ctx,
shader_lib_ctx.src0 = src;
shader_lib_ctx.src1 = nullptr;
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;

@CISC CISC Jun 5, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Still conflicting versions of clang-format?

@yomaytk yomaytk Jun 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Hmm, my environment (clangd 22.1.6) produces the same format changes as this PR, plus some additional format changes in this file.
@nikhilJain17 , what version of clang-format or clangd are you using? (see #23760 (comment))

@nikhilJain17 nikhilJain17 Jun 6, 2026

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.

I'm on 21.1.8

@yomaytk yomaytk Jun 7, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks! For reference, since the formatting output seems to differ slightly between versions 21 and 22, it might be good to use version 22 or later for future PRs.

@nikhilJain17 nikhilJain17 Jun 7, 2026

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.

Sounds good, will do


webgpu_pipeline pipeline = ctx->shader_lib->get_get_rows_pipeline(shader_lib_ctx);
auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
Expand Down Expand Up @@ -2103,8 +2104,10 @@ static webgpu_encoded_op ggml_webgpu_unary_op(webgpu_context & ctx, ggml_tensor
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, dst));
}

uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x);
uint32_t wg_x, wg_y;
uint32_t total_wg = CEIL_DIV(ggml_nelements(dst), decisions->wg_size);
compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}

static webgpu_encoded_op ggml_webgpu_binary_op(webgpu_context & ctx,
Expand Down Expand Up @@ -2178,8 +2181,10 @@ static webgpu_encoded_op ggml_webgpu_binary_op(webgpu_context & ctx,
}
}

uint32_t wg_x = CEIL_DIV(ne, decisions->wg_size);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x);
uint32_t wg_x, wg_y;
uint32_t total_wg = CEIL_DIV(ggml_nelements(dst), decisions->wg_size);
compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}

static webgpu_encoded_op ggml_webgpu_add_id(webgpu_context & ctx,
Expand Down Expand Up @@ -2607,8 +2612,10 @@ static webgpu_encoded_op ggml_webgpu_scale(webgpu_context & ctx, ggml_tensor * s
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, dst));
}

uint32_t wg_x = CEIL_DIV(ggml_nelements(dst), decisions->wg_size);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x);
uint32_t wg_x, wg_y;
uint32_t total_wg = CEIL_DIV(ggml_nelements(dst), decisions->wg_size);
compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}

static webgpu_encoded_op ggml_webgpu_soft_max(webgpu_context & ctx,
Expand Down Expand Up @@ -3712,7 +3719,8 @@ static ggml_guid_t ggml_backend_webgpu_guid(void) {

static void ggml_webgpu_init_memset_pipeline(webgpu_global_context & ctx) {
// we use the maximum workgroup size for the memset pipeline
size_t max_threads = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup * ctx->capabilities.limits.maxComputeWorkgroupsPerDimension;
size_t max_threads = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup *
ctx->capabilities.limits.maxComputeWorkgroupsPerDimension;
// Size the bytes_per_thread so that the largest buffer size can be handled
ctx->capabilities.memset_bytes_per_thread =
CEIL_DIV(ctx->capabilities.limits.maxStorageBufferBindingSize, max_threads);
Expand Down
13 changes: 8 additions & 5 deletions ggml/src/ggml-webgpu/wgsl-shaders/binary.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,10 +130,13 @@ fn update(dst_i: u32, src0_i: u32, src1_i: u32) {
}

@compute @workgroup_size(WG_SIZE)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
if (gid.x < params.ne) {
let src0_i = params.offset_src0 + src0_index(gid.x);
let src1_i = params.offset_src1 + src1_index(gid.x);
update(params.offset_dst + gid.x, src0_i, src1_i);
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_wg: vec3<u32>) {
let threads_per_group = u32(WG_SIZE);
let i = gid.x + (num_wg.x * threads_per_group) * gid.y;
if (i < params.ne) {
let src0_i = params.offset_src0 + src0_index(i);
let src1_i = params.offset_src1 + src1_index(i);
update(params.offset_dst + i, src0_i, src1_i);
}
}
10 changes: 6 additions & 4 deletions ggml/src/ggml-webgpu/wgsl-shaders/scale.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,12 +43,14 @@ struct Params {
var<storage, read_write> src: array<f32>;

@compute @workgroup_size(WG_SIZE)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
if (gid.x >= params.ne) {
fn main(
@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_wg: vec3<u32>) {
let threads_per_group = u32(WG_SIZE);
var i = gid.x + (num_wg.x * threads_per_group) * gid.y;
if (i >= params.ne) {
return;
}

var i = gid.x;
let i3 = i / (params.ne2 * params.ne1 * params.ne0);
i = i % (params.ne2 * params.ne1 * params.ne0);
let i2 = i / (params.ne1 * params.ne0);
Expand Down
11 changes: 7 additions & 4 deletions ggml/src/ggml-webgpu/wgsl-shaders/unary.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,14 @@ fn erf_approx(x: TYPE) -> TYPE {
}

@compute @workgroup_size(WG_SIZE)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
if (gid.x >= params.ne) {
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(num_workgroups) num_wg: vec3<u32>) {
let threads_per_group = u32(WG_SIZE);
let flat_i = gid.x + (num_wg.x * threads_per_group) * gid.y;
if (flat_i >= params.ne) {
return;
}
var i = gid.x;
var i = flat_i;
let ne2 = params.ne2;
#ifdef DIAG
let ne1 = params.ne0;
Expand Down Expand Up @@ -205,6 +208,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
#ifdef INPLACE
src[params.offset_src + src_idx] = res;
#else
dst[params.offset_dst + gid.x] = res;
dst[params.offset_dst + flat_i] = res;
#endif
}
Loading