Skip to content

Commit bbe97e5

Browse files
committed
opencl: stop writing zeros into the padded MoE activation slots
kernel_moe_reorder_b zero-filled every padded token slot. The MoE GEMMs accumulate per output column and scatter only the real columns, so whatever sits in a padded slot never reaches dst and the fill is pure bandwidth. Established by a poison test rather than by reading the closed-source kernel: filling padded slots with 1e30 leaves MUL_MAT_ID at 383 OK / 0 FAIL. The positive control is what gives that meaning -- poisoning the real gather as well fails 188 of 383 with NaN at index 0, and 188 is exactly the moe_reorder_b dispatch count, so the buffer is read and the test can fail. The destination is the reused prealloc_act_trans scratch, so padded slots now retain stale activations from a previous layer rather than zeros. That is safe for the same reason. Worth about 5% of the kernel on granite-3b MXFP4. Perplexity is bit-identical on three models, which covers the cross-layer stale-buffer case test-backend-ops cannot reproduce.
1 parent cda33c2 commit bbe97e5

1 file changed

Lines changed: 9 additions & 5 deletions

File tree

ggml/src/ggml-opencl/kernels/moe_reorder_b.cl

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,15 @@ kernel void kernel_moe_reorder_b(
2020

2121
uint router_idx = router[post_router_idx];
2222

23-
float4 out = (float4)(0);
24-
if (router_idx != 0xFFFFFFFF) {
25-
ushort activation_idx = router_idx / map_ratio;
26-
out = src[activation_idx * K / 4 + k_4];
23+
// Padded slots need not be written at all. The MoE GEMMs accumulate per output
24+
// column and scatter only the real columns, so whatever sits in a padded slot
25+
// never reaches dst -- verified by filling them with 1e30 and re-running
26+
// test-backend-ops MUL_MAT_ID (383 OK / 0 FAIL, unchanged), against a positive
27+
// control that poisons the real gather too and fails 188 of 383.
28+
if (router_idx == 0xFFFFFFFF) {
29+
return;
2730
}
2831

29-
dst[post_router_idx * K / 4 + k_4] = out;
32+
ushort activation_idx = router_idx / map_ratio;
33+
dst[post_router_idx * K / 4 + k_4] = src[activation_idx * K / 4 + k_4];
3034
}

0 commit comments

Comments
 (0)