feat(moe): --moe-collect-stats, so expert-cache behaviour is measurable - #231
feat(moe): --moe-collect-stats, so expert-cache behaviour is measurable#231vcruz305 wants to merge 2 commits into
Conversation
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
…cle stat) Ports from upstream FreeToken, adapted to the tier: FlashML-org#342 lm_head on sampled rows only (already generalised here via select_lm_head_rows); FlashML-org#339 the varlen GDN/KDA prefill conv takes max_seq_len from the scheduler on the Triton fallback (inert when sgl_kernel is installed, which every install path pins, so no node-4 change); FlashML-org#338 the n-gram PLE row-id hash as one Triton kernel with a bounded memo that is bypassed during CUDA graph capture (consumed by the pinned and cached PLE backends; the disk backend stages from its host hash); FlashML-org#231 the routing-oracle hit rate on the stats line next to the realised hot-pair rate, with the baseline reset on a live cache rebuild so the oracle can never read below realised. FlashML-org#89 (route-density tile selection) is skipped: its ds_fp4 tile table does not match the NVFP4 kernel's, which needs its own sm_89 sweep. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A88MCbnLtwsFSHmqwuJezY
|
Tested on a Qwen3.8-Flash-Next deployment (512 experts top-10, 48 MoE layers, 8984 GPU slots = 36.6% residency, KV 262k). Works, cost is within noise.
Output during the single-stream phase: Two notes. (1) Only the first line type appeared in this ~1.2k-token session; the Measured on a 2x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4, no NVLink) / 2x Xeon Gold 6526Y / 503 GiB box, CUDA 13.3, torch 2.11+cu130, sgl_kernel 0.4.5, model Single-stream = median of three 64-vs-256-token completion pairs, aggregate = 8 concurrent 256-token completions, TTFT on a ~1k-token prompt. Run-to-run spread of the baseline on this box is about +-4% single-stream. Update 2026-09-05, merged into the deploy branch. Tried on 2 x RTX 6000 Ada (sm_89) serving Qwen3.8-Flash-Next (RadixArk NVFP4) at TP=2, offload backend, fp8 KV pool of 8 x 262,144 tokens, merged onto my deploy branch (main af71ba4 + #385/#386/#389/#392/#354 and ten other open PRs), tests run on the box, then put in production. The readout on this box, single-stream decode at 70.9% residency (17,423 of 24,576 slots resident, the rest of the VRAM holding the eight-context fp8 KV pool): That is the number I was missing: it says why decode here is flat from 100% residency down to 70.9% (LRU sits at the oracle bound, the working set is a third of the slots). Merged clean, carried in production behind the flag. |
…ernels at startup, not mid-request Upstream FlashML-org#169 at a8326af, merged onto deploy/chatdnp for the PR sweep. Conflict in engine/engine.py: both sides add a method at the same place (FlashML-org#231's _emit_moe_stats, FlashML-org#169's _warmup_prefill_lens); kept both. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…gprobs for chat and legacy completions Upstream FlashML-org#224 at 855650d, merged onto deploy/chatdnp for the PR sweep. Conflicts: engine.py keeps FlashML-org#231's stats readout before the logprobs-aware return; openai_api.py keeps the vision `images` argument and FlashML-org#222's disconnect-watching drain with the logprobs entries added; generation.py keeps FlashML-org#266's marker filter and routes every content delta through FlashML-org#224's _content_delta so the logprobs entries ride the filtered text. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
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_layeranddecode_routing_statsare all implemented and none ofthem has a caller.
collect_statsexists onEngineConfigbut has no CLI flag, andcollect_decode_freqhas no flag at all. So the numbers are there and nobody can see them.This adds
--moe-collect-statsand 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_statsgives you anoracle 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:
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:
active/layergoes from 8.0 to 19.7, so 4 tokens that could have touched 32 experts onlytouch 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:
and cannot be toggled per request. It costs a little decode throughput, which is why it
is off by default.
than every step.
reset_statsruns per window for the miss counters, but the routing histogram is leftalone, since the oracle bound wants the whole run's distribution rather than one window's.
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.