Commit cf659f240 for llama.cpp
commit cf659f240b232e15700b1721bbb43c1d973c52e9
Author: Nikhil Jain <nikhil.jain0987@gmail.com>
Date: Sat Nov 7 00:01:27 2026 -0800
webgpu : Implement group norm kernel (#31490)
* Only run webgpu CI on my fork
* Add webgpu only workflow
* Add WebGPU group_norm kernel
* clean up
* resolve merge conflicts
* parameterize data type; clean up cmake changes
* add nullptr check
* clean up webkit experiments
* more cleanup
* more cleanup
* clang format
* Remove unused feature check for WebGPU
* clean up supports_op condition and update documentation
* add perf tests for group_norm
* clean up doc with script and undo change in test-backend-ops
* Do 2d dispatch of workgroups
* Do 2d dispatch of workgroups
* Unstage unrelated changes
* undo docs churn
* add new test case
* undo small change
* undo small change
diff --git a/docs/ops.md b/docs/ops.md
index e9f4e8adb..4eec04e65 100644
--- a/docs/ops.md
+++ b/docs/ops.md
@@ -62,7 +62,7 @@ Legend:
| GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
| GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
-| GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
diff --git a/docs/ops/WebGPU.csv b/docs/ops/WebGPU.csv
index b57710424..3867ec891 100644
--- a/docs/ops/WebGPU.csv
+++ b/docs/ops/WebGPU.csv
@@ -12309,8 +12309,9 @@
"WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,1024,1,1],permute=0,slice=0","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
-"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[64,64,320,1],num_groups=32,eps=0.000001","support","0","no","WebGPU"
-"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[9,9,1280,1],num_groups=32,eps=0.000001","support","0","no","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[64,64,320,1],num_groups=32,eps=0.000001","support","1","yes","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[9,9,1280,1],num_groups=32,eps=0.000001","support","1","yes","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[3,1,1,65537],num_groups=1,eps=0.000001","support","1","yes","WebGPU"
"WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,1,1],ne_b=[256,16,1,1],stride_dim=-1","support","0","no","WebGPU"
"WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,2,3],ne_b=[256,16,2,3],stride_dim=-1","support","0","no","WebGPU"
"WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,2,3],ne_b=[128,16,2,3],stride_dim=-1","support","0","no","WebGPU"
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
index 85d51a018..30653fccf 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -251,6 +251,31 @@ struct ggml_webgpu_get_rows_pipeline_key_hash {
}
};
+/** Group Norm **/
+
+struct ggml_webgpu_group_norm_pipeline_key {
+ ggml_type src_type;
+ ggml_type dst_type;
+ bool inplace;
+ uint32_t wg_size;
+
+ bool operator==(const ggml_webgpu_group_norm_pipeline_key & other) const {
+ return src_type == other.src_type && dst_type == other.dst_type && inplace == other.inplace &&
+ wg_size == other.wg_size;
+ }
+};
+
+struct ggml_webgpu_group_norm_pipeline_key_hash {
+ size_t operator()(const ggml_webgpu_group_norm_pipeline_key & key) const {
+ size_t seed = 0;
+ ggml_webgpu_hash_combine(seed, key.src_type);
+ ggml_webgpu_hash_combine(seed, key.dst_type);
+ ggml_webgpu_hash_combine(seed, key.inplace);
+ ggml_webgpu_hash_combine(seed, key.wg_size);
+ return seed;
+ }
+};
+
/** Row Norm **/
struct ggml_webgpu_row_norm_pipeline_key {
@@ -1225,6 +1250,9 @@ inline bool ggml_webgpu_can_use_mmvq(const ggml_tensor * src0,
}
class ggml_webgpu_shader_lib {
+ std::unordered_map<ggml_webgpu_group_norm_pipeline_key, webgpu_pipeline, ggml_webgpu_group_norm_pipeline_key_hash>
+ group_norm_pipelines;
+
wgpu::Device device;
pre_wgsl::Preprocessor preprocessor;
@@ -1335,6 +1363,40 @@ class ggml_webgpu_shader_lib {
return sum_rows_pipelines[mean];
}
+ webgpu_pipeline get_group_norm_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ const ggml_webgpu_group_norm_pipeline_key key = {
+ context.src0->type,
+ context.dst->type,
+ ggml_webgpu_tensor_equal(context.src0, context.dst),
+ std::min(context.max_wg_size, 128u),
+ };
+ auto it = group_norm_pipelines.find(key);
+ if (it != group_norm_pipelines.end()) {
+ return it->second;
+ }
+
+ std::vector<std::string> defines = {
+ "WG_SIZE=" + std::to_string(key.wg_size),
+ key.src_type == GGML_TYPE_F16 ? "SRC_F16" : "SRC_F32",
+ key.dst_type == GGML_TYPE_F16 ? "DST_F16" : "DST_F32",
+ };
+ std::string variant =
+ std::string("group_norm_src_") + ggml_type_name(key.src_type) + "_dst_" + ggml_type_name(key.dst_type);
+ if (key.inplace) {
+ GGML_ASSERT(key.src_type == key.dst_type);
+ defines.push_back("INPLACE");
+ variant += "_inplace";
+ }
+ auto processed = preprocessor.preprocess(wgsl_group_norm, defines);
+ auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+ decisions->wg_size = key.wg_size;
+ decisions->inplace = key.inplace;
+ auto pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
+ pipeline.context = decisions;
+ group_norm_pipelines[key] = pipeline;
+ return pipeline;
+ }
+
webgpu_pipeline get_row_norm_pipeline(const ggml_webgpu_shader_lib_context & context) {
ggml_webgpu_row_norm_pipeline_key key = {};
key.op = context.dst->op;
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 768eb96f5..31d67fbf8 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -2759,6 +2759,47 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_rms_norm_mul(webgpu_context
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}
+static webgpu_encoded_op ggml_webgpu_group_norm(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) {
+ const size_t src_type_size = ggml_type_size(src->type);
+ const size_t dst_type_size = ggml_type_size(dst->type);
+ const uint32_t num_groups = (uint32_t) ggml_get_op_params_i32(dst, 0);
+ std::vector<uint32_t> params = {
+ (uint32_t) src->ne[0],
+ (uint32_t) src->ne[1],
+ (uint32_t) src->ne[2],
+ (uint32_t) src->ne[3],
+ num_groups,
+ ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 1)),
+ (uint32_t) (src->nb[1] / src_type_size),
+ (uint32_t) (src->nb[0] / src_type_size),
+ (uint32_t) (src->nb[2] / src_type_size),
+ (uint32_t) (src->nb[3] / src_type_size),
+ (uint32_t) (dst->nb[1] / dst_type_size),
+ (uint32_t) (dst->nb[0] / dst_type_size),
+ (uint32_t) (dst->nb[2] / dst_type_size),
+ (uint32_t) (dst->nb[3] / dst_type_size),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / src_type_size),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / dst_type_size),
+ };
+
+ ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+ shader_lib_ctx.src0 = src;
+ shader_lib_ctx.dst = dst;
+ shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+ webgpu_pipeline pipeline = ctx->shader_lib->get_group_norm_pipeline(shader_lib_ctx);
+ auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+ std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src) };
+ if (!decisions->inplace) {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, dst));
+ }
+ uint32_t wg_x;
+ uint32_t wg_y;
+ compute_2d_workgroups(num_groups * (uint32_t) src->ne[3],
+ 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_row_norm(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) {
std::vector<uint32_t> params = {
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)),
@@ -3487,6 +3528,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
return ggml_webgpu_concat(ctx, src0, src1, node);
case GGML_OP_REPEAT:
return ggml_webgpu_repeat(ctx, src0, node);
+ case GGML_OP_GROUP_NORM:
+ return ggml_webgpu_group_norm(ctx, src0, node);
case GGML_OP_RMS_NORM:
if (ggml_webgpu_can_fuse_rms_norm_mul(cgraph, node_idx)) {
num_encoded_ops = 2;
@@ -4928,6 +4971,18 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
supports_op = max_kv_tile > 0;
break;
}
+ case GGML_OP_GROUP_NORM:
+ {
+ bool valid_type = (src0->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) ||
+ (src0->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F16);
+ // One workgroup per normalization group within each batch item.
+ const uint64_t num_workgroups = uint64_t(ggml_get_op_params_i32(op, 0)) * uint64_t(src0->ne[3]);
+ const uint64_t max_per_dim =
+ ctx->webgpu_global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension;
+ const bool dispatch_fits = num_workgroups <= max_per_dim * max_per_dim;
+ supports_op = valid_type && dispatch_fits;
+ break;
+ }
case GGML_OP_RMS_NORM:
case GGML_OP_NORM:
case GGML_OP_L2_NORM:
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/group_norm.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/group_norm.wgsl
new file mode 100644
index 000000000..a2dcd42c0
--- /dev/null
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/group_norm.wgsl
@@ -0,0 +1,154 @@
+#if defined(SRC_F16) || defined(DST_F16)
+enable f16;
+#endif
+
+#ifdef SRC_F16
+#define SRC_TYPE f16
+#else
+#define SRC_TYPE f32
+#endif
+
+#ifdef DST_F16
+#define DST_TYPE f16
+#else
+#define DST_TYPE f32
+#endif
+
+struct Params {
+ width: u32,
+ height: u32,
+ num_channels: u32,
+ batch_size: u32,
+ num_groups: u32,
+ epsilon: f32,
+ // input strides
+ row_stride: u32,
+ col_stride: u32,
+ channel_stride: u32,
+ batch_stride: u32,
+ // output strides
+ output_row_stride: u32,
+ output_col_stride: u32,
+ output_channel_stride: u32,
+ output_batch_stride: u32,
+ // offsets
+ input_offset: u32,
+ output_offset: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> input: array<SRC_TYPE>;
+#ifdef INPLACE
+@group(0) @binding(1) var<uniform> params: Params;
+#else
+@group(0) @binding(1) var<storage, read_write> output: array<DST_TYPE>;
+@group(0) @binding(2) var<uniform> params: Params;
+#endif
+
+var<workgroup> partials: array<f32, WG_SIZE>;
+var<workgroup> mean: f32;
+var<workgroup> variance: f32;
+
+fn input_index(logical_position: u32, group_base: u32) -> u32 {
+ let col: u32 = logical_position % params.width;
+ let row: u32 = (logical_position / params.width) % params.height;
+ let channel: u32 = (logical_position / params.width) / params.height;
+ return group_base
+ + (row * params.row_stride)
+ + (col * params.col_stride)
+ + (channel * params.channel_stride);
+}
+
+fn output_index(logical_position: u32, group_base: u32) -> u32 {
+ let col: u32 = logical_position % params.width;
+ let row: u32 = (logical_position / params.width) % params.height;
+ let channel: u32 = (logical_position / params.width) / params.height;
+ return group_base
+ + (row * params.output_row_stride)
+ + (col * params.output_col_stride)
+ + (channel * params.output_channel_stride);
+}
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+ @builtin(workgroup_id) workgroup_id: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>,
+ @builtin(local_invocation_index) lane: u32
+) {
+ if (params.num_groups == 0u) {
+ return;
+ }
+ let linear_group_id = workgroup_id.x + workgroup_id.y * num_workgroups.x;
+ let batch_index = linear_group_id / params.num_groups;
+ if (batch_index >= params.batch_size) {
+ return;
+ }
+ let group_index = linear_group_id % params.num_groups;
+ let channels_per_group = params.num_channels / params.num_groups
+ + select(0u, 1u, params.num_channels % params.num_groups != 0u);
+ if (group_index > params.num_channels / max(channels_per_group, 1u)) {
+ return;
+ }
+ let first_channel = group_index * channels_per_group;
+ if (first_channel >= params.num_channels) {
+ return;
+ }
+ let group_channels = min(channels_per_group, params.num_channels - first_channel);
+ let num_elems_per_group = group_channels * params.width * params.height;
+ let batch_base = params.input_offset + batch_index * params.batch_stride;
+ let group_base = batch_base + first_channel * params.channel_stride;
+
+ var partial_sum : f32 = 0.0f;
+ for (var i = lane; i < num_elems_per_group; i += WG_SIZE) {
+ partial_sum += f32(input[input_index(i, group_base)]);
+ }
+ partials[lane] = partial_sum;
+
+ workgroupBarrier();
+
+ if (lane == 0) {
+ var total_sum : f32 = 0.0f;
+ for (var i: u32 = 0u; i < WG_SIZE; i++) {
+ total_sum += partials[i];
+ }
+ mean = total_sum / f32(num_elems_per_group);
+ }
+
+ workgroupBarrier();
+
+ var partial_sq_dev : f32 = 0.0f;
+ for (var i = lane; i < num_elems_per_group; i += WG_SIZE) {
+ let delta = f32(input[input_index(i, group_base)]) - mean;
+ partial_sq_dev += delta * delta;
+ }
+ partials[lane] = partial_sq_dev;
+
+ workgroupBarrier();
+
+ if (lane == 0) {
+ var total_sum : f32 = 0.0f;
+ for (var i = 0u; i < WG_SIZE; i++) {
+ total_sum += partials[i];
+ }
+ variance = total_sum / f32(num_elems_per_group);
+ }
+
+ workgroupBarrier();
+
+ let scale = inverseSqrt(variance + params.epsilon);
+#ifndef INPLACE
+ let output_group_base = params.output_offset
+ + batch_index * params.output_batch_stride
+ + first_channel * params.output_channel_stride;
+#endif
+ for (var i = lane; i < num_elems_per_group; i += WG_SIZE) {
+ let index = input_index(i, group_base);
+ let normalized = (f32(input[index]) - mean) * scale;
+#ifdef INPLACE
+ input[index] = SRC_TYPE(normalized);
+#else
+ output[output_index(i, output_group_base)] = DST_TYPE(normalized);
+#endif
+ }
+
+
+}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 91b0135c2..cd9d2b894 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12415,6 +12415,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 17, 65537, 1, 1 }));
test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {64, 64, 320, 1}));
test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {9, 9, 1280, 1}));
+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {3, 1, 1, 65537}, 1));
test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {64, 64, 320, 1}));
test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {9, 9, 1280, 1}));
test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1));
@@ -13238,6 +13239,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 1, 1}));
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 4, 1}));
+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {64, 64, 320, 1}));
+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {9, 9, 1280, 1}));
+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {3, 1, 1, 65537}, 1));
+
// SNAKE activation fusion at BigVGAN scale (T=7680 = 24 kHz x 320 ms, C=192)
test_cases.emplace_back(new test_snake_fuse(GGML_TYPE_F32, {7680, 192, 1, 1}));
test_cases.emplace_back(new test_snake_fuse(GGML_TYPE_F16, {7680, 192, 1, 1}));