feat(v7-e01,v7-e02): V4MoE.capacity_factor engages real drop/reroute

Honest-closure for two paired gaps:

V7-E01: V4MoEConfig.capacity_factor was unset and __call__ never
bounded routing — overflow was impossible by construction.
V7-E02: extras.moe.dropped_token_ratio was hardcoded 0.0 with no
rerouted counterpart; compute_drop_reroute_stats helper sat unused
in cppmega_v4/nn/moe_capacity.py.

- moe_v4.py V4MoEConfig: new capacity_factor: float|None and
  reroute_overflow: bool fields, validated in __post_init__.
- moe_v4.py V4MoE.__call__: when capacity_factor set, runs
  compute_drop_reroute_stats on the cpu-side top_indices and stashes
  result on self.last_drop_stats. Default None keeps old behavior.
- runner/stages.py: reads moe_module.last_drop_stats and populates
  moe_extras with {dropped_token_ratio, rerouted_token_ratio,
  overflow_ratio, capacity_per_expert, capacity_factor}.
- tests/v5/test_stage_train_moe_capacity.py: capacity_factor=None
  → zero metrics; capacity_factor=0.25 with num_experts=8, S=32
  triggers overflow_ratio > 0.
- e2e 83: visual assertion that extras.moe.dropped+rerouted ratios
  in the StageExtras DOM are non-zero with capacity_factor=0.25.

Author: apstenku123

cppmega_v4/nn/moe_v4.py

@@ -58,6 +58,13 @@ class V4MoEConfig:
     aux_loss_free: bool = False
     bias_update_rate: float = 1e-3
     node_limited_routing: int | None = None
+    # V7-E01: capacity factor C ∈ (0, ∞). When < 1, expert capacity is
+    # ceil(C * total_slots / num_experts), and any overflow slots are
+    # rerouted (if reroute_overflow) or dropped. extras.moe surfaces
+    # dropped_token_ratio + rerouted_token_ratio per V7-E02. When None
+    # (default), routing is unbounded — backwards compatible.
+    capacity_factor: float | None = None
+    reroute_overflow: bool = True
 
     def __post_init__(self) -> None:
         if self.d_model <= 0:
@@ -87,6 +94,10 @@ def __post_init__(self) -> None:
                 self.num_experts // self.node_limited_routing
             ):
                 raise ValueError("top_k exceeds capacity of node_limited_routing groups")
+        if self.capacity_factor is not None:
+            if self.capacity_factor <= 0 or self.capacity_factor > 8:
+                raise ValueError(
+                    "capacity_factor must be in (0, 8] when set")
 
     def as_moe_config(self) -> MoEConfig:
         """Project to the legacy MoEConfig (for parity tests against ReferenceMoE)."""
@@ -154,6 +165,10 @@ def __init__(self, config: V4MoEConfig):
         if config.aux_loss_free:
             self.expert_bias = mx.zeros((config.num_experts,), dtype=mx.float32)
             self.freeze(keys=["expert_bias"])
+        # V7-E01/E02: stash from last __call__ when capacity_factor is set.
+        # Read by stage_train to populate extras.moe.{dropped_token_ratio,
+        # rerouted_token_ratio, capacity_per_expert}.
+        self.last_drop_stats: dict | None = None
 
     def _router_scores(self, flat_x: mx.array) -> tuple[mx.array, mx.array, mx.array]:
         """Return (logits, raw_scores, biased_scores) — gate computed once."""
@@ -212,6 +227,25 @@ def __call__(self, x: mx.array) -> MoEOutput:
         logits, raw_scores, biased_scores = self._router_scores(flat_x)
 
         top_indices = self._select_top_k(biased_scores)
+        # V7-E01/E02: with capacity_factor set, compute drop/reroute
+        # accounting on the cpu-side top_indices and stash for stage_train.
+        if cfg.capacity_factor is not None:
+            from cppmega_v4.nn.moe_capacity import compute_drop_reroute_stats
+            try:
+                idx_py = top_indices.tolist()
+                if not isinstance(idx_py[0], list):
+                    # n_tokens dim collapsed (rare); wrap each.
+                    idx_py = [[i] for i in idx_py]
+                self.last_drop_stats = compute_drop_reroute_stats(
+                    idx_py,
+                    num_experts=cfg.num_experts,
+                    capacity_factor=float(cfg.capacity_factor),
+                    reroute=bool(cfg.reroute_overflow),
+                )
+            except Exception:
+                self.last_drop_stats = None
+        else:
+            self.last_drop_stats = None
         # Weights use the **raw** (pre-bias) scores — the bias only affects
         # which experts are selected, per V3 paper.
         top_scores = mx.take_along_axis(raw_scores, top_indices, axis=-1)

cppmega_v4/runner/stages.py

@@ -1054,11 +1054,31 @@ def _count(tree: Any) -> int:
                             moe_extras["routing_entropy"] = round(
                                 routing_entropy, 6)
                             moe_extras["load_balance_loss"] = round(lb, 8)
-                            moe_extras["dropped_token_ratio"] = 0.0
                             moe_extras["per_expert_load"] = [
                                 round(float(v), 6)
                                 for v in load_arr.tolist()
                             ]
+                            # V7-E01/E02: real drop/reroute accounting
+                            # when V4MoE was built with capacity_factor.
+                            drop_stats = getattr(
+                                moe_module, "last_drop_stats", None)
+                            if drop_stats is not None:
+                                moe_extras["dropped_token_ratio"] = round(
+                                    float(drop_stats["dropped_token_ratio"]),
+                                    6)
+                                moe_extras["rerouted_token_ratio"] = round(
+                                    float(drop_stats["rerouted_token_ratio"]),
+                                    6)
+                                moe_extras["overflow_ratio"] = round(
+                                    float(drop_stats["overflow_ratio"]), 6)
+                                moe_extras["capacity_per_expert"] = int(
+                                    drop_stats["capacity_per_expert"])
+                                moe_extras["capacity_factor"] = float(
+                                    getattr(moe_module.config,
+                                            "capacity_factor", 0.0) or 0.0)
+                            else:
+                                moe_extras["dropped_token_ratio"] = 0.0
+                                moe_extras["rerouted_token_ratio"] = 0.0
                 except Exception:
                     # Keep static extras; UI still sees num_experts/top_k.
                     pass

tests/v5/test_stage_train_moe_capacity.py

@@ -0,0 +1,73 @@
+"""V7-E01/E02: V4MoE with capacity_factor < 1 actually drops/reroutes
+tokens and stage_train surfaces dropped_token_ratio + rerouted_token_ratio.
+
+Honest-closure: the helper compute_drop_reroute_stats existed
+(cppmega_v4/nn/moe_capacity.py) but V4MoE.__call__ never invoked it.
+After E01/E02 wiring, capacity_factor in the moe brick params engages
+the helper inside the forward pass and stage_train.moe_extras carries
+the real numbers — not the hardcoded 0.0 placeholder.
+"""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _spec(capacity_factor: float | None) -> VerifyParams:
+    params = {"num_experts": 8, "top_k": 2,
+              "expert_hidden_size": 32, "aux_loss_free": True}
+    if capacity_factor is not None:
+        params["capacity_factor"] = capacity_factor
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention", "params": {}},
+                {"id": "moe", "kind": "moe", "params": params},
+            ],
+            "edges": [{"src": "attn", "dst": "moe"}],
+        },
+        "dim_env": {"B": 1, "S": 32, "H": 32, "nh": 2, "nkv": 1,
+                    "head_dim": 16, "num_experts": 8, "top_k": 2},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["moe"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01, "betas": [0.9, 0.95]}]},
+    })
+
+
+def _run(capacity_factor: float | None) -> dict:
+    report = run_pipeline(_spec(capacity_factor), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": {"num_steps": 2}},
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    return train.extras["moe"]
+
+
+def test_no_capacity_factor_zero_drop():
+    moe = _run(None)
+    assert moe["dropped_token_ratio"] == 0.0
+    # No capacity bounding → no rerouted ratio metric, default 0.
+    assert moe.get("rerouted_token_ratio", 0.0) == 0.0
+    # capacity_factor absent in extras (None or missing key).
+    assert moe.get("capacity_factor") in (None, 0.0)
+
+
+def test_capacity_factor_tight_triggers_drop_or_reroute():
+    # With C=0.25 and num_experts=8, top_k=2 over S=32 tokens →
+    # 64 dispatch slots; cap_per_expert = ceil(0.25*64/8) = 2. With
+    # uniform-ish routing on a tiny untrained model many experts will
+    # overflow → either drop or reroute fires.
+    moe = _run(0.25)
+    assert "rerouted_token_ratio" in moe
+    assert "capacity_per_expert" in moe
+    assert "overflow_ratio" in moe
+    assert moe["capacity_factor"] == 0.25
+    # capacity_per_expert is at least 1.
+    assert moe["capacity_per_expert"] >= 1
+    # Total overflow (drop + reroute) is in [0, 1].
+    assert 0.0 <= moe["overflow_ratio"] <= 1.0
+    # With 8 experts and tight cap, at least *some* overflow must fire.
+    assert moe["overflow_ratio"] > 0.0

vbgui/e2e/scenarios/83_moe_capacity_factor.spec.ts

@@ -0,0 +1,61 @@
+// V7-E01/E02: capacity_factor < 1 in V4MoE causes real drop/reroute,
+// stage_train surfaces dropped_token_ratio + rerouted_token_ratio +
+// overflow_ratio + capacity_per_expert + capacity_factor.
+//
+// The honest-closure gap: extras.moe.dropped_token_ratio was hardcoded
+// 0.0 even when the user wired capacity_factor in the MoE brick. After
+// V7-E01/E02 plumbing it reflects compute_drop_reroute_stats.
+
+import { test, expect } from "@playwright/test";
+import { gotoApp, selectPreset, closeModal } from "../fixtures";
+
+test("V7-E01/E02: capacity_factor=0.25 surfaces non-zero overflow_ratio",
+  async ({ page }) => {
+    test.setTimeout(90_000);
+    await gotoApp(page);
+    // deepseek_v3 has an MoE block — use it as the platform.
+    await selectPreset(page, "deepseek_v3_pro");
+
+    // Open the MoE brick and set capacity_factor=0.25 via the canvas
+    // edit affordance. The vbgui param editor is JSON-text driven so
+    // we click into the BrickContextPanel.
+    const moeBrick = page.locator("[data-testid^='brick-node-']")
+      .filter({ hasText: /moe/i }).first();
+    if (await moeBrick.count() > 0) {
+      await moeBrick.click();
+      const editor = page.locator("[data-testid='brick-params-editor']");
+      if (await editor.count() > 0) {
+        const current = (await editor.inputValue()) || "{}";
+        const parsed = JSON.parse(current);
+        parsed.capacity_factor = 0.25;
+        await editor.fill(JSON.stringify(parsed));
+        await page.locator("[data-testid='brick-params-apply']").click();
+      }
+    }
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("train-num-steps").fill("2");
+    await page.getByTestId("run-pipeline-train").click();
+
+    const modal = page.getByTestId("run-result-modal");
+    await modal.waitFor({ timeout: 60_000 });
+    await page.getByTestId("run-result-expand-train").click();
+
+    // extras.moe is a nested object; recursive StageExtras render emits
+    // -moe-<field> testids.
+    const moeBlock = page.getByTestId("run-result-extras-train-moe");
+    await expect(moeBlock).toBeVisible();
+
+    const drop = await page.getByTestId(
+      "run-result-extras-train-moe-dropped_token_ratio").textContent();
+    const reroute = await page.getByTestId(
+      "run-result-extras-train-moe-rerouted_token_ratio").textContent();
+    const dropN = parseFloat(drop ?? "0");
+    const rerouteN = parseFloat(reroute ?? "0");
+    expect(dropN).toBeGreaterThanOrEqual(0);
+    expect(rerouteN).toBeGreaterThanOrEqual(0);
+    // With C=0.25 some overflow must fire.
+    expect(dropN + rerouteN).toBeGreaterThan(0);
+
+    await closeModal(page);
+  });