feat(moe): env-gated memory-efficient sparse-gather MoE (default OFF)

ReferenceMoE.__call__ was a dense loop running ALL num_experts experts on
ALL tokens then masking — O(num_experts x tokens x expert_hidden) activation,
which OOMs MLX-eager at seq=4096 (121 GB on gb10).

Add a sparse-gather routed-combine gated by CPPMEGA_MOE_EFFICIENT (default OFF
so existing paths are byte-identical). For each expert it gathers only the
tokens routed to it (top_k routing), runs the FFN on that subset, and
scatter-adds the weighted output back. Math is identical to the dense path
(a non-routed token contributes expert_out*0, exactly the dropped term).
Routing indices are stop_gradient'd, so deriving per-expert row lists on the
host does not perturb autograd; gradients flow through the gathered FFN +
scatter-add. Unrouted experts are zero-touched so their params keep a zero
(never missing) gradient.

Measured single-MoE-layer fwd+bwd (d=3584, 16 experts, top_k=4, bf16, Mac):
seq=4096 dense=3.84GB/223ms -> eff=2.20GB/83ms (1.75x less mem, 2.7x faster);
ratio grows with seq (4x less routed-activation per token).

tests/test_moe_efficient.py pins exactness: fp32 forward <1e-5, input-grad
<1e-5, param-grads <1e-4 vs dense; unrouted-expert grads present; bf16 within
reduction-order tolerance. Existing tests/test_moe.py pass with flag ON and OFF.

RULE #1: env-gated default-OFF, exact same math, fail-loud (numpy conversion
of a traced/compiled index array would raise, not silently degrade).

Author: apstenku123

cppmega_mlx/nn/moe.py

@@ -2,15 +2,29 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import Literal, NamedTuple
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
 
 ActivationName = Literal["gelu", "relu2", "swiglu"]
 
+# Env gate for the memory-efficient sparse-gather routed compute. Default OFF so
+# every existing training/inference path stays byte-identical to the dense
+# reference unless the operator explicitly opts in.
+_MOE_EFFICIENT_ENV = "CPPMEGA_MOE_EFFICIENT"
+
+
+def _moe_efficient_enabled() -> bool:
+    """Return True iff the env flag opts into the sparse-gather routed compute."""
+
+    value = os.environ.get(_MOE_EFFICIENT_ENV, "").strip().lower()
+    return value in {"1", "true", "yes", "on"}
+
 
 @dataclass(frozen=True)
 class MoEConfig:
@@ -210,12 +224,10 @@ def __call__(self, x: mx.array) -> MoEOutput:
         flat_indices = router.top_indices.reshape(-1, self.config.top_k)
         flat_weights = router.top_weights.reshape(-1, self.config.top_k)
 
-        routed = mx.zeros_like(flat_x)
-        for expert_id, expert in enumerate(self.experts):
-            expert_out = expert(flat_x)
-            mask = mx.equal(flat_indices, mx.array(expert_id, dtype=flat_indices.dtype))
-            weight = mx.sum(mx.where(mask, flat_weights, mx.zeros_like(flat_weights)), axis=-1)
-            routed = routed + expert_out * weight[:, None]
+        if _moe_efficient_enabled():
+            routed = self._routed_combine_sparse(flat_x, flat_indices, flat_weights)
+        else:
+            routed = self._routed_combine_dense(flat_x, flat_indices, flat_weights)
 
         routed = routed.reshape(x.shape)
         shared = self.shared_expert(x) if self.shared_expert is not None else None
@@ -227,6 +239,113 @@ def __call__(self, x: mx.array) -> MoEOutput:
             shared_output=shared,
         )
 
+    def _expert_weight(
+        self,
+        flat_indices: mx.array,
+        flat_weights: mx.array,
+        expert_id: int,
+    ) -> mx.array:
+        """Per-token combine weight for ``expert_id`` (0 when not selected).
+
+        Equals ``sum_k top_weights[t, k]`` over the slots ``k`` whose
+        ``top_indices[t, k] == expert_id``.  Identical in both the dense and
+        sparse routed-combine paths so the math is unchanged.
+        """
+
+        mask = mx.equal(flat_indices, mx.array(expert_id, dtype=flat_indices.dtype))
+        return mx.sum(
+            mx.where(mask, flat_weights, mx.zeros_like(flat_weights)), axis=-1
+        )
+
+    def _routed_combine_dense(
+        self,
+        flat_x: mx.array,
+        flat_indices: mx.array,
+        flat_weights: mx.array,
+    ) -> mx.array:
+        """Dense reference: run every expert on all tokens, then mask+combine.
+
+        Peak activation is ``O(num_experts x tokens x expert_hidden)`` because
+        every one of ``num_experts`` experts materializes a full-token hidden.
+        Kept as the default path for byte-identical behaviour.
+        """
+
+        routed = mx.zeros_like(flat_x)
+        for expert_id, expert in enumerate(self.experts):
+            expert_out = expert(flat_x)
+            weight = self._expert_weight(flat_indices, flat_weights, expert_id)
+            routed = routed + expert_out * weight[:, None]
+        return routed
+
+    def _routed_combine_sparse(
+        self,
+        flat_x: mx.array,
+        flat_indices: mx.array,
+        flat_weights: mx.array,
+    ) -> mx.array:
+        """Memory-efficient sparse-gather routed compute (env-gated, exact).
+
+        For each expert, gather ONLY the tokens routed to it (the rows where the
+        expert appears in ``top_indices``), run the FFN on that subset, scale by
+        the per-token combine weight, and scatter-add the contribution back.
+
+        This is mathematically identical to :meth:`_routed_combine_dense`: a
+        non-routed token contributes ``expert_out * 0`` in the dense path, which
+        is exactly the term dropped here.  The routing indices are produced under
+        ``mx.stop_gradient`` (see :class:`TopKRouter`), so deriving the per-expert
+        row lists on the host from them does not perturb autograd — gradients
+        flow only through the gathered FFN compute and the scatter-add.
+
+        Peak activation is bounded by the *largest* expert's token count, which
+        is on average ``tokens x top_k / num_experts`` instead of
+        ``num_experts x tokens`` for the dense path — a ``num_experts / top_k``
+        reduction in the routed MoE activation footprint (e.g. 16/4 = 4x), with
+        experts processed sequentially so only one expert's hidden is live.
+        """
+
+        num_tokens, d_model = flat_x.shape
+        routed = mx.zeros_like(flat_x)
+
+        # ``top_indices`` is stop_gradient'd routing metadata; materializing it on
+        # the host to build per-expert row lists is a metadata read, not a
+        # gradient-bearing op.  Gather/scatter on the device carry the gradient.
+        host_indices = np.asarray(flat_indices)
+        for expert_id, expert in enumerate(self.experts):
+            rows = np.nonzero(np.any(host_indices == expert_id, axis=1))[0]
+            if rows.size == 0:
+                # No token routed to this expert in this batch: it contributes
+                # exactly zero in the dense path. Touch its params with a zero
+                # scalar so the grad tree stays populated (grad == 0, not missing).
+                routed = routed + self._zero_touch_expert(expert).astype(flat_x.dtype)
+                continue
+            row_idx = mx.array(rows.astype(np.int32))
+            x_gathered = mx.take(flat_x, row_idx, axis=0)
+            out_gathered = expert(x_gathered)
+            weight = self._expert_weight(flat_indices, flat_weights, expert_id)
+            weight_gathered = mx.take(weight, row_idx, axis=0)
+            contrib = out_gathered * weight_gathered[:, None]
+            scattered = mx.zeros((num_tokens, d_model), dtype=flat_x.dtype)
+            scattered = scattered.at[row_idx].add(contrib)
+            routed = routed + scattered
+        return routed
+
+    def _zero_touch_expert(self, expert: "FeedForwardExpert") -> mx.array:
+        """Return a scalar that depends on every expert param but equals 0.
+
+        Used only when an expert receives no tokens in a batch so its parameters
+        still appear in the gradient tree (with a zero gradient) — never silently
+        dropped from the optimizer state.  This mirrors the dense path, where an
+        unrouted expert still runs (and gets a zero combine weight), so its
+        params always carry a (zero) gradient.
+        """
+
+        from mlx.utils import tree_flatten
+
+        acc = mx.zeros((), dtype=mx.float32)
+        for _leaf_name, leaf in tree_flatten(expert.parameters()):
+            acc = acc + leaf.astype(mx.float32).sum() * mx.zeros((), dtype=mx.float32)
+        return acc
+
 
 def _require_positive(name: str, value: int | float) -> None:
     if value <= 0:

tests/test_moe_efficient.py

@@ -0,0 +1,184 @@
+"""Numeric parity for the env-gated memory-efficient sparse-gather MoE.
+
+The efficient routed-combine (``CPPMEGA_MOE_EFFICIENT=1``) must produce the same
+forward output *and* the same gradients (w.r.t. inputs and every parameter) as
+the dense reference, otherwise the env flag would silently change training
+numerics. These tests pin that exactness in fp32 (where the only remaining delta
+is float epsilon) and document the bf16 accumulation-order tolerance.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from cppmega_mlx.nn import moe as moe_mod
+from cppmega_mlx.nn.moe import MoEConfig, ReferenceMoE
+
+
+def _config(*, d_model: int = 3584, num_experts: int = 16, top_k: int = 4) -> MoEConfig:
+    return MoEConfig(
+        d_model=d_model,
+        num_experts=num_experts,
+        top_k=top_k,
+        expert_hidden_size=128,
+        shared_expert_hidden_size=96,
+        activation="swiglu",
+    )
+
+
+def _build(config: MoEConfig, dtype: mx.Dtype) -> ReferenceMoE:
+    mx.random.seed(20260601)
+    moe = ReferenceMoE(config)
+    if dtype != mx.float32:
+        moe.update(
+            tree_unflatten(
+                [
+                    (name, leaf.astype(dtype) if leaf.dtype == mx.float32 else leaf)
+                    for name, leaf in tree_flatten(moe.parameters())
+                ]
+            )
+        )
+    mx.eval(moe.parameters())
+    return moe
+
+
+def _run_both(moe: ReferenceMoE, x: mx.array):
+    """Return (dense_output, efficient_output) for the same module + input."""
+
+    moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+    dense_out = moe(x).output
+    moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = "1"
+    try:
+        eff_out = moe(x).output
+    finally:
+        moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+    mx.eval(dense_out, eff_out)
+    return dense_out, eff_out
+
+
+def test_env_flag_default_off() -> None:
+    moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+    assert moe_mod._moe_efficient_enabled() is False
+    for value in ("1", "true", "on", "YES"):
+        moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = value
+        assert moe_mod._moe_efficient_enabled() is True
+    moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = "0"
+    assert moe_mod._moe_efficient_enabled() is False
+    moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+
+
+def test_efficient_forward_matches_dense_fp32() -> None:
+    config = _config()
+    moe = _build(config, mx.float32)
+    mx.random.seed(11)
+    x = mx.random.normal((1, 128, config.d_model)).astype(mx.float32)
+
+    dense_out, eff_out = _run_both(moe, x)
+    max_diff = float(mx.max(mx.abs(dense_out - eff_out)))
+    assert max_diff < 1e-5, f"forward fp32 max_diff={max_diff:.3e} exceeds 1e-5"
+
+
+def test_efficient_input_grad_matches_dense_fp32() -> None:
+    config = _config()
+    moe = _build(config, mx.float32)
+    mx.random.seed(12)
+    x = mx.random.normal((1, 128, config.d_model)).astype(mx.float32)
+
+    def dense_loss(xx: mx.array) -> mx.array:
+        moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+        return moe(xx).output.square().sum()
+
+    def eff_loss(xx: mx.array) -> mx.array:
+        moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = "1"
+        try:
+            return moe(xx).output.square().sum()
+        finally:
+            moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+
+    gd = mx.grad(dense_loss)(x)
+    gs = mx.grad(eff_loss)(x)
+    mx.eval(gd, gs)
+    max_diff = float(mx.max(mx.abs(gd - gs)))
+    assert max_diff < 1e-5, f"input-grad fp32 max_diff={max_diff:.3e} exceeds 1e-5"
+
+
+def test_efficient_param_grads_match_dense_fp32() -> None:
+    config = _config()
+    mx.random.seed(13)
+    x = mx.random.normal((1, 128, config.d_model)).astype(mx.float32)
+
+    def loss(model: ReferenceMoE, xx: mx.array) -> mx.array:
+        return model(xx).output.square().sum()
+
+    # Dense
+    moe_dense = _build(config, mx.float32)
+    moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+    _, gd = nn.value_and_grad(moe_dense, loss)(moe_dense, x)
+    # Efficient (same seed -> same weights)
+    moe_eff = _build(config, mx.float32)
+    moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = "1"
+    try:
+        _, gs = nn.value_and_grad(moe_eff, loss)(moe_eff, x)
+    finally:
+        moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+
+    flat_d = dict(tree_flatten(gd))
+    flat_s = dict(tree_flatten(gs))
+    assert set(flat_d) == set(flat_s), "gradient trees differ in keys"
+    worst = 0.0
+    worst_name = ""
+    for name, gd_leaf in flat_d.items():
+        d = float(mx.max(mx.abs(gd_leaf - flat_s[name])))
+        if d > worst:
+            worst, worst_name = d, name
+    assert worst < 1e-4, f"param-grad fp32 worst={worst:.3e} ({worst_name}) exceeds 1e-4"
+
+
+def test_efficient_unrouted_expert_keeps_zero_param_grad() -> None:
+    """An expert with no routed tokens still appears in the grad tree (grad 0)."""
+
+    config = _config(num_experts=16, top_k=1)
+    mx.random.seed(99)
+    moe = _build(config, mx.float32)
+    # Force the router so all tokens pick expert 0 -> experts 1..15 unrouted.
+    weight = np.zeros((config.num_experts, config.d_model), dtype=np.float32)
+    weight[0, 0] = 50.0  # huge logit for expert 0 on channel 0
+    moe.router.gate.weight = mx.array(weight)
+    mx.eval(moe.parameters())
+    x = mx.zeros((1, 8, config.d_model), dtype=mx.float32)
+    x = x + 0.0
+    x_idx = np.zeros((8, config.d_model), dtype=np.float32)
+    x_idx[:, 0] = 1.0  # ensure channel 0 active so expert 0 wins
+    x = mx.array(x_idx[None])
+
+    def loss(model: ReferenceMoE, xx: mx.array) -> mx.array:
+        return model(xx).output.square().sum()
+
+    moe_mod.os.environ[moe_mod._MOE_EFFICIENT_ENV] = "1"
+    try:
+        _, grads = nn.value_and_grad(moe, loss)(moe, x)
+    finally:
+        moe_mod.os.environ.pop(moe_mod._MOE_EFFICIENT_ENV, None)
+    flat = dict(tree_flatten(grads))
+    # Every expert param must be present (zero grad allowed, missing not).
+    for e in range(config.num_experts):
+        for proj in ("gate_proj", "up_proj", "down_proj"):
+            key = f"experts.{e}.{proj}.weight"
+            assert key in flat, f"missing grad for {key}"
+            assert np.isfinite(np.array(flat[key])).all(), key
+
+
+def test_efficient_bf16_within_accumulation_tolerance() -> None:
+    """bf16 differs only by reduction-order rounding (documented, not a bug)."""
+
+    config = _config()
+    moe = _build(config, mx.bfloat16)
+    mx.random.seed(14)
+    x = mx.random.normal((1, 128, config.d_model)).astype(mx.bfloat16)
+    dense_out, eff_out = _run_both(moe, x)
+    diff = float(mx.max(mx.abs(dense_out.astype(mx.float32) - eff_out.astype(mx.float32))))
+    # Reordered bf16 reductions: tolerance ~ a few ULPs of the output scale.
+    assert diff < 5e-2, f"bf16 max_diff={diff:.3e} unexpectedly large"