pub fn forward_moe(
gpu: &dyn GpuBackend,
k: &Glm5NextMlpKernels,
cfg: &Glm5NextMlpConfig,
w: &Glm5NextMoeWeights,
x: DevicePtr,
out: DevicePtr,
rows: usize,
ws: &Glm5NextMlpWorkspace,
stream: u64,
) -> Result<()>Expand description
One routed MoE site, one token. Leaves a partial sum in out whenever this rank shares
the experts (EP) or the shared expert (TP) with anyone else.
§🪤 The device→host round trip
The expert loop reads the selected ids back to the host to decide which experts are local.
That is a synchronising copy_d2h on the decode critical path, once per routed layer.
It is deliberate for this slice: correctness first, and it is exactly what the gated
microtest does. The upgrade path is the pointer-table grouped GEMM
(layers::moe::ptr_table_build), which keeps the routing on device — not a change to
this math.
Routed MoE over rows rows.
🔴 The routed experts amortize PARTIALLY over a verify’s rows. Measured on the live routing trace (42 series x 406 steps), the union of selected experts over K consecutive tokens is 8.00 / 13.74 / 18.76 / 23.35 at K = 1..4, so their weight traffic grows with K however the loop is written — but it grows along that curve, not along 8K.
🔴 RETRACTED (2026-08-29): this comment used to say the routed experts “stay one row at a
time” and that deduplicating the union “needs a device-side sort the dispatch kernel does
not have yet”. Both are wrong. top_k * rows <= 64 ids resolve in ONE block by pairwise
scan — no sort — and w4a16_gemv_sw_moe_batchm_mR then sweeps each union expert once.
Measured on t69, K=3, six probes byte-identical either way: open512 20.25 -> 22.12 tok/s
(+9.2%), a 125.4 -> 114.8 ms step. At K=2 (the serving default) 18.91 -> 19.75.
đź”´ WIDENED (2026-08-31) from rows <= 4 to rows <= 8. The 4 was the compiled tier family,
not the union’s limit — 8 * 8 == 64 fits its single block exactly. This is what makes the
8-row batched PREFILL sub-chunk (ANOMALIES A65) amortize its routed experts too; before it,
prefill batched every stage EXCEPT the experts, which by then were most of the traffic left.
The SHARED expert is a different animal again: it is the same weights for every row, so it runs once over all of them regardless.