Commit fcac19079 for llama.cpp
commit fcac19079f3f2985fe74ed845f0880e7e9a8030a
Author: Masashi Yoshimura <yoshimura.masashi.frbs@gmail.com>
Date: Wed Oct 14 08:24:23 2026 +0900
webgpu: (Fix CI) fix mul_mat_id with duplicate expert ids (#30451)
* Fix mm-id of webgpu in the case of ids duplication
* formatting
* fix comments to avoid breaking up sentences in comments
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
index 65556f352..6efff9b37 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -1053,6 +1053,22 @@ struct ggml_webgpu_mul_mat_shader_decisions {
/** MUL_MAT_ID **/
+struct ggml_webgpu_mul_mat_id_gather_pipeline_key {
+ uint32_t n_experts;
+
+ bool operator==(const ggml_webgpu_mul_mat_id_gather_pipeline_key & other) const {
+ return n_experts == other.n_experts;
+ }
+};
+
+struct ggml_webgpu_mul_mat_id_gather_pipeline_key_hash {
+ size_t operator()(const ggml_webgpu_mul_mat_id_gather_pipeline_key & key) const {
+ size_t seed = 0;
+ ggml_webgpu_hash_combine(seed, key.n_experts);
+ return seed;
+ }
+};
+
struct ggml_webgpu_mul_mat_id_pipeline_key {
ggml_type src0_type;
ggml_type src1_type;
@@ -1263,12 +1279,15 @@ class ggml_webgpu_shader_lib {
std::unordered_map<ggml_webgpu_mul_mat_pipeline_key, webgpu_pipeline, ggml_webgpu_mul_mat_pipeline_key_hash>
mul_mat_fast_pipelines; // fast mat-mat (reg-tile or subgroup)
std::unordered_map<ggml_webgpu_quantize_q8_pipeline_key, webgpu_pipeline, ggml_webgpu_quantize_q8_pipeline_key_hash>
- quantize_q8_pipelines;
- std::unordered_map<int, webgpu_pipeline> mul_mat_id_gather_pipelines; // key is fixed
+ quantize_q8_pipelines;
+ std::unordered_map<ggml_webgpu_mul_mat_id_gather_pipeline_key,
+ webgpu_pipeline,
+ ggml_webgpu_mul_mat_id_gather_pipeline_key_hash>
+ mul_mat_id_gather_pipelines; // n_experts
std::unordered_map<ggml_webgpu_mul_mat_id_pipeline_key, webgpu_pipeline, ggml_webgpu_mul_mat_id_pipeline_key_hash>
- mul_mat_id_pipelines; // src0_type/src1_type
+ mul_mat_id_pipelines; // src0_type/src1_type
std::unordered_map<ggml_webgpu_mul_mat_id_pipeline_key, webgpu_pipeline, ggml_webgpu_mul_mat_id_pipeline_key_hash>
- mul_mat_id_vec_pipelines; // src0_type/src1_type
+ mul_mat_id_vec_pipelines; // src0_type/src1_type
std::unordered_map<ggml_webgpu_set_rows_pipeline_key, webgpu_pipeline, ggml_webgpu_set_rows_pipeline_key_hash>
set_rows_pipelines;
@@ -2325,20 +2344,26 @@ class ggml_webgpu_shader_lib {
}
webgpu_pipeline get_mul_mat_id_gather_pipeline(const ggml_webgpu_shader_lib_context & context) {
- auto it = mul_mat_id_gather_pipelines.find(1);
+ ggml_webgpu_mul_mat_id_gather_pipeline_key key = {};
+ key.n_experts = (uint32_t) context.src0->ne[2];
+
+ auto it = mul_mat_id_gather_pipelines.find(key);
if (it != mul_mat_id_gather_pipelines.end()) {
return it->second;
}
std::vector<std::string> defines;
- defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+ defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size) + "u");
+ defines.push_back(std::string("N_EXPERT=") + std::to_string(key.n_experts) + "u");
+ defines.push_back(std::string("TILE_N_SLOTS=") +
+ std::to_string(WEBGPU_MUL_MAT_TILE_N * WEBGPU_MUL_MAT_WG_SIZE_N) + "u");
auto processed = preprocessor.preprocess(wgsl_mul_mat_id_gather, defines);
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
decisions->wg_size = context.max_wg_size;
- webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "mul_mat_id_gather");
- pipeline.context = decisions;
- mul_mat_id_gather_pipelines[1] = pipeline;
+ webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "mul_mat_id_gather");
+ pipeline.context = decisions;
+ mul_mat_id_gather_pipelines[key] = pipeline;
return pipeline;
}
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 3fbac9675..15e6c784b 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -1817,6 +1817,19 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id_vec(webgpu_context & ctx,
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}
+// upper bound on the number of slot tiles emitted by mul_mat_id_gather.wgsl
+static uint32_t ggml_webgpu_mul_mat_id_max_tiles(const ggml_tensor * dst) {
+ const uint32_t n_expert = (uint32_t) dst->src[0]->ne[2];
+ const uint32_t n_slots = (uint32_t) (dst->ne[1] * dst->ne[2]); // n_expert_used * n_tokens
+
+ return CEIL_DIV(n_slots, WEBGPU_MUL_MAT_TILE_N * WEBGPU_MUL_MAT_WG_SIZE_N) + std::min(n_expert, n_slots);
+}
+
+// tile count followed by (expert, first slot, n slots) per tile
+static size_t ggml_webgpu_mul_mat_id_tiles_nbytes(const ggml_tensor * dst) {
+ return (1 + 3 * (size_t) ggml_webgpu_mul_mat_id_max_tiles(dst)) * sizeof(uint32_t);
+}
+
static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
ggml_tensor * src0,
ggml_tensor * src1,
@@ -1856,39 +1869,29 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
(uint32_t) (src2->nb[1] / ggml_type_size(src2->type)),
};
- const size_t dst_offset = ggml_webgpu_tensor_offset(dst);
- const size_t gathered_buf_nbytes = src0->ne[2] * src1->ne[2] * sizeof(uint32_t);
+ const size_t dst_offset = ggml_webgpu_tensor_offset(dst);
+ const size_t ids_cache_nbytes = src2->ne[0] * src1->ne[2] * sizeof(uint32_t);
+ const size_t ids_tiles_nbytes = ggml_webgpu_mul_mat_id_tiles_nbytes(dst);
- const size_t gathered_expert_used_align_offset = ROUNDUP_POW2(
+ const size_t ids_cache_align_offset = ROUNDUP_POW2(
dst_offset + ggml_nbytes(dst), ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
- const size_t gathered_tokens_align_offset =
- ROUNDUP_POW2(gathered_expert_used_align_offset + gathered_buf_nbytes,
- ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
- const size_t gathered_count_ids_align_offset =
- ROUNDUP_POW2(gathered_tokens_align_offset + gathered_buf_nbytes,
+ const size_t ids_tiles_align_offset =
+ ROUNDUP_POW2(ids_cache_align_offset + ids_cache_nbytes,
ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
- const size_t gathered_binding_size = ROUNDUP_POW2(gathered_buf_nbytes, WEBGPU_STORAGE_BUF_BINDING_MULT);
- const size_t gathered_count_ids_binding_size =
- ROUNDUP_POW2(src0->ne[2] * sizeof(uint32_t), WEBGPU_STORAGE_BUF_BINDING_MULT);
+ const size_t ids_cache_bind_size = ROUNDUP_POW2(ids_cache_nbytes, WEBGPU_STORAGE_BUF_BINDING_MULT);
+ const size_t ids_tiles_bind_size = ROUNDUP_POW2(ids_tiles_nbytes, WEBGPU_STORAGE_BUF_BINDING_MULT);
// bind group entries for mul_mat_id_gather.wgsl
std::vector<wgpu::BindGroupEntry> gather_entries = {
ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src2), ggml_webgpu_tensor_align_offset(ctx, src2),
ggml_webgpu_tensor_binding_size(ctx, src2)),
- ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(dst), gathered_expert_used_align_offset,
- gathered_binding_size),
- ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(dst), gathered_tokens_align_offset,
- gathered_binding_size),
- ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), gathered_count_ids_align_offset,
- gathered_count_ids_binding_size),
+ ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(dst), ids_cache_align_offset, ids_cache_bind_size),
+ ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(dst), ids_tiles_align_offset, ids_tiles_bind_size),
};
- // n_expert is much less than maxComputeWorkgroupsPerDimension (e.g., n_exeprt=256 at Qwen3.5-35B-A3B)
- const uint32_t gather_wg_x = param_n_expert;
-
dispatches.push_back({
- gather_pipeline, std::move(gather_params), std::move(gather_entries), { gather_wg_x, 1 }
+ gather_pipeline, std::move(gather_params), std::move(gather_entries), { 1, 1 }
});
// params for mul_mat_id.wgsl
@@ -1916,12 +1919,8 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
ggml_webgpu_tensor_binding_size(ctx, src1)),
ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(dst), ggml_webgpu_tensor_align_offset(ctx, dst),
ggml_webgpu_tensor_binding_size(ctx, dst)),
- ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), gathered_expert_used_align_offset,
- gathered_binding_size),
- ggml_webgpu_make_bind_group_entry(4, ggml_webgpu_tensor_buf(dst), gathered_tokens_align_offset,
- gathered_binding_size),
- ggml_webgpu_make_bind_group_entry(5, ggml_webgpu_tensor_buf(dst), gathered_count_ids_align_offset,
- gathered_count_ids_binding_size),
+ ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), ids_cache_align_offset, ids_cache_bind_size),
+ ggml_webgpu_make_bind_group_entry(4, ggml_webgpu_tensor_buf(dst), ids_tiles_align_offset, ids_tiles_bind_size),
};
// Calculate workgroup dimensions
@@ -1930,15 +1929,12 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
auto * main_decisions = static_cast<ggml_webgpu_mul_mat_shader_decisions *>(main_pipeline.context.get());
- uint32_t wg_m;
+ // must match `TILE_N_SLOTS` of mul_mat_id_gather.wgsl
+ GGML_ASSERT(main_decisions->tile_n * main_decisions->wg_size_n == WEBGPU_MUL_MAT_TILE_N * WEBGPU_MUL_MAT_WG_SIZE_N);
- uint32_t tile_m_s = main_decisions->tile_m * main_decisions->wg_size_m;
- uint32_t tile_n_s = main_decisions->tile_n * main_decisions->wg_size_n;
- wg_m = CEIL_DIV(dst->ne[0], tile_m_s);
- uint32_t total_gathered = dst->ne[1] * dst->ne[2];
- uint32_t max_active_experts = std::min((uint32_t) src0->ne[2], total_gathered);
- uint32_t max_wg_n = CEIL_DIV(total_gathered, tile_n_s) + max_active_experts;
- uint32_t total_wg = wg_m * max_wg_n;
+ const uint32_t tile_m_s = main_decisions->tile_m * main_decisions->wg_size_m;
+ const uint32_t wg_m = CEIL_DIV(dst->ne[0], tile_m_s);
+ const uint32_t total_wg = wg_m * ggml_webgpu_mul_mat_id_max_tiles(dst);
compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
@@ -4065,11 +4061,11 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
const ggml_tensor * src0 = tensor->src[0];
const ggml_tensor * src1 = tensor->src[1];
if (src0 && src1) {
- const size_t gathered_size = sizeof(uint32_t) * tensor->src[0]->ne[2] * tensor->src[1]->ne[2];
- const size_t gathered_count_ids_size = sizeof(uint32_t) * tensor->src[0]->ne[2];
- res = ROUNDUP_POW2(
- res + gathered_size * 2 + gathered_count_ids_size +
- ctx->webgpu_global_ctx->capabilities.limits.minStorageBufferOffsetAlignment * 3,
+ const size_t ids_cache_size = sizeof(uint32_t) * tensor->ne[1] * tensor->ne[2];
+ const size_t ids_tiles_size = ggml_webgpu_mul_mat_id_tiles_nbytes(tensor);
+ res = ROUNDUP_POW2(
+ res + ids_cache_size + ids_tiles_size +
+ ctx->webgpu_global_ctx->capabilities.limits.minStorageBufferOffsetAlignment * 2,
WEBGPU_STORAGE_BUF_BINDING_MULT);
}
}
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id.wgsl
index 91039ff25..a3c25f656 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id.wgsl
@@ -38,11 +38,11 @@ struct MulMatIdParams {
@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // [cols, rows, n_expert]
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // [cols, b_ne1, n_tokens]
@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>; // [rows, n_expert_used, n_tokens]
-@group(0) @binding(3) var<storage, read_write> global_gathered_expert_used: array<u32>; // [n_expert][n_tokens]
-@group(0) @binding(4) var<storage, read_write> global_gathered_tokens: array<u32>; // [n_expert][n_tokens]
-@group(0) @binding(5) var<storage, read_write> gathered_count_ids: array<u32>; // [n_expert]
-@group(0) @binding(6) var<uniform> params: MulMatIdParams;
+@group(0) @binding(3) var<storage, read_write> ids_slot_pos: array<u32>; // [n_expert_used * n_tokens]
+@group(0) @binding(4) var<storage, read_write> ids_tiles: array<u32>; // [n_tiles, (expert, first slot, n slots) * n_tiles]
+
+@group(0) @binding(5) var<uniform> params: MulMatIdParams;
fn get_local_n(thread_id: u32) -> u32 {
return thread_id / WORKGROUP_SIZE_M;
@@ -90,44 +90,29 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
let local_m = get_local_m(thread_id);
let local_n = get_local_n(thread_id);
- var expert_idx:u32 = 0xFFFFFFFFu;
- var wg_in_batch:u32 = 0;
- var wg_sum:u32 = 0;
let wg_m_count = (params.m + WORKGROUP_SIZE_M * TILE_M - 1u) / (WORKGROUP_SIZE_M * TILE_M);
let wg_linear = wg_id.y * num_wg.x + wg_id.x;
- for (var i = 0u;i < params.n_expert;i += 1) {
- let wg_n_count = (gathered_count_ids[i] + WORKGROUP_SIZE_N * TILE_N - 1u) / (WORKGROUP_SIZE_N * TILE_N);
- let wg_per_matrix = wg_m_count * wg_n_count;
- if (wg_sum <= wg_linear && wg_linear < wg_sum + wg_per_matrix) {
- expert_idx = i;
- wg_in_batch = wg_linear - wg_sum;
- break;
- }
- wg_sum += wg_per_matrix;
- }
+ // the slot tile of this workgroup, described by mul_mat_id_gather.wgsl
+ let tile = wg_linear / wg_m_count;
+ let wg_m = wg_linear % wg_m_count;
- let is_valid = expert_idx != 0xFFFFFFFFu;
+ let is_valid = tile < ids_tiles[0];
- var wg_m: u32 = 0;
- var wg_n: u32 = 0;
- var offset_wg_m: u32 = 0;
- var offset_wg_n: u32 = 0;
+ var expert_idx: u32 = 0;
var rest_token_n: u32 = 0;
var src0_batch_offset: u32 = 0;
- wg_m = wg_in_batch % wg_m_count;
- wg_n = wg_in_batch / wg_m_count;
-
- offset_wg_m = wg_m * WORKGROUP_SIZE_M * TILE_M;
- offset_wg_n = wg_n * WORKGROUP_SIZE_N * TILE_N;
+ let offset_wg_m = wg_m * WORKGROUP_SIZE_M * TILE_M;
if (is_valid) {
- rest_token_n = gathered_count_ids[expert_idx] - offset_wg_n;
- let global_gathered_base = expert_idx * params.n_tokens + offset_wg_n;
- for (var i = thread_id; i < TILE_N * WORKGROUP_SIZE_N && offset_wg_n + i < gathered_count_ids[expert_idx]; i += TOTAL_WORKGROUP_SIZE) {
- gathered_expert_used[i] = global_gathered_expert_used[global_gathered_base + i];
- gathered_tokens[i] = global_gathered_tokens[global_gathered_base + i];
+ expert_idx = ids_tiles[1 + 3 * tile];
+ let slot_base = ids_tiles[1 + 3 * tile + 1];
+ rest_token_n = ids_tiles[1 + 3 * tile + 2];
+ for (var i = thread_id; i < rest_token_n; i += TOTAL_WORKGROUP_SIZE) {
+ let pos_in_ids = ids_slot_pos[slot_base + i];
+ gathered_expert_used[i] = pos_in_ids % params.n_expert_used;
+ gathered_tokens[i] = pos_in_ids / params.n_expert_used;
}
src0_batch_offset = params.offset_src0 + expert_idx * params.stride_02;
}
@@ -135,7 +120,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
workgroupBarrier();
let output_row_base = offset_wg_m + local_m * TILE_M;
- let output_col_base = offset_wg_n + local_n * TILE_N;
+ let output_col_base = local_n * TILE_N;
let dst2_stride = params.m * params.n_expert_used;
let dst1_stride = params.m;
@@ -178,9 +163,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
if (is_valid) {
for (var tn = 0u; tn < TILE_N; tn++) {
let n_idx = output_col_base + tn;
- if (n_idx < gathered_count_ids[expert_idx]) {
- let dst1_idx = gathered_expert_used[n_idx - offset_wg_n];
- let dst2_idx = gathered_tokens[n_idx - offset_wg_n];
+ if (n_idx < rest_token_n) {
+ let dst1_idx = gathered_expert_used[n_idx];
+ let dst2_idx = gathered_tokens[n_idx];
let dst12_offset = params.offset_dst + dst2_idx * dst2_stride + dst1_idx * dst1_stride;
for (var tm = 0u; tm < TILE_M; tm += VEC_SIZE) {
let global_row = output_row_base + tm;
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
index 581e92270..1439a12d0 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
@@ -11,42 +11,68 @@ struct MulMatIdGatherParams {
};
@group(0) @binding(0) var<storage, read_write> ids: array<i32>; // [n_expert_used, n_tokens]
-@group(0) @binding(1) var<storage, read_write> global_gathered_expert_used: array<u32>; // [n_expert][n_tokens]
-@group(0) @binding(2) var<storage, read_write> global_gathered_tokens: array<u32>; // [n_expert][n_tokens]
-@group(0) @binding(3) var<storage, read_write> gathered_count_ids: array<u32>; // [n_expert]
+@group(0) @binding(1) var<storage, read_write> ids_slot_pos: array<u32>; // [n_expert_used * n_tokens]
+@group(0) @binding(2) var<storage, read_write> ids_tiles: array<u32>; // [n_tiles, (expert, first slot, n slots) * n_tiles]
-@group(0) @binding(4) var<uniform> params: MulMatIdGatherParams;
+@group(0) @binding(3) var<uniform> params: MulMatIdGatherParams;
-var<workgroup> count:atomic<u32>;
+// per expert: the slot count, replaced by the first slot of the expert in ids_slot_pos by the prefix sum, then advanced as its write cursor
+var<workgroup> ids_count: array<atomic<u32>, N_EXPERT>;
+// launched as a single workgroup, where the threads process the elements of ids in parallel with atomics
@compute @workgroup_size(WG_SIZE)
-fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
- @builtin(local_invocation_id) local_id: vec3<u32>) {
+fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
let thread_id = local_id.x;
- let own_expert = wg_id.x; // the expert assigned to this workgroup
+ let ids_slots_num = params.n_expert_used * params.n_tokens;
- if (thread_id == 0u) {
- atomicStore(&count, 0);
+ for (var e = thread_id; e < N_EXPERT; e += WG_SIZE) {
+ atomicStore(&ids_count[e], 0u);
}
workgroupBarrier();
- for (var i = thread_id;i < params.n_expert_used * params.n_tokens;i += WG_SIZE) {
+ // count the slots of each expert
+ for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
let row = i / params.n_expert_used;
let col = i % params.n_expert_used;
let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
- if (own_expert == expert) {
- let pos = atomicAdd(&count, 1u);
- let gathered_id = own_expert * params.n_tokens + pos;
- global_gathered_expert_used[gathered_id] = col;
- global_gathered_tokens[gathered_id] = row;
- }
+ atomicAdd(&ids_count[expert], 1u);
}
workgroupBarrier();
+ // exclusive prefix sum over the counts, and emit one tile descriptor per `TILE_N_SLOTS` slots of each expert
if (thread_id == 0u) {
- gathered_count_ids[own_expert] = atomicLoad(&count);
+ var base: u32 = 0;
+ var tile: u32 = 0;
+
+ for (var e: u32 = 0; e < N_EXPERT; e += 1) {
+ let n = atomicLoad(&ids_count[e]);
+
+ atomicStore(&ids_count[e], base);
+
+ for (var k: u32 = 0; k < n; k += TILE_N_SLOTS) {
+ ids_tiles[1 + 3 * tile] = e;
+ ids_tiles[1 + 3 * tile + 1] = base + k;
+ ids_tiles[1 + 3 * tile + 2] = min(n - k, TILE_N_SLOTS);
+ tile += 1;
+ }
+
+ base += n;
+ }
+
+ ids_tiles[0] = tile;
+ }
+
+ workgroupBarrier();
+
+ // store the position in ids tensor to ids_slot_pos
+ for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
+ let row = i / params.n_expert_used;
+ let col = i % params.n_expert_used;
+ let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
+ let pos = atomicAdd(&ids_count[expert], 1u);
+ ids_slot_pos[pos] = i;
}
}
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
index 78ae955e6..cb95783bc 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
@@ -59,45 +59,20 @@ fn main(
let thread_id = local_id.x;
- for (var i = thread_id;i < params.n_expert;i += WG_SIZE) {
- gathered_count_ids[i] = 0;
- }
-
- workgroupBarrier();
-
- // gather the selected experts for the target token.
- for (var col = thread_id;col < params.n_expert_used;col += WG_SIZE) {
- let expert = ids[params.offset_ids + col];
- gathered_count_ids[expert] = 1;
- gathered_expert_used[expert] = col;
- }
-
- workgroupBarrier();
-
- let output_groups:u32 = (params.m + OUTPUTS_PER_WG - 1u) / OUTPUTS_PER_WG;
+ let output_groups: u32 = (params.m + OUTPUTS_PER_WG - 1u) / OUTPUTS_PER_WG;
let wg_linear = wg_id.y * num_wg.x + wg_id.x;
- var own_expert:u32 = 0;
- var wg_in_batch:u32 = 0;
- var wg_sum:u32 = 0;
-
- for (var i = 0u;i < params.n_expert;i += 1) {
- let wg_vec_count = gathered_count_ids[i]; // 1 or 0
- let wg_per_matrix = output_groups * wg_vec_count;
- if (wg_sum <= wg_linear && wg_linear < wg_sum + wg_per_matrix) {
- own_expert = i;
- wg_in_batch = wg_linear - wg_sum;
- break;
- }
- wg_sum += wg_per_matrix;
+ if (wg_linear >= output_groups * params.n_expert_used) {
+ return;
}
- let row_base = (wg_linear % output_groups) * OUTPUTS_PER_WG;
- let dst1_stride = params.m;
+ let slot = wg_linear / output_groups;
+ let own_expert = ids[params.offset_ids + slot];
+ let row_base = (wg_linear % output_groups) * OUTPUTS_PER_WG;
let src0_batch_offset = params.offset_src0 + own_expert * params.stride_02;
- let src1_idx_base = params.offset_src1 + (gathered_expert_used[own_expert] % params.b_ne1) * params.stride_11;
- let dst_idx_base = params.offset_dst + gathered_expert_used[own_expert] * dst1_stride + row_base;
+ let src1_idx_base = params.offset_src1 + (slot % params.b_ne1) * params.stride_11;
+ let dst_idx_base = params.offset_dst + slot * params.m + row_base;
let acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);