feat(moe): --moe-collect-stats, so expert-cache behaviour is measurable

python/freetoken/engine/engine.py

@@ -30,6 +30,11 @@
 
 logger = init_logger(__name__)
 
+# Decode steps between --moe-collect-stats reports. Reading the counters costs a host sync,
+# so the window is coarse; it is short enough that a normal-length reply still produces a
+# few reports rather than one at the very end.
+MOE_STATS_INTERVAL = 256
+
 
 def _require_offload_cache_size(cache_size: int, num_experts: int) -> None:
     """The offload MoE cache needs at least one slot per expert per layer. A too-small size
@@ -332,6 +337,7 @@ def __init__(self, config: EngineConfig):
         # graphs, or other processes. Cross-rank MIN, deterministic across ranks.
         self._post_weights_free = post_weights_free
         self.moe_offload_cache = None
+        self._moe_stats_step = 0
         self.cpu_moe_executor = None
         if is_offload_moe_backend(config.moe_backend):
             self._init_offload_moe_cache(config)
@@ -628,6 +634,12 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         # Must be set before CUDA graph capture so the (device-side) accumulation ops are
         # captured and re-run on every decode replay.
         cache.collect_stats = config.moe_collect_stats
+        # The routing histogram rides the same switch: on its own the miss rate says how
+        # often we fetch, but not whether a smarter policy could have avoided the fetch.
+        # decode_routing_stats turns it into an oracle hit rate -- the ceiling any policy
+        # holding this many slots could reach on the observed routing -- which is the number
+        # worth having before anyone rewrites eviction.
+        cache.collect_decode_freq = config.moe_collect_stats
         # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's
         # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4).
         layers = attach_offload_moe_cache(self.model, cache)
@@ -932,8 +944,63 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True)
         copy_done_event = torch.cuda.Event()
         copy_done_event.record(self.stream)
+        if self.moe_offload_cache is not None and self.moe_offload_cache.collect_stats and batch.is_decode:
+            self._moe_stats_step += 1
+            if self._moe_stats_step >= MOE_STATS_INTERVAL:
+                self._moe_stats_step = 0
+                self._emit_moe_stats()
         return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
 
+    def _emit_moe_stats(self) -> None:
+        """Report one window of expert-cache behaviour, then reset the miss counters.
+
+        Accumulation is device-side and captured into the decode graph, so it is free to
+        leave running; reading it is not (the counters have to come back to the host), which
+        is why this only fires every MOE_STATS_INTERVAL decode steps. The routing histogram
+        is deliberately *not* reset -- the oracle bound wants the whole run's distribution,
+        not one window's.
+        """
+        cache = self.moe_offload_cache
+        agg = cache.decode_miss_stats()
+        if not agg["layer_calls"]:
+            return
+        parts = [
+            f"miss_rate={agg['miss_rate']:.3f}",
+            f"active/layer={agg['active_per_layer']:.1f}",
+            f"missing/layer={agg['missing_per_layer']:.1f}",
+        ]
+        if cache.decode_target == "hybrid":
+            # How the misses split: PCIe-fetched to the GPU vs handed to the CPU kernels.
+            parts.append(f"fetch_rate={agg['fetch_rate']:.3f}")
+            parts.append(f"cpu/layer={agg['cpu_per_layer']:.1f}")
+        logger.info_rank0(f"MoE cache ({MOE_STATS_INTERVAL} decode steps): " + ", ".join(parts))
+
+        per_layer = [L for L in cache.decode_miss_stats_per_layer()["per_layer"] if L["steps"]]
+        worst = sorted(per_layer, key=lambda L: -L["miss_rate"])[:5]
+        if worst:
+            logger.info_rank0(
+                "MoE cache worst layers: "
+                + ", ".join(f"L{L['layer']}={L['miss_rate']:.3f}" for L in worst)
+            )
+        routing = cache.decode_routing_stats()
+        if routing:
+            # oracle_hit_at_slots is the upper bound on hit rate for *any* policy with this
+            # many slots per layer. If it sits near the realized hit rate, the cache is
+            # already doing as well as the routing allows and the win has to come from
+            # somewhere else (more slots, more bandwidth); if it sits far above, eviction
+            # policy is leaving something on the table.
+            logger.info_rank0(
+                "MoE routing: "
+                f"oracle_hit={routing['oracle_hit_at_slots']:.3f} "
+                f"(realized {1.0 - agg['miss_rate']:.3f}), "
+                f"slots/layer={routing['slots_per_layer']:.1f}, "
+                f"working_set={routing['working_set_mean']:.1f}"
+                f"/{routing['working_set_max']}, "
+                f"experts_for_90pct={routing['experts_for_90pct']:.1f}, "
+                f"norm_entropy={routing['norm_entropy']:.3f}"
+            )
+        cache.reset_stats()
+
     @torch.inference_mode()
     def _warmup_prefill(self) -> None:
         """Compile the Triton prefill path before the first real request.

python/freetoken/server/args.py

@@ -533,6 +533,19 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="The unified MoE cache eviction policy.",
     )
 
+    parser.add_argument(
+        "--moe-collect-stats",
+        action="store_true",
+        default=ServerArgs.moe_collect_stats,
+        help=(
+            "Log MoE expert-cache miss rate and routing skew during decode. The counters "
+            "are captured into the decode CUDA graph, so this can only be chosen at startup. "
+            "Measured cost is below noise (46.1 vs 45.8 tok/s median on Ornith-35B-A3B "
+            "IQ3_S), but it stays off by default since it is a diagnostic and the readout "
+            "costs a host sync."
+        ),
+    )
+
     parser.add_argument(
         "--moe-cpu-threads",
         type=int,

tests/moe/test_moe_collect_stats.py

@@ -0,0 +1,109 @@
+"""--moe-collect-stats: the flag, and the report it produces.
+
+The counters themselves are accumulated device-side inside ``ensure_experts`` and were
+already covered; what was missing until this flag existed was any way to turn them on from
+the command line or read them back. These tests cover that wiring -- the flag reaching
+``ServerArgs``, and the emit formatting the numbers and resetting the window afterwards.
+"""
+
+import contextlib
+import io
+from types import SimpleNamespace
+
+from freetoken.engine.engine import MOE_STATS_INTERVAL, Engine
+from freetoken.server.args import ServerArgs, parse_args
+
+
+def test_flag_is_registered_and_defaults_off():
+    """``--help`` short-circuits before the model is resolved, so this needs no checkpoint."""
+    buf = io.StringIO()
+    with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit):
+        parse_args(["--help"])
+    assert "--moe-collect-stats" in buf.getvalue()
+    # Off unless asked for: the counters ride in the decode CUDA graph and cost throughput.
+    assert ServerArgs.moe_collect_stats is False
+
+
+class _StubCache:
+    """Just enough cache to exercise the emit: the four readers plus the window reset."""
+
+    def __init__(self, decode_target="gpu", layer_calls=512):
+        self.collect_stats = True
+        self.decode_target = decode_target
+        self._layer_calls = layer_calls
+        self.reset_calls = 0
+
+    def decode_miss_stats(self):
+        return {
+            "layer_calls": self._layer_calls,
+            "active_per_layer": 8.0,
+            "missing_per_layer": 2.0,
+            "miss_rate": 0.25,
+            "fetched_per_layer": 1.5,
+            "cpu_per_layer": 0.5,
+            "fetch_rate": 0.75,
+            "prefill_hit_rows": 0,
+            "prefill_rows": 0,
+        }
+
+    def decode_miss_stats_per_layer(self):
+        return {
+            "per_layer": [
+                {"layer": 0, "steps": 4, "miss_rate": 0.5},
+                {"layer": 1, "steps": 4, "miss_rate": 0.1},
+                # steps == 0 means the layer never ran in this window; it must not be
+                # ranked as a 0.0-miss-rate "best" layer.
+                {"layer": 2, "steps": 0, "miss_rate": 0.0},
+            ]
+        }
+
+    def decode_routing_stats(self):
+        return {
+            "slots_per_layer": 56.7,
+            "working_set_mean": 173.1,
+            "working_set_max": 243,
+            "experts_for_90pct": 92.3,
+            "oracle_hit_at_slots": 0.764,
+            "norm_entropy": 0.813,
+        }
+
+    def reset_stats(self):
+        self.reset_calls += 1
+
+
+def _emit(cache, caplog):
+    engine = SimpleNamespace(moe_offload_cache=cache, _emit_moe_stats=None)
+    with caplog.at_level("INFO"):
+        Engine._emit_moe_stats(engine)
+    return "\n".join(r.getMessage() for r in caplog.records)
+
+
+def test_emit_reports_and_resets_the_window(caplog):
+    cache = _StubCache()
+    out = _emit(cache, caplog)
+    assert "miss_rate=0.250" in out
+    # The oracle bound is the whole point of the report: it says how much room a different
+    # eviction policy could possibly have.
+    assert "oracle_hit=0.764" in out
+    assert "(realized 0.750)" in out
+    # Ranked worst-first, and the layer that never ran is left out entirely.
+    assert "L0=0.500, L1=0.100" in out
+    assert "L2" not in out
+    assert cache.reset_calls == 1
+
+
+def test_hybrid_split_only_reported_for_hybrid(caplog):
+    assert "fetch_rate" not in _emit(_StubCache(decode_target="gpu"), caplog)
+    caplog.clear()
+    assert "fetch_rate=0.750" in _emit(_StubCache(decode_target="hybrid"), caplog)
+
+
+def test_idle_window_emits_nothing_and_keeps_counters(caplog):
+    """No decode ran, so there is nothing to report -- and nothing to reset either."""
+    cache = _StubCache(layer_calls=0)
+    assert _emit(cache, caplog) == ""
+    assert cache.reset_calls == 0
+
+
+def test_interval_is_a_sane_window():
+    assert MOE_STATS_INTERVAL >= 1