feat(engine): E1 DeepSeek-style MoE gate (sigmoid, group top-k, bias) (#108)

[tlkahn]

Summary

Implements #108 (parent #104 E1 MoE infrastructure): DeepSeek V3/V3.2/V4-style MoE gate as a family-agnostic routing primitive.

Plan: doc/plans/imp-108-e1-deepseek-style-moe-gate.md

API

• fwd_moe_route_deepseek_params_t + fwd_moe_route_deepseek(...)
• Inputs: precomputed router logits [..., E], required e_score_correction_bias [E]
• Outputs: owned inds [..., K], scores [..., K] (same contract as softmax route)

Algorithm (mlx-lm group_expert_select)

scores = sigmoid(f32(gates)); orig = scores
scores = scores + bias                    # selection only
if n_group > 1:
  unflatten -> top-2 sum group scores -> zero worst groups -> flatten
inds = topk(biased scores)
scores = take_along(orig, inds)           # pre-bias weights (S1)
if top_k > 1 and norm_topk_prob: renormalize
scores *= routed_scaling_factor

Scope decisions

• Softmax fwd_moe convenience wrapper unchanged (DeepSeek composes route + switch + combine)
• No family enablement / weight-map / config parse (deferred to E2)
• Bias required (callers pass zeros when absent)
• k_drop == 0 skips group masking rather than argpartition with kth=-1

Tests

Issue task	Coverage
Sigmoid scoring	test_route_deepseek_sigmoid_topk_n_group_1
Group expert selection	test_route_deepseek_group_*
e_score_correction_bias (S1 split)	test_route_deepseek_scores_ignore_bias_for_values
routed_scaling_factor	test_route_deepseek_routed_scaling_factor
norm_topk_prob (incl. top_k==1 no-op)	test_route_deepseek_norm_topk_prob_*
mlx-lm goldens	test_route_deepseek_matches_mlx_lm_golden_vectors
Compose with #107 switch/combine	test_fwd_moe_deepseek_compose_*

Goldens frozen from mlx-lm group_expert_select (n_group=1 norm+scale; n_group=4 + non-zero bias).

Test plan

• make test (CPU) green including test_fwd_moe_api
• ./tests/test_fwd_moe_deepseek_gpu green (21 cases)
• ./tests/test_fwd_moe_gpu still green (softmax path untouched)
• ASan/UBSan clean on new GPU suite
• CI green

Closes #108.

[lgsunnyvale]

Code review: feat(engine): E1 DeepSeek-style MoE gate (#123)

Thorough review of the full PR: forward.h / forward.c (implementation), test_fwd_moe_api.c (CPU validation), and test_fwd_moe_deepseek_gpu.c (1206-line GPU suite with 21 test cases).

Summary

The implementation is solid. It follows the plan (imp-108-e1-deepseek-style-moe-gate.md) precisely, respects all design decisions (D1-D9), and has comprehensive test coverage including frozen mlx-lm goldens. No critical issues found.

Positive observations

• Clean separation from the softmax route. fwd_moe stays softmax-only per D4; the deepseek route is a fully independent function with its own params struct. This preserves the architecture invariant.
• The S1 "bias for selection, pre-bias sigmoid for weights" split is correctly implemented and explicitly tested (test_route_deepseek_scores_ignore_bias_for_values).
• S2 norm guard (top_k > 1 && p->norm_topk_prob) correctly differs from the softmax route's always-norm behavior. Tested both branches including the top_k == 1 no-op case.
• Group path correctly implements top-2 sum per group (S5), not max-pooling. Tested with a case where max expert is in group A but group B wins on top-2 sum.
• D7 k_drop == 0 skip avoids calling argpartition with kth=-1.
• Ownership contract (*inds/*scores unchanged on failure, freed+replaced on success) is correctly implemented and tested in Cycle 9.
• All 21 tests have clear, documented purpose. The mlx-lm golden test (Cycle 7) covers both n_group=1 and n_group=4 with non-zero bias.
• Makefile wildcards (test_*_gpu.c) automatically pick up the new test binary; no build system changes needed.
• Cleanup label frees all temporary arrays regardless of success/error path; no early-return leak (all validation before allocation block is correct).

Issues found

1. (Medium) No guard against zero-sum division in norm path

/* S2: norm only when top_k > 1 and flag set */
if (top_k > 1 && p->norm_topk_prob) {
    if (!MLXB_CHECK(mlx_sum_axis(&weight_sum, top_w, -1, true, s)))
        goto cleanup;
    if (!MLXB_CHECK(mlx_divide(&scores_tmp, top_w, weight_sum, s)))
        goto cleanup;
}

If all selected pre-bias sigmoid weights are zero (theoretically possible with extremely negative logits combined with bf16 underflow), weight_sum will be zero and mlx_divide will produce inf/nan scores that propagate downstream. The test helper host_group_expert_select guards against this (if (sum != 0.f)) but the production code does not.

This is consistent with mlx-lm's group_expert_select which also has no zero-sum guard, so it is not an mlx-lm parity divergence. However, unlike the softmax route (where sum-is-1 makes zero-sum near-impossible), sigmoid scores can be arbitrarily close to zero. A defensive check that skips normalization when sum < eps (or that prevents downstream NaN propagation) would be safer.

Suggestion (optional follow-up, not blocking): add something like:

mlx_array zero_check = mlx_array_new();
if (!MLXB_CHECK(mlx_less(&zero_check, weight_sum, epsilon_scalar, s)))
    goto cleanup;
// ... broadcast zeros to scores_tmp where zero_check ...

Or simply document in the header that callers with potentially all-zero sigmoid scores should pre-validate.

2. (Low) e_score_correction_bias last-dim-only validation is permissive

The validation checks:

int bias_ndim = (int)mlx_array_ndim(e_score_correction_bias);
if (bias_ndim < 1) return -1;
if (mlx_array_dim(e_score_correction_bias, bias_ndim - 1) != E)
    return -1;

This only verifies the last dimension equals E, but does not reject multi-dimensional biases. A caller could pass bias shaped [1, E, 1] and it would pass validation. While mlx broadcasting will handle this correctly in the mlx_add call, the documented contract says [E] (1D). Consider tightening to bias_ndim == 1 to catch mis-shaped biases early, or update the contract comment to say "last dim E" if multi-dim is intentional.

3. (Style/Low) zero_v free-then-recreate in group path

/* In the k_drop > 0 block: */
mlx_array_free(zero_v);
zero_v = mlx_array_new_float(0.0f);

zero_v was already initialized to mlx_array_new() at function top. This pattern mutates it in-place from "empty" to "scalar 0.0". It is functionally correct because both states are valid inputs to mlx_array_free in cleanup. But it is mildly surprising on first read: a reader might wonder why we free an empty array to allocate a new one in the same variable. A comment noting that we are replacing the empty placeholder with the scalar 0.0 used by put_along_axis would help, or alternatively a separate variable.

4. (Nit) Rounding in test_route_deepseek_norm_topk_prob_ignored_when_top_k_1

assert(fabsf(sh[0] - expect) < 1e-4f);

All other score-comparison tests use 1e-5f tolerance. This test uses 1e-4f, which is 10x looser. The looser tolerance is fine (scale=2.5 introduces some rounding) but inconsistent with the rest of the suite. Consider tightening or adding a comment explaining the looser bound.

5. (Test coverage gap) No test for n_group=1, topk_group > 1

The validation has if (p->topk_group > p->n_group) return -1; which catches this case when n_group=1 and topk_group=2. The test_route_deepseek_rejects_bad_topk_group test covers topk_group=0 and topk_group=3 with n_group=2 but does not have a case for n_group=1, topk_group=2. The logic is trivially covered by the comparison check, so this is negligible, but adding a third variant would make the validation matrix more complete.

Architecture and design

The composition tests (Cycle 8) demonstrate the intended DeepSeek body pattern:

route_deepseek -> fwd_switch_glu -> fwd_moe_combine

This naturally composes with the shared-expert residual path already in fwd_moe_combine (has_shared=true, has_shared_expert_gate=false). The always-active shared expert path is independently verified with a manual swiglu reference, confirming end-to-end correctness for the DeepSeek V3/V3.2/V4 MoE body structure.

Verdict

Approve. The implementation is correct, well-tested, and faithfully ports the mlx-lm group_expert_select algorithm. The issues listed are minor and do not block merge. The medium zero-sum guard is a nice-to-have defensive improvement that can land as a follow-up.

[lgsunnyvale]

Code review: PR 123 (E1 DeepSeek-style MoE gate)

Reviewed as senior staff against issue #108, plan doc/plans/imp-108-e1-deepseek-style-moe-gate.md, sibling #107 / PR 122 layering (fwd_moe_route_softmax / fwd_switch_glu / fwd_moe_combine), and mlx-lm 0.31.x group_expert_select in deepseek_v3.py / deepseek_v32.py (bodies identical).

Scope verdict: In-scope for #108. Family-agnostic route primitive only. No engine_model_check_supported flip, no weight-map / config parse, no fwd_moe() softmax behavior change. D1-D9 honored. Softmax convenience path untouched.

Verification performed:

• Read fwd_moe_route_deepseek against frozen mlx-lm algorithm (S1-S10 / D7).
• Re-derived Case A host math; re-ran Case A + Case B through live mlx_lm.models.deepseek_v3.group_expert_select - inds/scores match the frozen vectors in test_route_deepseek_matches_mlx_lm_golden_vectors exactly.
• Locally re-ran tests/test_fwd_moe_api and tests/test_fwd_moe_deepseek_gpu (21 cases) - green.
• CI build-and-test green (CPU make test only; GPU suite is author/runtime attestation, same as E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107).

Overall: Approve with nits. No merge-blocking correctness bug found. This is a clean E1 port and the right plug-in surface for E2 composition (route_deepseek -> switch_glu -> combine).

---

What looks solid

1. Algorithm fidelity. Cast-to-f32 -> sigmoid -> save orig_scores -> add bias -> optional group mask -> top-k on biased -> take_along(orig_scores) -> norm only if top_k > 1 && flag -> * routed_scaling_factor matches mlx-lm line-for-line. S1 (bias selection vs unbiased weights) and S2 (top_k == 1 skips norm; differs from softmax route) are the two easy parity footguns and both are implemented and tested explicitly.

2. Group path details. Unflatten (G, Eg), group score = top-2 sum (not max), argpartition on axis -2, stop_gradient on indices, put_along_axis zero, flatten - correct. D7 (k_drop == 0 skips masking rather than kth=-1) is the right defensive choice; Python still builds an empty slice no-op, C just avoids the dead ops. Eg < 2 when n_group > 1 rejected - matches Python still calling topk(..., 2) whenever n_group > 1.

3. API / layering. Separate fwd_moe_route_deepseek_params_t (D1), precomputed logits + required bias (D2/D3), no fwd_moe_deepseek convenience (D4 default). Compose tests exercise the E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107 extract fwd_moe_combine with always-active shared - exactly the E2 call shape. Softmax fwd_moe remains dumb. Good.

4. Ownership / cleanup. Mirrors fwd_moe_route_softmax: many temps, single cleanup label, *inds/*scores replaced only on success after full compute, hand-off via inds_tmp = mlx_array_new() so cleanup does not free the returned arrays. Failure-leaves-untouched and success-replaces covered.

5. Test suite quality. Cycles 2-9 map cleanly to E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108 tasks. Dedicated GPU file keeps E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107/E1: DeepSeek-style MoE gate (sigmoid, group top-k, correction bias) #108 separable. Goldens are real mlx-lm outputs (not only host-ref self-consistency). Partition-unstable order handled via set equality + sort-by-id score compare. Validation matrix covers the plan table. This is the right shape of coverage for a routing primitive.

---

Findings

Low - nits / follow-ups (non-blocking)

L1. Dim-based validation is mostly GPU-suite-only; CI will not see regressions.

make test / CI runs test_fwd_moe_api only. That file covers nulls + top_k == 0. The rest of the prologue matrix lives in test_fwd_moe_deepseek_gpu.c:

• top_k > E
• E % n_group != 0
• topk_group bounds
• n_group > 1 && Eg < 2
• null bias (also in CPU null-args - good)

Those checks are pure metadata (mlx_array_ndim / mlx_array_dim) before any stream op. Prefer a few host-array cases in test_fwd_moe_api.c via mlx_array_new_data + dummy stream so CI guards the matrix without Metal. Not a functional hole today (implementation matches plan; local GPU green), but it is the main process gap relative to "CI green" on a GPU-heavy PR.

L2. Bias rank is under-validated.

int bias_ndim = (int)mlx_array_ndim(e_score_correction_bias);
if (bias_ndim < 1) return -1;
if (mlx_array_dim(e_score_correction_bias, bias_ndim - 1) != E)
    return -1;

MoEGate stores [E]. Any rank with last dim E is accepted and left to mlx_add broadcast. That is usually fine, but a mis-wired [B, E] with wrong B fails opaquely inside the add rather than at the gate boundary. Prefer bias_ndim == 1 (or document intentional broadcast and add a last-dim mismatch test either way). Bias last-dim != E is currently untested in both suites.

L3. Golden Case B seed comment is slightly off for bias.

Gates match numpy.RandomState(42).randn(16). The checked-in bias vector is 0.5 * the next randn(16), not the raw second draw. Scores still match live mlx-lm with the frozen arrays (verified). Please tweak the comment to the exact generation recipe so regeneration does not silently drift selection while keeping score tables.

L4. test_route_deepseek_batch_seq_shapes is shape/finite only.

Fine as a smoke, but a single-row compare against host_group_expert_select (already in the file) would catch axis mistakes that only show up with ndim > 2 leading dims. Group axis -2 is the classic place that passes [1,1,E] and fails [B,S,E].

L5. Test helper duplication vs test_fwd_moe_gpu.c.

Accepted by the plan for suite separation. If E2 adds more MoE GPU tests, extract tests/moe_test_util.h rather than a third copy of make_bf16 / build_bf16_switch / read_f32. No action required to close #108.

---

Non-issues / explicitly not requesting changes

• Zeroing dropped groups to 0.0 can theoretically beat kept groups with large negative biased scores. Identical to mlx-lm; do not "fix" in the port.
• n_group > 1 && k_drop == 0 still unflatten/flatten. Harmless identity; skipping entirely would be a micro-refactor only.
• No fwd_moe_deepseek wrapper. Correct deferral; composition tests are enough for E1.
• No config knobs / bias weight load / family enable. Correctly E2.
• CI lacks make test-gpu. Pre-existing project constraint; not introduced here. Author checkbox + local green is the current bar (same as E1: MoE block (fwd_moe) - router, gather_qmm/mm, shared experts #107).

---

Issue #108 checkbox map (reviewer ack)

Task	Evidence
Sigmoid scoring	test_route_deepseek_sigmoid_topk_n_group_1
Group selection n_group/topk_group	test_route_deepseek_group_* + golden B
e_score_correction_bias (S1 split)	test_route_deepseek_scores_ignore_bias_for_values
routed_scaling_factor	test_route_deepseek_routed_scaling_factor + golden A
norm_topk_prob incl. top_k==1	test_route_deepseek_norm_topk_prob_*
mlx-lm synthetic parity	test_route_deepseek_matches_mlx_lm_golden_vectors (re-verified live)
Compose with #107 switch/combine	test_fwd_moe_deepseek_compose_*

---

Merge recommendation

Approve. Land as-is, or optionally fold L1 (CPU validation hoist) and L3 (seed comment) if you want a tiny follow-up commit before merge. Nothing here should block #108 close or parent #104 gate-line acceptance once #107 is already in.

E2 callers: pass zeros when bias absent, compose fwd_linear(router) -> fwd_moe_route_deepseek -> fwd_switch_glu -> fwd_moe_combine with has_shared=true / has_shared_expert_gate=false for the DeepSeek body.

[tlkahn]

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

Took (landed in fa59b13):

• L1 - hoisted dim/param prologue rejects into tests/test_fwd_moe_api.c (test_route_deepseek_cpu_validation_matrix) so CI/make test guards them without Metal
• L2 - tightened bias to rank-1 [E] in fwd_moe_route_deepseek; CPU + GPU tests for rank != 1 and last-dim != E
• L3 - Case B golden comment now states the verified recipe (gates = rs.randn(16); bias = 0.5 * rs.randn(16)); frozen vectors unchanged
• L4 - batch/seq test renamed to test_route_deepseek_batch_seq_matches_host_rows and compares each [B,S] row to host_group_expert_select (not shape/finite only)

Parked (per review):

• L5 helper extract until E2 would force a third copy

Not mixing in sibling comment 5020570697 (zero-sum norm guard / style nits) - parity-sensitive; separate note if ever pursued.

Plan/TDD cycles: doc/plans/fix-123-5020600218.md. Verified: make test, ./tests/test_fwd_moe_deepseek_gpu, ./tests/test_fwd_moe_gpu, ASan/UBSan on API + DeepSeek GPU suites.