refactor: stateless MoE routing with grouped dispatch

- replace per-expert mask scan with sort+bincount grouped dispatch
- carry router stats in forward output instead of module state
- keep MoE diagnostics working under DDP/FSDP wrappers
- remove unused _load_balancing_loss helper

tests/module/test_forward_configs.py

@@ -265,10 +265,9 @@ def test_moe_defaults_preserve_normalized_routing():
     assert model.layers[0].mlp.norm_topk_prob is True
 
 
-def test_moe_router_probs_populated_after_forward():
-    """Verify DeepSeekMoE._router_probs is set after forward in training mode."""
+def test_moe_router_stats_in_output_during_training():
+    """Verify forward output carries per-layer router_stats in training mode."""
     from astrai.config.model_config import AutoRegressiveLMConfig
-    from astrai.model.components.mlp import DeepSeekMoE
 
     config = AutoRegressiveLMConfig(
         **TINY_CONFIG,
@@ -283,19 +282,18 @@ def test_moe_router_probs_populated_after_forward():
     input_ids = torch.randint(0, config.vocab_size, (2, 8))
 
     with torch.enable_grad():
-        model(input_ids)
+        outputs = model(input_ids)
 
-    # All MoE layers should have router_probs set
-    moe_layers = [m for m in model.modules() if isinstance(m, DeepSeekMoE)]
-    assert len(moe_layers) > 0
-    for layer in moe_layers:
-        assert layer._router_probs is not None
-        assert layer._router_probs.ndim == 2
-        assert layer._router_probs.shape[-1] == 4  # n_routed_experts
+    stats = outputs["router_stats"]
+    assert isinstance(stats, list)
+    assert len(stats) == config.num_hidden_layers
+    for s in stats:
+        assert s["probs"].shape == (2 * 8, 4)  # (N, n_routed_experts)
+        assert s["topk_indices"].shape == (2 * 8, 2)  # (N, n_activated_experts)
 
 
-def test_get_moe_router_probs_moe_model():
-    """Verify get_moe_router_probs() returns a list of tensors for MoE models."""
+def test_moe_router_stats_absent_in_eval():
+    """Verify no router_stats are emitted outside training."""
     from astrai.config.model_config import AutoRegressiveLMConfig
 
     config = AutoRegressiveLMConfig(
@@ -306,60 +304,27 @@ def test_get_moe_router_probs_moe_model():
         n_activated_experts=2,
     )
     model = AutoRegressiveLM(config)
-    model.train()
+    model.eval()
 
-    with torch.enable_grad():
-        model(torch.randint(0, config.vocab_size, (2, 8)))
+    with torch.no_grad():
+        outputs = model(torch.randint(0, config.vocab_size, (2, 8)))
 
-    probs = model.get_moe_router_probs()
-    assert isinstance(probs, list)
-    assert len(probs) == 2  # num_hidden_layers
-    for p in probs:
-        assert p.ndim == 2
-        assert p.shape[-1] == 4
+    assert "router_stats" not in outputs
 
 
-def test_get_moe_router_probs_non_moe_model():
-    """Verify get_moe_router_probs() returns empty list for non-MoE models."""
+def test_no_router_stats_for_mlp_model():
+    """Verify pure MLP models emit no router_stats and no aux_loss."""
     from astrai.config.model_config import AutoRegressiveLMConfig
 
     config = AutoRegressiveLMConfig(**TINY_CONFIG, ffn_type="mlp")
     model = AutoRegressiveLM(config)
-
-    probs_untrained = model.get_moe_router_probs()
-    assert probs_untrained == []
-
     model.train()
+
     with torch.enable_grad():
-        model(torch.randint(0, config.vocab_size, (2, 8)))
+        outputs = model(torch.randint(0, config.vocab_size, (2, 8)))
 
-    probs = model.get_moe_router_probs()
-    assert probs == []
-
-
-def test_collect_router_probs_static_method():
-    """Verify DeepSeekMoE.collect_router_probs static method."""
-    from astrai.model.components.mlp import DeepSeekMoE
-
-    moe = DeepSeekMoE(
-        dim=8,
-        dim_ffn=16,
-        n_routed_experts=4,
-        n_shared_experts=1,
-        n_activated_experts=2,
-    )
-    moe.train()
-    with torch.enable_grad():
-        moe(torch.randn(2, 8, 8))
-
-    # collect_router_probs should find the MoE layer
-    probs = DeepSeekMoE.collect_router_probs(moe)
-    assert len(probs) == 1
-    assert probs[0].shape[-1] == 4
-
-    # On a plain MLP module, should return empty
-    mlp_module = MLP(8, 16)
-    assert DeepSeekMoE.collect_router_probs(mlp_module) == []
+    assert "router_stats" not in outputs
+    assert "aux_loss" not in outputs
 
 
 def test_moe_aux_loss_only_emitted_during_training():

tests/module/test_model_forward.py

@@ -9,19 +9,15 @@ import pytest
 import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.model.components.mlp import DeepSeekMoE
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.trainer.strategy import (
     SEQStrategy,
     SFTStrategy,
     StrategyFactory,
     _collect_moe_diagnostics,
-    _load_balancing_loss,
 )
 from tests.helpers import TINY_CONFIG
 
-# ── helpers ──────────────────────────────────────────────────────────
-
 
 def _make_tiny_moe_config(**overrides) -> AutoRegressiveLMConfig:
     return AutoRegressiveLMConfig(
@@ -43,14 +39,16 @@ def _make_model(config=None) -> AutoRegressiveLM:
     return AutoRegressiveLM(config)
 
 
-# ── _collect_moe_diagnostics unit tests ─────────────────────────────
+def _router_stats(probs, topk_indices):
+    return {"probs": probs, "topk_indices": topk_indices}
 
 
 def test_collect_moe_diagnostics_returns_all_keys():
     """_collect_moe_diagnostics should return the four expected keys."""
     # Simulate two MoE layers with uniform routing probabilities
     probs = torch.ones(128, 4) / 4.0
-    diag = _collect_moe_diagnostics([probs, probs], top_k=2)
+    topk = torch.zeros(128, 2, dtype=torch.long)
+    diag = _collect_moe_diagnostics([_router_stats(probs, topk)] * 2)
 
     assert set(diag.keys()) == {
         "router_entropy",
@@ -64,11 +62,11 @@ def test_collect_moe_diagnostics_returns_all_keys():
 
 def test_collect_moe_diagnostics_empty_list():
     """Empty list returns empty dict."""
-    assert _collect_moe_diagnostics([], top_k=2) == {}
+    assert _collect_moe_diagnostics([]) == {}
 
 
 def test_collect_moe_diagnostics_uniform_routing():
-    """Uniform routing probabilities with top_k=2 → tie-breaking by index.
+    """Uniform routing with top_k=2 → tie-breaking by index.
 
     torch.topk breaks ties by index, so with equal probabilities
     experts 0 and 1 always win over experts 2 and 3:
@@ -77,7 +75,8 @@ def test_collect_moe_diagnostics_uniform_routing():
       - load_imbalance_max = 2.0
     """
     probs = torch.ones(128, 4) / 4.0
-    diag = _collect_moe_diagnostics([probs], top_k=2)
+    topk = torch.tensor([[0, 1]] * 128)
+    diag = _collect_moe_diagnostics([_router_stats(probs, topk)])
 
     assert diag["dead_expert_fraction"] == pytest.approx(0.5, abs=1e-6)
     assert diag["load_imbalance_mean"] == pytest.approx(1.0, abs=1e-6)
@@ -88,38 +87,41 @@ def test_collect_moe_diagnostics_max_entropy():
     """Uniform probabilities should give log(num_experts) entropy."""
     num_experts = 4
     probs = torch.ones(128, num_experts) / num_experts
-    diag = _collect_moe_diagnostics([probs], top_k=2)
+    topk = torch.zeros(128, 2, dtype=torch.long)
+    diag = _collect_moe_diagnostics([_router_stats(probs, topk)])
     expected_entropy = float(torch.log(torch.tensor(num_experts, dtype=torch.float32)))
     assert diag["router_entropy"] == pytest.approx(expected_entropy, abs=1e-5)
 
 
-# ── _load_balancing_loss unit tests ──────────────────────────────────
+def test_moe_metrics_flow_through_wrapped_model(device):
+    """DDP-like wrappers (no .config / get_moe_router_probs) still collect MoE metrics."""
+    import torch.nn as nn
 
+    from astrai.trainer.strategy import SEQStrategy
 
-def test_load_balancing_loss_shape_and_range():
-    """Verify _load_balancing_loss returns a non-negative scalar tensor."""
-    probs = torch.randn(64, 8).softmax(dim=-1)
-    loss = _load_balancing_loss(probs)
-    assert loss.ndim == 0
-    assert loss.item() >= 0
+    class ForwardOnlyWrapper(nn.Module):
+        def __init__(self, model):
+            super().__init__()
+            self.module = model
 
+        def forward(self, *args, **kwargs):
+            return self.module(*args, **kwargs)
 
-def test_load_balancing_loss_uniform_minimum():
-    """Uniform routing gives the lowest possible load balancing loss."""
-    probs = torch.ones(64, 8) / 8.0
-    loss = _load_balancing_loss(probs).item()
+    config = _make_tiny_moe_config()
+    model = AutoRegressiveLM(config).to(device)
+    wrapped = ForwardOnlyWrapper(model)
+    wrapped.train()
 
-    # Very skewed routing should give higher loss
-    skewed = torch.zeros(64, 8)
-    skewed[:, 0] = 1.0
-    skewed[:, 1] = 1.0
-    skewed = skewed / skewed.sum(dim=-1, keepdim=True)
-    skewed_loss = _load_balancing_loss(skewed).item()
+    strategy = SEQStrategy(wrapped, device, moe_aux_loss_coef=0.01)
+    output = strategy.compute_loss_output(
+        {
+            "input_ids": torch.randint(0, config.vocab_size, (2, 8)),
+            "target_ids": torch.randint(0, config.vocab_size, (2, 8)),
+        }
+    )
 
-    assert loss < skewed_loss
-
-
-# ── SEQStrategy integration tests ────────────────────────────────────
+    assert "moe_aux_loss" in output["metrics"]
+    assert "router_entropy" in strategy._moe_metrics
 
 
 class TestSEQStrategyMoE:
@@ -255,9 +257,6 @@ class TestSEQStrategyMoE:
         assert strategy._moe_metrics == {}
 
 
-# ── SFTStrategy integration tests ────────────────────────────────────
-
-
 class TestSFTStrategyMoE:
     """End‑to‑end tests for SFTStrategy with MoE aux loss."""