ggml-webgpu: add mulmat with overlapping src0/src1 (e.g., for minimax-01) (#27321)

This commit is contained in:
Masashi Yoshimura
2026-08-19 16:29:33 +09:00
committed by GitHub
parent 0adcc3bb57
commit 5112b9738b
10 changed files with 189 additions and 92 deletions
@@ -954,10 +954,11 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key {
int vectorized;
uint32_t num_cols;
bool use_mmvq;
bool src_overlap;
bool operator==(const ggml_webgpu_mul_mat_vec_pipeline_key & other) const {
return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
num_cols == other.num_cols && use_mmvq == other.use_mmvq;
num_cols == other.num_cols && use_mmvq == other.use_mmvq && src_overlap == other.src_overlap;
}
};
@@ -969,6 +970,7 @@ struct ggml_webgpu_mul_mat_vec_pipeline_key_hash {
ggml_webgpu_hash_combine(seed, key.vectorized);
ggml_webgpu_hash_combine(seed, key.num_cols);
ggml_webgpu_hash_combine(seed, key.use_mmvq);
ggml_webgpu_hash_combine(seed, key.src_overlap);
return seed;
}
};
@@ -977,6 +979,7 @@ struct ggml_webgpu_mul_mat_vec_shader_decisions {
uint32_t wg_size;
uint32_t outputs_per_wg;
uint32_t vec_size;
bool src_overlap = false;
};
struct ggml_webgpu_quantize_q8_pipeline_key {
@@ -998,10 +1001,11 @@ struct ggml_webgpu_mul_mat_pipeline_key {
ggml_type src1_type;
int vectorized;
int use_subgroup_matrix;
bool src_overlap;
bool operator==(const ggml_webgpu_mul_mat_pipeline_key & other) const {
return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized &&
use_subgroup_matrix == other.use_subgroup_matrix;
use_subgroup_matrix == other.use_subgroup_matrix && src_overlap == other.src_overlap;
}
};
@@ -1012,6 +1016,7 @@ struct ggml_webgpu_mul_mat_pipeline_key_hash {
ggml_webgpu_hash_combine(seed, key.src1_type);
ggml_webgpu_hash_combine(seed, key.vectorized);
ggml_webgpu_hash_combine(seed, key.use_subgroup_matrix);
ggml_webgpu_hash_combine(seed, key.src_overlap);
return seed;
}
};
@@ -1034,6 +1039,7 @@ struct ggml_webgpu_mul_mat_shader_decisions {
uint32_t subgroup_matrix_n;
uint32_t mul_mat_wg_size;
bool src_overlap = false;
};
/** MUL_MAT_ID **/
@@ -1950,7 +1956,7 @@ class ggml_webgpu_shader_lib {
return quantize_q8_pipelines[key];
}
webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context) {
webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_mul_mat_vec_pipeline_key key = {};
key.src0_type = context.src0->type;
key.src1_type = context.src1->type;
@@ -1961,6 +1967,7 @@ class ggml_webgpu_shader_lib {
key.num_cols = context.dst->ne[1];
key.use_mmvq =
ggml_webgpu_can_use_mmvq(context.src0, context.src1, context.supports_dot_product, context.vendor);
key.src_overlap = src_overlap;
auto it = mul_mat_vec_pipelines.find(key);
if (it != mul_mat_vec_pipelines.end()) {
@@ -2068,6 +2075,11 @@ class ggml_webgpu_shader_lib {
defines.push_back("Q8_1_T");
}
if (key.src_overlap) {
defines.push_back("SRC_OVERLAP");
variant += "_src_overlap";
}
defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));
defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg));
defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION");
@@ -2089,7 +2101,7 @@ class ggml_webgpu_shader_lib {
return mul_mat_vec_pipelines[key];
}
webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context) {
webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_mul_mat_pipeline_key key = {};
key.src0_type = context.src0->type;
key.src1_type = context.src1->type;
@@ -2098,6 +2110,7 @@ class ggml_webgpu_shader_lib {
1 :
0;
key.use_subgroup_matrix = context.supports_subgroup_matrix;
key.src_overlap = src_overlap;
auto it = mul_mat_fast_pipelines.find(key);
if (it != mul_mat_fast_pipelines.end()) {
@@ -2216,6 +2229,11 @@ class ggml_webgpu_shader_lib {
variant += "_vectorized";
}
if (key.src_overlap) {
defines.push_back("SRC_OVERLAP");
variant += "_src_overlap";
}
if (!key.use_subgroup_matrix) {
defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u");
defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u");