forward_moe

Function forward_moe 

Source
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.