Struct ExpertCache 

pub struct ExpertCache { /* private fields */ }

The GPU expert cache’s residency map.

Implementations

impl ExpertCache

pub fn new(num_layers: usize, num_experts: usize, cache_size: usize) -> Self

A cold cache of cache_size slots shared by every layer.

cache_size must hold at least one whole layer: a prefill materializes a layer’s experts into slots 0..num_experts, and a cache that cannot hold one layer cannot serve a prefill at all.

pub fn with_fetch_order(self, order: FetchOrder) -> Self

pub fn num_layers(&self) -> usize

pub fn num_experts(&self) -> usize

pub fn cache_size(&self) -> usize

pub fn stats(&self) -> ExpertCacheStats

pub fn layer_stats(&self, layer: u32) -> ExpertCacheStats

The same counters as stats, attributed to one MoE layer over the current reset_stats-delimited window.

A single global miss rate cannot distinguish one layer thrashing from every layer being uniformly a little over budget: the two average out to the same number and want opposite fixes – move slots between layers, versus give the pool more slots. This is the breakdown that tells them apart. Sums over every layer equal stats exactly.

pub fn per_layer_stats(&self) -> &[ExpertCacheStats]

Every layer’s counters at once, indexed by MoE layer id.

pub fn reset_stats(&mut self)

Close the current stats window and open a new one.

The routing histogram is deliberately not cleared here: it estimates a distribution, and its value grows with the number of steps behind it, so a /metrics scrape that resets the rate counters must not also destroy it. Use reset_routing when the distribution itself is what went stale.

pub fn total_experts(&self) -> usize

Total routed experts across the model – the denominator the cache’s residency rate is quoted against.

pub fn slot_of(&self, layer: u32, expert: u32) -> Option<u32>

Which slot holds (layer, expert), if any. For tests and reporting; the hot path uses ensure.

pub fn resident_in(&self, slot: u32) -> Option<ExpertId>

Which expert occupies slot, if any.

pub fn resident_slots(&self) -> usize

Slots currently holding an expert.

pub fn forget_slot(&mut self, slot: u32) -> Option<ExpertId>

Drop one slot’s residency, returning what it used to hold.

The map here is a record of copies the caller was asked to make, and a copy can fail after the map has recorded it – see crate::expert_slots::SlotFault::Device. Without this the next step reads that slot as a hit and multiplies whatever the failed copy left in it; with it, the expert simply misses again and is re-fetched.

The slot becomes the first eviction candidate rather than merely an empty one: its LRU stamp goes to the beginning of time, so a pool under pressure spends it before evicting an expert that is really there. Idempotent, and None for a slot that already held nothing.

pub fn reset(&mut self)

Forget everything. A rebuild is a cold start, so the counters go too – carrying them across would skew every rate that is quoted per call.

pub fn rebuild(&mut self, cache_size: usize) -> Result<(), RebuildRejected>

Resize the pool, keeping the model geometry.

Everything resident is dropped: slot ids are positions in an allocation that no longer exists, so keeping the map would point at other experts’ bytes. Refusing an impossible target before touching anything is deliberate – the caller can then keep serving from the cache it already has.

pub fn ensure(&mut self, layer: u32, expert_ids: &[u32]) -> EnsurePlan

Make every expert this layer routed to resident, evicting by LRU.

Every route gets a slot: this is the pure-offload path, where the CPU computes nothing.

pub fn ensure_hybrid( &mut self, layer: u32, expert_ids: &[u32], policy: &QStarPolicy, ) -> EnsurePlan

The bandwidth-adaptive path: fetch what the q* split says to fetch, and hand the rest back for the CPU.

A route that comes back None is the caller’s to compute from host RAM. Together with the Some routes it covers every routed expert exactly once – which is what makes it safe to simply add the two partial results.

pub fn materialize_layer(&mut self, layer: u32) -> CopyPlan

Put a whole layer’s experts in slots 0..num_experts, in expert order, for a prefill.

A prefill touches every expert, so streaming them one miss at a time is pointless – the layer is copied whole, and position equals expert id, which lets routing ids index the buffer directly with no slot lookup at all.

The bookkeeping still has to be exact: any other layer’s expert that was living in one of those slots loses its residency here, or the next decode step would hit on a slot that now holds someone else’s bytes.

pub fn set_collect_routing(&mut self, on: bool)

Start (or stop) accumulating the decode routing histogram.

Off by default: the counters are cheap but they are only meaningful over a long, stationary window, and a histogram that silently spans a model swap or a rebuild is worse than none.

pub fn collects_routing(&self) -> bool

pub fn routing_histogram(&self, layer: u32) -> &[u64]

One layer’s raw histogram row, indexed by expert id.

pub fn reset_routing(&mut self)

Throw the observed routing distribution away.

pub fn routing_skew(&self) -> Option<RoutingSkewReport>

Per-layer routing concentration over the observed histogram, or None if no routing has been observed at all.

This describes the traffic, not the cache: nothing here depends on which policy ran or on what it happened to keep. That is the point. The realized miss rate says how a policy did; LayerRoutingSkew::oracle_hit_at_slots says how well the best possible policy could have done on the same traffic with the same fair share of slots. A layer where the two are close is not a cache-sizing problem however bad it looks, because the routing is flat and the slots are not where the hit rate went.

pub fn prefill_buffer_slots(&self) -> usize

The slot range the two prefill buffers borrow: [0, 2 * num_experts).

pub fn prefill_overlap_fits(&self) -> bool

Whether this pool is large enough to lend the buffers their slots at all.

A cache of exactly 2 * num_experts fits the buffers but leaves no hit region above them, so every prefill row is a miss by the classification rule below. That is correct, merely slow; below 2 * num_experts the buffers do not fit and overlap must be off.

pub fn prefill_buffer_layer(&self, buffer_id: u32) -> Option<u32>

Which layer a buffer currently stages, if any.

pub fn prefill_hit_rows(&self) -> u64

Expert rows served from the cache since the last reset_stats.

pub fn prefill_rows(&self) -> u64

All expert rows staged into the buffers over the same window. The ratio against prefill_hit_rows is how much of a prefill never touched the link.

pub fn begin_prefill(&mut self)

Open a prefill chunk: both buffers empty, and take the residency snapshot the chunk classifies against.

The snapshot is what makes the classification stable for the whole chunk. Hits are decided once, against the map as it stood before any of this chunk’s staging ran, and the only writer inside the chunk – buffer invalidation – only ever rewrites slots that are already below the buffer threshold and therefore misses under both the live map and the snapshot. So the two can never disagree about a row, and a caller may issue the gathers and the host transfers in any order.

pub fn prefetch_prefill_layer( &mut self, layer: u32, bank_feat_bytes: &[u64], ) -> PrefillPlan

Stage one layer into its buffer and say what that costs.

bank_feat_bytes is each host bank’s per-expert row size, in the caller’s bank order; BankEntry::bank indexes it.

Three rules decide the answer, and each of them is a correctness rule rather than a tuning choice:

• A row is resident only if its slot is at or above 2 * num_experts. Anything below that – including -1 – is a miss by definition, because the buffers own those slots and rewrite them every other layer within the chunk. Widening the test to “has a slot at all” makes a prefill gather from slots the other buffer is concurrently overwriting, and it loads some other expert’s bytes without any error.
• Invalidation clears the residency map and zeroes usage for the buffer’s slots. See prefill_buffer_slots.
• Misses coalesce into contiguous expert runs, and a bank under SMALL_BANK_FEAT_BYTES ships its whole layer as one entry even when nothing missed.

Nothing staged here is registered as resident: the buffer’s bytes are volatile within the chunk, so recording them would hand the next decode step a hit on a slot that is about to be overwritten.

pub fn release_prefill_layer(&mut self, layer: u32)

Mark a layer’s buffer done, freeing it for the next layer of the same parity. A layer that is not the one staged is ignored.

Trait Implementations

impl Debug for ExpertCache

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.