Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml-webgpu: add layer norm ops (#22406)

Browse Source 
* shader(norm): add layer norm ops

* shader(norm): stablize floating point computation with Kahan summation and handle mixed types

* shader(norm): remove the non-contiguous strides

* shader(norm): use the original implementation rather than the kahan summation
This commit is contained in:
Chen Yuan
2026-05-03 20:52:53 -07:00
committed by GitHub
parent e48034dfc9
commit d4b0c22f9e
3 changed files with 107 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+29 -3
View File
@@ -228,11 +228,13 @@ struct ggml_webgpu_get_rows_pipeline_key_hash {
/** Row Norm **/
struct ggml_webgpu_row_norm_pipeline_key {
ggml_op op;
bool inplace;
ggml_op op;
ggml_type src_type;
ggml_type dst_type;
bool inplace;
bool operator==(const ggml_webgpu_row_norm_pipeline_key & other) const {
return op == other.op && inplace == other.inplace;
return op == other.op && src_type == other.src_type && dst_type == other.dst_type && inplace == other.inplace;
}
};
@@ -240,6 +242,8 @@ struct ggml_webgpu_row_norm_pipeline_key_hash {
size_t operator()(const ggml_webgpu_row_norm_pipeline_key & key) const {
size_t seed = 0;
ggml_webgpu_hash_combine(seed, key.op);
ggml_webgpu_hash_combine(seed, key.src_type);
ggml_webgpu_hash_combine(seed, key.dst_type);
ggml_webgpu_hash_combine(seed, key.inplace);
return seed;
}
@@ -1097,6 +1101,8 @@ class ggml_webgpu_shader_lib {
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;
key.src_type = context.src0->type;
key.dst_type = context.dst->type;
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
auto it = row_norm_pipelines.find(key);
@@ -1111,6 +1117,10 @@ class ggml_webgpu_shader_lib {
defines.push_back("RMS_NORM");
variant = "rms_norm";
break;
case GGML_OP_NORM:
defines.push_back("NORM");
variant = "norm";
break;
case GGML_OP_L2_NORM:
defines.push_back("L2_NORM");
variant = "l2_norm";
@@ -1124,6 +1134,22 @@ class ggml_webgpu_shader_lib {
variant += "_inplace";
}
if (key.src_type == GGML_TYPE_F32) {
defines.push_back("SRC_F32");
variant += "_src_f32";
} else if (key.src_type == GGML_TYPE_F16) {
defines.push_back("SRC_F16");
variant += "_src_f16";
}
if (key.dst_type == GGML_TYPE_F32) {
defines.push_back("DST_F32");
variant += "_dst_f32";
} else if (key.dst_type == GGML_TYPE_F16) {
defines.push_back("DST_F16");
variant += "_dst_f16";
}
const uint32_t row_norm_wg_size = 128u;
uint32_t wg_size = std::min(context.max_wg_size, row_norm_wg_size);
defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));