Commit 553d2c2e7 for llama.cpp
commit 553d2c2e7ec67307b44d6264ded5e3d6fce5b1b9
Author: Leon <118652405+LeON-Nie-code@users.noreply.github.com>
Date: Thu Nov 5 16:29:14 2026 +0800
webgpu : add MEAN support (#31591)
Assisted-by: Codex
diff --git a/docs/ops.md b/docs/ops.md
index 3e14b06cb..351093cfa 100644
--- a/docs/ops.md
+++ b/docs/ops.md
@@ -71,7 +71,7 @@ Legend:
| LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
| LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
+| MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
| MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
| MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | 🟡 |
diff --git a/docs/ops/WebGPU.csv b/docs/ops/WebGPU.csv
index d402092c4..b57710424 100644
--- a/docs/ops/WebGPU.csv
+++ b/docs/ops/WebGPU.csv
@@ -12278,13 +12278,23 @@
"WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,2,1,3]","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,3,2,1]","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,1,3,2]","support","1","yes","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[10,5,4,3]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[33,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[33,256,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32769,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32,256,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32768,1,1,1]","support","0","no","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[10,5,4,3],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[33,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[33,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32769,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32768,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=1,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=0,slice=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=1,slice=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1,7,3,2],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1023,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1024,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1025,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[17,65537,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[512,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[281,1536,1,1],permute=0,slice=0","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM","type=f32,ne=[33,1,1,1]","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM","type=f32,ne=[33,1024,1,1]","support","1","yes","WebGPU"
"WebGPU: WebGPU","SUM","type=f32,ne=[33,256,1,1]","support","1","yes","WebGPU"
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
index d1cd0818e..85d51a018 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -1228,7 +1228,7 @@ class ggml_webgpu_shader_lib {
wgpu::Device device;
pre_wgsl::Preprocessor preprocessor;
- std::unordered_map<int, webgpu_pipeline> sum_rows_pipelines; // key is fixed, no variants yet
+ std::unordered_map<int, webgpu_pipeline> sum_rows_pipelines; // key is whether to compute the mean
std::unordered_map<int, webgpu_pipeline> argmax_pipelines; // key is vec4
std::unordered_map<int, webgpu_pipeline> argsort_pipelines; // key is order
std::unordered_map<int, webgpu_pipeline> argsort_merge_pipelines; // key is order
@@ -1319,16 +1319,20 @@ class ggml_webgpu_shader_lib {
ggml_webgpu_shader_lib(wgpu::Device device) { this->device = device; }
webgpu_pipeline get_sum_rows_pipeline(const ggml_webgpu_shader_lib_context & context) {
- auto it = sum_rows_pipelines.find(1);
+ const bool mean = context.dst->op == GGML_OP_MEAN;
+ auto it = sum_rows_pipelines.find(mean);
if (it != sum_rows_pipelines.end()) {
return it->second;
}
std::vector<std::string> defines;
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+ if (mean) {
+ defines.push_back("MEAN");
+ }
- auto processed = preprocessor.preprocess(wgsl_sum_rows, defines);
- sum_rows_pipelines[1] = ggml_webgpu_create_pipeline(device, processed, "sum_rows");
- return sum_rows_pipelines[1];
+ auto processed = preprocessor.preprocess(wgsl_sum_rows, defines);
+ sum_rows_pipelines[mean] = ggml_webgpu_create_pipeline(device, processed, mean ? "mean" : "sum_rows");
+ return sum_rows_pipelines[mean];
}
webgpu_pipeline get_row_norm_pipeline(const ggml_webgpu_shader_lib_context & context) {
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index ea3d91f3b..6f68d4623 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -3265,7 +3265,8 @@ static webgpu_encoded_op ggml_webgpu_sum_rows(webgpu_context & ctx, ggml_tensor
total_sum ? 0 : (uint32_t) (src->nb[3] / ggml_type_size(src->type)),
total_sum ? static_cast<uint32_t>(ggml_nelements(src)) : (uint32_t) src->ne[0],
total_sum ? 1 : (uint32_t) src->ne[1],
- total_sum ? 1 : (uint32_t) src->ne[2] };
+ total_sum ? 1 : (uint32_t) src->ne[2],
+ total_sum ? 1 : (uint32_t) ggml_nrows(dst) };
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src),
ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, dst) };
@@ -3533,6 +3534,7 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
return ggml_webgpu_cumsum(ctx, src0, node);
case GGML_OP_SUM:
case GGML_OP_SUM_ROWS:
+ case GGML_OP_MEAN:
return ggml_webgpu_sum_rows(ctx, src0, node);
case GGML_OP_CONV_2D:
return ggml_webgpu_conv_2d(ctx, src0, src1, node);
@@ -4837,6 +4839,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
supports_op = op->type == GGML_TYPE_F32 && src0->type == op->type && ggml_is_contiguously_allocated(src0);
break;
case GGML_OP_SUM_ROWS:
+ case GGML_OP_MEAN:
supports_op = op->type == GGML_TYPE_F32 && src0->type == op->type && ggml_is_contiguous_rows(src0);
break;
case GGML_OP_UPSCALE:
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/sum_rows.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/sum_rows.wgsl
index 1ead3ac3f..e4f175262 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/sum_rows.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/sum_rows.wgsl
@@ -15,7 +15,8 @@ struct Params {
ne0: u32,
ne1: u32,
- ne2: u32
+ ne2: u32,
+ nrows: u32
};
@group(0) @binding(2)
@@ -28,7 +29,11 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(num_workgroups) num_wg: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
- var i = wid.x + wid.y * num_wg.x;
+ let row = wid.x + wid.y * num_wg.x;
+ if (row >= params.nrows) {
+ return;
+ }
+ var i = row;
let i3 = i / (params.ne2 * params.ne1);
i = i % (params.ne2 * params.ne1);
let i2 = i / params.ne1;
@@ -51,6 +56,10 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
}
if (lid.x == 0) {
- dst[params.offset_dst + wid.x] = shared_sum[0];
+#ifdef MEAN
+ dst[params.offset_dst + row] = shared_sum[0] / f32(params.ne0);
+#else
+ dst[params.offset_dst + row] = shared_sum[0];
+#endif
}
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 427dce1e0..987d5488b 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12285,6 +12285,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, false));
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, false, true));
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, true));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 1, 7, 3, 2 }));
+ for (int64_t n : {1023, 1024, 1025}) {
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { n, 5, 1, 1 }));
+ }
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 17, 65537, 1, 1 }));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 512, 1, 1, 1 }));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 281, 1536, 1, 1 }));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 1, 1, 1 }));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 1024, 1, 1 }));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 }));
@@ -12303,6 +12310,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 1, 1, 1 }));
test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 1024, 1, 1 }));
test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 256, 1, 1 }));
+ 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_mul_add(GGML_TYPE_F32, {64, 64, 320, 1}));
@@ -13304,6 +13312,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_col2im_1d(GGML_TYPE_F16, 16, 512, 2048, 8, 0));
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, {256, 256, 3, 1}));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, {512, 1, 1, 1}));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, {281, 1536, 1, 1}));
for (int n_token : {1, 512}) {