feat(engine): E1 MoE block fwd_moe - router, gather_qmm/mm, shared experts (#107)

[tlkahn]

Summary

Implements #107 (parent #104 E1 MoE infrastructure): a family-agnostic MoE FFN block used later by DeepSeek (E2) and qwen3_5_moe (E3).

Plan: doc/plans/impl-107-e1-moe-block-fwd_moe.md

API (layered, R1)

Symbol	Role
fwd_moe_route_softmax	precise softmax + top-k via argpartition on negated logits; optional norm_topk_prob
fwd_gather_expert_linear	quantized mlx_gather_qmm / bf16 mlx_gather_mm
fwd_switch_glu	expert SwiGLU gather; decode + sorted multi-token paths
fwd_moe	route + switch + score-combine + optional shared residual
fwd_array_contiguous	row-contiguous materialize (gather_qmm silent-zeros gotcha)

Scope decisions

• R2 sort when S > 1 || B*S*K >= 64 (mlx-serve rule)
• R3 caller supplies layout: quant [E,out,in]+transpose; bf16 [E,in,out] no internal transpose
• R4 gather_mm always sorted_indices=false (mlx 0.31.2 dense sorted bug); gather_qmm honors sort
• R5 contiguous helper + gotcha GPU test; no packed gate_up load sanitize (E2: DeepSeek weight sanitization + weights_expected_names #112)
• Shared experts: always-active (DeepSeek body) and sigmoid-gated (qwen3_5 / qwen3_next)
• No family enablement, expected-names, config norm_topk_prob field, or DeepSeek group gate (E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108)

Issue #107 checkbox map

Task	Test(s)
Router top-k + softmax	test_route_softmax_topk_indices_and_scores, test_route_softmax_norm_topk_prob, test_route_softmax_batch_seq_shapes
Quantized expert dispatch via mlx_gather_qmm	test_gather_qmm_single_token_matches_manual_expert, switch/moe quant paths
Contiguous-materialization gotcha	test_gather_qmm_lazy_slice_zeros_without_contiguous
bf16 path via mlx_gather_mm	test_gather_mm_bf16_matches_per_expert_matmul, test_fwd_moe_all_bf16_experts
Shared experts always-active + residual	test_fwd_moe_shared_always_active_adds_dense_swiglu
Sigmoid-gated shared (plan extension)	test_fwd_moe_shared_sigmoid_gate

Files

• src/engine/forward.h / forward.c - MoE structs + implementation
• tests/test_fwd_moe_api.c - CPU API/NULL smoke
• tests/test_fwd_moe_gpu.c - primary GPU suite

Out of scope (unchanged)

DeepSeek sigmoid/group gate (#108), MLA/DSA, family wiring (#113/#118), expert packing sanitize (#112), parity gates.

Test plan

• make test green (includes new test_fwd_moe_api)
• ./tests/test_fwd_moe_gpu green (15 cases)
• ASan/UBSan clean on MoE API + GPU suites
• make test-gpu - MoE suite OK; pre-existing test_weights_gpu DeepSeek V4 GGUF assert failure is unrelated
• Reviewer: spot-check R2/R3/R4 comments and ownership/goto cleanup in sort path
• After merge: E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108 can consume inds/scores shape via a new fwd_moe_route_deepseek

[lgsunnyvale]

Code review: PR 122 (E1 MoE block fwd_moe)

Reviewed as senior staff against issue #107, plan doc/plans/impl-107-e1-moe-block-fwd_moe.md, mlx-serve moeMLP2 / gatherExpertMm / moeRoutingChain, and mlx-lm switch_layers.SwitchGLU + qwen3_moe / qwen3_next / deepseek_v3 shared paths.

Scope verdict: In-scope for #107. Family-agnostic, no enablement/expected-names/parity creep. R1-R5 honored. Locally re-ran tests/test_fwd_moe_api and tests/test_fwd_moe_gpu (15 cases) - green.

Overall: Approve with nits. No merge-blocking correctness bug found in the stated synthetic-expert scope. A few API/test footguns worth fixing or explicitly parking before E2/E3 consume this.

---

What looks solid

1. Layering matches R1. fwd_moe_route_softmax / fwd_gather_expert_linear / fwd_switch_glu / fwd_moe / fwd_array_contiguous is the right split for E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108 to plug a DeepSeek route without rewriting dispatch.
2. Port fidelity. Sort rule S > 1 || B*S*K >= 64 (R2), quant [E,out,in] + gather_qmm(transpose=true), bf16 [E,in,out] + gather_mm with sorted_indices forced false (R3/R4), and contiguous helper (R5) match the mlx-serve production path and the known 0.31.2 dense-sorted bug.
3. mlx_squeeze_axis(..., -2) is an improvement over mlx-serve's full squeeze. Full squeeze on decode [1,1,K,1,H] collapses leading ones and only survives via broadcast; axis-only squeeze keeps [B,S,K,H] honest at B=1/S=1. Good comment in the sort path.
4. Shared-expert matrix is complete for named consumers. Always-active (DeepSeek body) and sigmoid-gated (qwen3_5 / qwen3_next) are both implemented and parity-tested against a manual residual.
5. Ownership / cleanup. Goto style matches fwd_attention / fwd_swiglu. Sort-path locals are freed on both success and failure; *out is only replaced on success; route hands off inds/scores without double-free.
6. GPU suite quality. Router set-equality (partition-unstable order documented), quant gather vs dequant+matmul, bf16 gather vs matmul, contiguous slice vs independently quantized ref (decode + sorted shapes), switch decode/sorted, full fwd_moe manual combine, error hygiene with sentinel *out. This is the right shape of coverage for an E1 primitive.

---

Findings

Medium - address before E2/E3 land (prefer in this PR if cheap)

M1. Score-combine + shared residual are not reusable without duplicating fwd_moe.

Plan R1 says #108 should either call lower-level switch + combine with external inds/scores, or add fwd_moe_route_deepseek and reuse (2)-(3). Today combine lives only inside fwd_moe:

mlx_expand_dims(&scores_exp, scores, -1, s);
mlx_multiply(&weighted, y_exp, scores_exp, s);
mlx_sum_axis(&routed, weighted, -2, false, s);
/* + optional shared residual */

#108 will either copy this block or grow fwd_moe with a precomputed-route mode under pressure. Please extract one of:

• fwd_moe_combine(out, y_exp, scores, mw, cfg, s) (weighted sum + optional shared), or
• an optional inds/scores override on fwd_moe that skips internal route.

Either keeps the DeepSeek gate PR from re-implementing residual math. Not a #107 functional hole, but it is the main API completeness gap relative to the written plan.

M2. p->num_experts is barely enforced against the actual router.

if (p->top_k < 1 || p->num_experts < 1 || p->top_k > p->num_experts)
    return -1;
/* ... */
fwd_linear(&logits, x, &mw->router, ...);
fwd_moe_route_softmax(&inds, &scores, logits, p->top_k, ...);

num_experts is not compared to logits's last dim. A mis-wired router (E_router != num_experts) succeeds whenever top_k <= E_router, and fails opaquely inside route otherwise. After the router linear (or inside route with an expected-E argument), please reject E != p->num_experts. Callers in #113/#118 will thank you.

---

Low - nits / test gaps (non-blocking)

L1. R2 boundary S == 1 && B*S*K >= 64 is untested.

Covered: S > 1 forces sort; S=1,K=2 stays decode. Missing: e.g. B=1,S=1,K=64 (or B=32,S=1,K=2) must take the sorted path and still match the unsorted reference. That is the other disjunct of R2 and the exact mlx-lm indices.size >= 64 case.

L2. Contiguous gotcha test name oversells the assertion.

test_gather_qmm_lazy_slice_zeros_without_contiguous only asserts contiguous-slice gather == independently quantized ref (decode + sorted). It never runs gather on the lazy view, so it does not demonstrate the zeros failure mode. Hard assert is correct per plan Cycle 5; rename to something like ..._contiguous_matches_independent_quant, or add a soft/optional lazy-vs-ref divergence check with a comment that it may be mlx-version-sensitive.

L3. has_shared_expert_gate without has_shared is silently ignored.

Validation only inspects the gate triplet when has_shared is set. Prefer if (has_shared_expert_gate && !has_shared) return -1 (or document that the gate flag is meaningless alone). Same class of footgun as M2 for struct-driven callers.

L4. Header comment vs struct shape.

/* Optional pieces: pass NULL triplet pointer, or a zeroed triplet with
   weight.ctx == NULL, to mean absent. */
typedef struct {
    weight_triplet_t router;  /* by value - cannot be NULL */

NULL triplet pointers apply to fwd_switch_glu / fwd_gather_expert_linear args, not to fwd_moe_weights_t fields. Tighten the comment so loaders do not think they can leave "null pointers" inside the struct.

L5. Quant mode hardcoded "affine".

Consistent with fwd_linear, so not a regression. MoE gather is still where nvfp4 / non-affine will break first (mlx-serve already threads mode per weight). A one-line TODO on the gather_qmm call site pointing at future mode-from-triplet/cfg is enough; do not scope-creep mode plumbing here.

L6. No expert bias path.

mlx-lm SwitchLinear supports optional bias[indices]. Qwen/DeepSeek MoE experts are bias-free today, so omission is fine - please note it next to the layout contract so E2/E3 does not discover it via fluent wrong logits.

L7. int total_inds = B * S * K (and B * S reshape extents).

Fine for current single-sequence serve sizes; overflows if this primitive is ever reused with large batch. Prefer size_t for the product and a range check before casting back into int shape arrays, or at least a comment that shapes are assumed to fit int.

L8. CPU API test is smoke-only.

Acceptable for Cycle 1 under the no-mlx-in-make test rule. Not asking for more CPU coverage.

---

Intentional non-findings (checked, OK)

Topic	Resolution
Top-k via -logits + argpartition(kth=k-1) vs mlx-lm argpartition(softmax, kth=-k)[..., -k:]	Monotonic; same set. Order unstable either way; score combine is order-invariant. Documented.
R2 vs mlx-lm indices.size >= 64 only	Plan explicitly chose mlx-serve rule for multi-position verify; correct.
Always-active shared vs mlx-serve shared_expert_gate_w == null early return	mlx-serve handles DeepSeek shared outside moeMLP2; PR correctly models always-active inside fwd_moe.
norm_topk_prob optional vs mlx-serve always-normalize	Matches mlx-lm qwen knobs; right for family-agnostic E1.
Sort path passes null lhs_indices after materializing x_rep	Matches mlx-serve; not a missed gather lhs feature.
Router on softmax probs vs logits for partition	Equivalent for top-k set.
fwd_array_contiguous public with no production call site yet	Intentional R5; loaders in #112/E3 must call it.
GELU-approx expert branch tied to cfg->hidden_act	Fine for silu families in Stage E; watch if a gelu-dense / silu-expert hybrid ever appears.

---

Test plan re-check

• Issue E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107 checkboxes map to named GPU tests (PR body table is accurate)
• test_fwd_moe_api + test_fwd_moe_gpu green locally
• Please confirm ASan/UBSan on the GPU suite on your machine (claimed in PR; not re-run under sanitizers here)
• Pre-existing test_weights_gpu DeepSeek V4 GGUF failure called out as unrelated - agree, keep it out of this PR

---

Merge recommendation

Approve once M1/M2 are either fixed in-thread or explicitly accepted as follow-up issues filed against #108/#113 (M1 especially should not be forgotten). L1-L4 are quick polish if you are already touching the file.

Nice work - this is a clean E1 primitive with the right gotchas pinned.

[tlkahn]

Thanks - agree with approve-with-nits and the intentional non-findings table.

Taking (landed in 8505f36):

• M1 extract fwd_moe_combine (score-combine + optional shared residual); fwd_moe is now route + switch + combine
• M2 reject when router logits last dim != p->num_experts
• L1 R2 boundary characterization: S==1 && B*S*K >= 64 sorted path
• L3 reject has_shared_expert_gate && !has_shared
• L4/L6 tighten fwd_moe_weights_t layout / presence / bias-free comment
• L5 one-line TODO above "affine" (matches fwd_linear; no mode plumbing)
• L2 rename contiguous gotcha test to test_gather_qmm_contiguous_slice_matches_independent_quant

Refuting:

• L7 overflow as current-arch false positive (single-seq serve; continuous batching deferred to v2)
• L8 not a defect (CPU smoke under no-mlx-in-make test invariant)
• L2 "missing zeros assert" is plan-intentional: hard assert is contiguous==independent ref; lazy zeros path was consciously soft/optional in E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107 Cycle 5 (rename only)

Cycle-level TDD plan: doc/plans/fix-122-5020153959.md (local). Suites: make test green; ./tests/test_fwd_moe_gpu green (prior + combine/M2/L3/L1); ASan/UBSan clean on MoE API + GPU.