From 3a1ee3cac9185352f8ec53c7f435c0be6ef20f8a Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Wed, 26 Aug 2026 19:25:20 -0700 Subject: [PATCH 1/2] feat(moe): --moe-collect-stats, so expert-cache behaviour is measurable The offload cache already accumulated everything needed to answer "is the expert cache doing well, and could a different policy do better" -- decode_miss_stats, decode_miss_stats_per_layer and decode_routing_stats were all implemented. None of them had a caller, and neither collect_stats nor collect_decode_freq had a way to be turned on from the command line, so in practice the numbers were unreachable. This wires them up behind one flag. Both collectors ride the same switch because the miss rate on its own only says how often we fetch, not whether the fetch was avoidable; 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. The counters are captured into the decode CUDA graph, so the flag has to be chosen at startup and costs a little decode throughput. Reading them costs a host sync, so the report fires every MOE_STATS_INTERVAL decode steps rather than every step. The routing histogram is deliberately not reset between windows -- the oracle bound wants the whole run's distribution. On Ornith-1.5-35B-A3B IQ3_S on an 8GB 4060 (2267 slots, 56.7 per layer) this reports a realized hit rate of 0.741-0.788 against an oracle of 0.710-0.764, i.e. LRU is already at the ceiling the routing allows and the remaining misses are capacity-bound, not policy-bound. That is the kind of conclusion the flag exists to let people reach on their own model instead of guessing. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/engine/engine.py | 67 +++++++++++++++++ python/freetoken/server/args.py | 11 +++ tests/moe/test_moe_collect_stats.py | 109 ++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+) create mode 100644 tests/moe/test_moe_collect_stats.py diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..a6a6dcb3 100644 --- a/python/freetoken/engine/engine.py +++ b/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. diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..b97624aa 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -533,6 +533,17 @@ 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. Counters are " + "captured into the decode CUDA graph, so this can only be chosen at startup, " + "and it costs a little decode throughput -- it is a diagnostic, not a default." + ), + ) + parser.add_argument( "--moe-cpu-threads", type=int, diff --git a/tests/moe/test_moe_collect_stats.py b/tests/moe/test_moe_collect_stats.py new file mode 100644 index 00000000..f960abbd --- /dev/null +++ b/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 From d5c9d958f0d72024dfde3705d4286ac45e161800 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Wed, 26 Aug 2026 19:37:05 -0700 Subject: [PATCH 2/2] docs(moe): replace the guessed --moe-collect-stats overhead with the measured one The help text claimed the flag "costs a little decode throughput". Measuring it on Ornith-1.5-35B-A3B IQ3_S puts single-request decode at a median 46.1 tok/s with the flag on against 45.8 with it off, i.e. the cost is below run-to-run noise. Say that instead of guessing. It stays off by default regardless, since it is a diagnostic and the periodic readout does cost a host sync. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/server/args.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b97624aa..ad43d2f1 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -538,9 +538,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None: action="store_true", default=ServerArgs.moe_collect_stats, help=( - "Log MoE expert-cache miss rate and routing skew during decode. Counters are " - "captured into the decode CUDA graph, so this can only be chosen at startup, " - "and it costs a little decode throughput -- it is a diagnostic, not a default." + "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." ), )