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..ad43d2f1 100644 --- a/python/freetoken/server/args.py +++ b/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, 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