Module expert_cache 

The global expert slot cache: which experts are resident on the GPU right now, and what one decode step has to move.

One id space, one pool

Experts are addressed by a flat id, layer * num_experts + expert, and all layers compete for one pool of cache_size slots. That is the “global” in global LRU, and it is the point: expert activation is not uniform across layers, so a per-layer cache of cache_size / num_layers slots wastes residency on layers whose routing is flat and starves the layers where a few experts take most of the traffic.

What a step produces

ExpertCache::ensure takes the experts a layer routed to and returns, for each, the slot it will be read from – plus a CopyPlan: the (slot, host row) pairs the caller must copy before the multiply. Nothing here moves bytes; the plan is the whole output.

Two rules that must not drift

• Victims are chosen by (usage, slot) ascending, and a slot touched by this very step is not a victim. Without the second rule a step that misses more experts than it hits can evict an expert it is about to read.
• Duplicate routes collapse. A batch routing twice to the same expert counts one miss, issues one copy, and shares one slot.

The hybrid entry point adds the crate::qstar split on top: only the first fetch misses get slots, and the rest come back as None, meaning “compute this one on the CPU”. Every route is assigned exactly once, to exactly one device.

Reading the cache back

ExpertCache::stats answers “is the cache big enough?” for the model as a whole. ExpertCache::layer_stats answers it per MoE layer, which is the only form that can tell one layer thrashing from every layer uniformly a little over budget: both show the same global miss rate and they have opposite fixes – rebalance versus grow. ExpertCache::routing_skew goes one level below the cache and reports what the routing distribution itself allows; a layer whose oracle_hit_at_slots is already low cannot be helped by any cache size, because no policy could do better on that traffic.

The prefill double buffers

A prefill chunk walks the layers in order, so it can stage layer L + 1‘s experts while layer L computes. The two staging buffers borrow slots [0, 2 * num_experts) of this very cache and rotate by layer % 2. Borrowed is not owned: the bytes in those slots are rewritten every other layer within a chunk, so a slot at or below 2 * num_experts can never be read as residency, and the buffers’ contents are never registered in the residency map. ExpertCache::prefetch_prefill_layer is the whole portable half of that – which expert rows are already on the device, which have to cross the link, and what the borrowed slots do to the LRU. No streams, no events and no copies live here.

Ported 1:1 from FreeToken’s moe/offload_cache.py and moe/offload_kernels.py, whose LRU is the lru_ensure kernel from flashlib (both Apache-2.0); see docs/THIRD_PARTY_NOTICES.md.

Structs

BankEntry

One host-to-device transfer entry: rows consecutive expert rows of one bank.

CopyPlan

The copies a step needs before it can multiply.

EnsurePlan

What one ensure decided.

ExpertCache

The GPU expert cache’s residency map.

ExpertCacheStats

Running counters, for /metrics and for deciding whether the cache is sized right.

ExpertId

A cached expert’s address in the flat id space.

GatherPlan

Device-to-device rows: what the prefill buffer can take from the cache instead of from the host.

LayerRoutingSkew

One MoE layer’s routing concentration over the observed decode histogram.

MissRun

A contiguous run of expert ids, [start, start + len).

PrefillPlan

What staging one layer into a prefill double buffer costs.

RebuildRejected

A slot resize the cache refused, having changed nothing.

RoutingSkewReport

The routing-skew report: per layer, plus the means over the layers that were actually routed to.

Enums

FetchOrder

Which miss gets the scarce fetch slots when the split caps them.

Constants

SMALL_BANK_FEAT_BYTES

Per-expert row size below which a host bank ships its whole layer as one transfer entry.