pub fn from_moe(
spec: &MoeSpec,
names: &StageNames,
target: &SpatialTarget,
) -> Result<Vec<SpatialProgram>, SpatialReject>Expand description
The Mixture-of-Experts FFN — route x to its top-k of ne experts and combine their SwiGLUs by the
renormalized router weight:
L = x·Wr router logits (1 × ne)
P = softmax(L)
b_1..b_k the k largest of P, first maximum on ties (k rounds of arg-max over a mask)
out = Σ_i P[b_i]·down_{b_i}(silu(gate_{b_i}·x) ⊙ up_{b_i}·x) / Σ_i P[b_i]★ ONLY THE CHOSEN EXPERTS’ RUNS MOVE — THAT IS WHAT THE GATHER IS FOR. The desugar computes every expert
and weights the unchosen ones by zero (a mask the VM can run); a chip that did the same would stream all
ne experts’ weights to use k of them. Expert b’s gate row o lives at (b·dff + o)·d + t — the
affine o·d + t past its run’s first element b·(dff·d) — so each round loads the whole tiles around its
expert’s run of each table and reads the run at its offset in the page (ExpertTable): k runs move,
and an expert index past the table faults. A table no more tiles than its k windows loads once, and the
rounds read their runs out of it through an IndirectTerm of stride the run and extent ne.
Three sections: the router (a reduction), the row softmax (flat_softmax_rows), and the routing plus
the experts. The weights sum is divided out once at the end: Σ P[b]·y_b / Σ P[b] is the desugar’s
Σ (P[b]/Σ P)·y_b.