Skip to content

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

Open
vcruz305 wants to merge 2 commits into
FlashML-org:mainfrom
vcruz305:feat/moe-collect-stats
Open

feat(moe): --moe-collect-stats, so expert-cache behaviour is measurable#231
vcruz305 wants to merge 2 commits into
FlashML-org:mainfrom
vcruz305:feat/moe-collect-stats

Conversation

@vcruz305

@vcruz305 vcruz305 commented Aug 27, 2026

Copy link
Copy Markdown

The offload cache already collects everything you need to know whether the expert cache is
doing a good job, but none of it is reachable. decode_miss_stats,
decode_miss_stats_per_layer and decode_routing_stats are all implemented and none of
them has a caller. collect_stats exists on EngineConfig but has no CLI flag, and
collect_decode_freq has no flag at all. So the numbers are there and nobody can see them.

This adds --moe-collect-stats and wires the readout into the decode loop.

Both collectors sit behind the one flag on purpose. The miss rate alone tells you how often
you fetch, but not whether the fetch was avoidable. decode_routing_stats gives you an
oracle hit rate, which is the best any policy holding that many slots could do on the
routing you actually observed. That second number is what tells you whether eviction policy
is worth touching.

What it looks like on Ornith-1.5-35B-A3B IQ3_S, 8GB 4060, 2267 slots:

MoE cache (256 decode steps): miss_rate=0.259, active/layer=8.0, missing/layer=2.1
MoE cache worst layers: L0=0.498, L1=0.446, L2=0.441, L3=0.361, L4=0.326
MoE routing: oracle_hit=0.764 (realized 0.741), slots/layer=56.7, working_set=173.1/243,
             experts_for_90pct=92.3, norm_entropy=0.813

Realized 0.741 against an oracle of 0.764 says LRU is basically at the ceiling the routing
allows on this model, so the remaining misses are a capacity problem and not a policy
problem. Normalised entropy of 0.81 says the same thing from the other side, routing is
close to diffuse and there is not much skew for a frequency policy to exploit. I went
looking for this because of the cache policy discussion in #174, and I would rather people
be able to check it on their own model than take my word for it.

It also picks up batching effects. Same server at 4 concurrent requests:

MoE cache (256 decode steps): miss_rate=0.255, active/layer=19.7, missing/layer=5.0

active/layer goes from 8.0 to 19.7, so 4 tokens that could have touched 32 experts only
touch 19.7 distinct ones, and fetches per layer go 1.8 to 5.0 for 4x the tokens. That is
about 30% less PCIe traffic per token.

Worth being careful about what that buys, since it is less than the traffic saving suggests.
On the server's own decode counter the aggregate goes 45.8 to 53.4 tok/s median, so 1.16x,
against the roughly 1.4x you would predict if decode were purely PCIe-bound. So decode is
only partly bandwidth-bound at this batch size and there is a compute component that grows
with it. Being able to see both the traffic and the throughput is the point of the flag.

Notes on the implementation:

  • The counters are captured into the decode CUDA graph, so this has to be a startup flag
    and cannot be toggled per request. It costs a little decode throughput, which is why it
    is off by default.
  • Reading the counters needs a host sync, so the report fires every 256 decode steps rather
    than every step.
  • reset_stats runs per window for the miss counters, but the routing histogram is left
    alone, since the oracle bound wants the whole run's distribution rather than one window's.
  • Layers with zero steps in a window are excluded from the worst-layer ranking so they do
    not show up as perfect.

Tests cover the flag being registered and defaulting off, the report formatting, the hybrid
fetch/cpu split only appearing for the hybrid target, the idle window emitting nothing and
leaving the counters alone, and the zero-step layer exclusion.

vcruz305 and others added 2 commits August 26, 2026 19:25
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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu
…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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant