Skip to main content

memra_engine/
moe_cache.rs

1//! EDGE-1 §B: SLRU GPU expert-residency cache (MOE-SLRU-PLAN §B).
2//!
3//! Stage-1 `moe_ffn` re-stages EVERY routed expert EVERY token over PCIe into one scratch slot.
4//! The same ~15-20% of experts recur (the "hot expert" mass), so an SLRU residency cache makes
5//! the steady-state re-stage count -> ~0. The cache holds N fixed-address GPU slots (never
6//! re-allocated, never fragmented), a `BlockId -> slot` residency table, an SLRU eviction policy
7//! (probation + protected segments; the second-miss "ghost" admission filter was measured a net
8//! loss in both regimes and removed 2026-07-08 — first-miss admit is the policy) so a one-off cold
9//! expert can never evict a genuinely hot one.
10//!
11//! THE bit-identity property (MOE-SLRU-PLAN §B.3): a cache HIT and a MISS feed `qmatvec_view` the
12//! *same* block bytes — the only difference is whether the `memcpy_htod` ran. So the cache-hit
13//! weight path is byte-for-byte identical to stage-every-token. TWO gates pin this, and they
14//! cover different classes: `src/bin/kernel_check.rs` `d2-cache-bit-identity` pins ONE block of a
15//! real GGUF checkpoint (dtypes `IQ3_S | IQ4_XS | Q6_K | Q8_0` — everything else, NVFP4 included,
16//! takes its `cells.skip` arm), and `tests/glm5_moe_residency_gpu.rs` pins it END TO END on a
17//! glm5_next fixture for the safetensors NVFP4 macro-carrying class, in CI, without a checkpoint.
18//!
19//! QUANT-FORMAT AGNOSTIC, and that is load-bearing for the safetensors NVFP4 class (glm5_next /
20//! GLM-5.3-Flash, Step-3.7-Flash-NVFP4, the unsloth 35B-A3B ST class). A slot is `max_block_bytes`
21//! of opaque bytes keyed by `BlockId`; nothing here reads a qtype, a block stride, or a scale.
22//! The loader has already repacked modelopt NVFP4 (`weight` + per-16 `weight_scale`) into ONE
23//! contiguous per-expert block in memra's internal `block_nvfp4` layout
24//! (`nvfp4_repack::repack_modelopt_to_gguf`, `row_bytes = in_f / 64 * 36`), so an NVFP4 block is
25//! staged and hit exactly like a k-quant GGUF block. The per-expert `weight_scale_2` MACRO scale
26//! is NOT in the block — it rides `HostExps::macros` and is folded post-matmul by the MoE forward
27//! — so residency can never move it, and hit/miss stay bit-identical for macro-carrying banks too.
28//!
29//! Gated behind `MEMRA_MOE_CACHE`, **default ON since 2026-07-08** (`docs/FLAGS.md`: the row is
30//! spelled `MEMRA_MOE_CACHE=0` = stage-every-token, i.e. `=0` is the ROLLBACK, not the default).
31//! `Engine::moe_cache_enabled()` is `var("MEMRA_MOE_CACHE") != Ok("0")`. This line previously read
32//! "default off => current stage-every-token behavior", which was the pre-2026-07-08 state and had
33//! been stale for seven weeks; it was corrected in the glm53-flash bring-up lane (2026-08-28) after
34//! it was quoted as the live default in a placement plan.
35
36use crate::Engine;
37use crate::model::{ExpertKeepalive, ExpertSource};
38use crate::spill_pread::{PreadPool, PreadStats, ReadTicket, SpillIoMode};
39use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, HostSlice, SyncOnDrop};
40use std::collections::{BTreeMap, HashMap, HashSet};
41use std::sync::Arc;
42
43/// Which projection of an expert (gate/up/down are three distinct GGUF blocks per expert).
44pub const PROJ_GATE: u8 = 0;
45pub const PROJ_UP: u8 = 1;
46pub const PROJ_DOWN: u8 = 2;
47
48/// Residency key: expert `ex` of layer `layer` projection `proj` is a distinct block.
49#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
50pub struct BlockId {
51    pub layer: u16,
52    pub proj: u8,
53    pub ex: u16,
54}
55impl BlockId {
56    #[inline]
57    pub fn new(layer: u16, proj: u8, ex: u16) -> Self {
58        BlockId { layer, proj, ex }
59    }
60}
61
62/// Where a dispatched block landed (always a retained resident slot since the first-miss-admit
63/// policy, 2026-07-08 — the transient staging tier went with the ghost filter).
64#[derive(Clone, Copy, Debug)]
65pub enum DispatchSlot {
66    Resident(usize),
67}
68
69/// Intrusive-list constants: `NIL` terminates a list; `seg` tags which segment holds a slot.
70const NIL: u32 = u32::MAX;
71const SEG_NONE: u8 = 0;
72const SEG_PROBATION: u8 = 1;
73const SEG_PROTECTED: u8 = 2;
74
75/// Per-slot intrusive doubly-linked node (slot indices are the arena — one node per GPU slot,
76/// shared by every class's two segments; a slot is in at most one segment at a time).
77#[derive(Clone, Copy, Debug, PartialEq, Eq)]
78struct SlotLink {
79    prev: u32,
80    next: u32,
81    seg: u8,
82}
83impl SlotLink {
84    const fn none() -> Self {
85        SlotLink {
86            prev: NIL,
87            next: NIL,
88            seg: SEG_NONE,
89        }
90    }
91}
92
93/// One SLRU segment as an intrusive doubly-linked list (front = LRU, back = MRU — the exact
94/// order contract the previous VecDeque carried). Every operation the hit path needs is O(1):
95/// `push_back` (MRU insert), `pop_front` (LRU evict), `unlink` (promotion removal by slot id).
96/// This is the Q5 audit fix (research/sweep-audits-20260805/AUDIT.md item 2): the VecDeque
97/// `position()+remove()` promotion was O(n_slots) PER HIT once the cache filled — ~46k slots x
98/// ~850 hits/token = ~40M host ops/token in the spill regime (measured 48.5 -> 46.0 tok/s on
99/// the 35B rtx6000 decode). Same eviction decisions for the same access pattern; only the cost of
100/// locating/removing a slot changed.
101#[derive(Debug)]
102struct SlruList {
103    head: u32,
104    tail: u32,
105    len: usize,
106}
107impl SlruList {
108    const fn new() -> Self {
109        SlruList {
110            head: NIL,
111            tail: NIL,
112            len: 0,
113        }
114    }
115
116    /// MRU insert. The slot must not currently be in any segment.
117    fn push_back(&mut self, slot: usize, seg: u8, links: &mut [SlotLink]) {
118        debug_assert_eq!(
119            links[slot].seg, SEG_NONE,
120            "slot {slot} already in a segment"
121        );
122        let s = slot as u32;
123        links[slot] = SlotLink {
124            prev: self.tail,
125            next: NIL,
126            seg,
127        };
128        if self.tail != NIL {
129            links[self.tail as usize].next = s;
130        } else {
131            self.head = s;
132        }
133        self.tail = s;
134        self.len += 1;
135    }
136
137    /// LRU removal.
138    fn pop_front(&mut self, links: &mut [SlotLink]) -> Option<usize> {
139        if self.head == NIL {
140            return None;
141        }
142        let s = self.head as usize;
143        self.unlink(s, links);
144        Some(s)
145    }
146
147    /// O(1) removal by slot id (the promotion path). The slot must be a member of THIS list.
148    fn unlink(&mut self, slot: usize, links: &mut [SlotLink]) {
149        let l = links[slot];
150        debug_assert_ne!(l.seg, SEG_NONE, "unlink of slot {slot} not in a segment");
151        if l.prev != NIL {
152            links[l.prev as usize].next = l.next;
153        } else {
154            debug_assert_eq!(self.head, slot as u32);
155            self.head = l.next;
156        }
157        if l.next != NIL {
158            links[l.next as usize].prev = l.prev;
159        } else {
160            debug_assert_eq!(self.tail, slot as u32);
161            self.tail = l.prev;
162        }
163        links[slot] = SlotLink::none();
164        self.len -= 1;
165    }
166
167    /// Front-to-back (LRU-to-MRU) iteration — the victim-scan order of the old VecDeque.
168    fn iter<'a>(&self, links: &'a [SlotLink]) -> SlruIter<'a> {
169        SlruIter {
170            links,
171            cur: self.head,
172        }
173    }
174}
175
176struct SlruIter<'a> {
177    links: &'a [SlotLink],
178    cur: u32,
179}
180impl Iterator for SlruIter<'_> {
181    type Item = usize;
182    fn next(&mut self) -> Option<usize> {
183        if self.cur == NIL {
184            return None;
185        }
186        let s = self.cur as usize;
187        self.cur = self.links[s].next;
188        Some(s)
189    }
190}
191
192/// One fixed-address size class with an independent SLRU. Separating queues by capacity prevents a
193/// small mixed-layout block from consuming the scarce slots that can hold a larger block.
194struct SlotClass {
195    capacity: usize,
196    probation: SlruList,
197    protected: SlruList,
198    free: Vec<usize>,
199    protected_cap: usize,
200}
201
202impl SlotClass {
203    /// HIT promotion under a FULL class (SLRU rules 4-6): probation hit -> protected MRU
204    /// (demoting protected LRU back to probation MRU while over cap); protected hit -> bump to
205    /// MRU; a slot in neither segment (defensive) inserts at protected MRU. Every arm is O(1).
206    fn on_hit_full(&mut self, slot: usize, links: &mut [SlotLink]) {
207        match links[slot].seg {
208            SEG_PROBATION => {
209                self.probation.unlink(slot, links);
210                self.push_protected(slot, links);
211            }
212            SEG_PROTECTED => {
213                self.protected.unlink(slot, links);
214                self.protected.push_back(slot, SEG_PROTECTED, links); // MRU
215            }
216            // not in either segment (shouldn't happen for a resident slot) — treat as protected MRU
217            _ => self.push_protected(slot, links),
218        }
219    }
220
221    /// Push a slot to protected MRU; if protected exceeds its cap, demote its LRU front to probation.
222    fn push_protected(&mut self, slot: usize, links: &mut [SlotLink]) {
223        self.protected.push_back(slot, SEG_PROTECTED, links);
224        while self.protected.len > self.protected_cap {
225            if let Some(demoted) = self.protected.pop_front(links) {
226                self.probation.push_back(demoted, SEG_PROBATION, links);
227            } else {
228                break;
229            }
230        }
231    }
232
233    /// LRU victim (rule 7 per-class): probation front first, else protected front. O(1).
234    fn pop_lru(&mut self, links: &mut [SlotLink]) -> Option<usize> {
235        self.probation
236            .pop_front(links)
237            .or_else(|| self.protected.pop_front(links))
238    }
239
240    /// Remove `slot` from whichever segment holds it (O(1) via the seg tag). No-op if in neither.
241    fn unlink_from_segment(&mut self, slot: usize, links: &mut [SlotLink]) {
242        match links[slot].seg {
243            SEG_PROBATION => self.probation.unlink(slot, links),
244            SEG_PROTECTED => self.protected.unlink(slot, links),
245            _ => {}
246        }
247    }
248}
249
250/// SLRU GPU expert-residency cache. Slots remain fixed-address for the cache lifetime. Uniform
251/// models use one class; mixed-layout models may preallocate several exact-capacity classes.
252pub struct MoeSlotCache {
253    slots: Vec<CudaSlice<u8>>, // fixed GPU buffers; capacities live in `classes`
254    slot_class: Vec<usize>,    // slot index -> size-class index
255    classes: Vec<SlotClass>,
256    /// Per-slot intrusive SLRU node (prev/next/segment). One arena for all classes: a slot
257    /// belongs to exactly one class, and to at most one of that class's two segments.
258    links: Vec<SlotLink>,
259    occupant: Vec<Option<BlockId>>, // slots[s] currently holds occupant[s]  (the residency bitmask)
260    table: HashMap<BlockId, usize>, // BlockId -> slot index (O(1) residency lookup)
261    /// Exponentially aged online access scores for the optional mixed-layout LFU victim policy.
262    /// Scores survive eviction and perf-counter resets; an opt-in decode-epoch decay prevents a
263    /// batched prompt from permanently outweighing recent token-to-token reuse.
264    frequencies: HashMap<BlockId, f32>,
265    /// Copy-stream prefetches that have reserved a slot but are not visible in `table` until the
266    /// consumer inserts an explicit compute-stream wait for `ready`. Pending slots are absent from
267    /// both SLRU queues, so neither synchronous admission nor another prefetch can evict them.
268    pending: HashMap<BlockId, PendingBlock>,
269    /// Source owners whose copy completed submission but not yet DMA completion. They are reaped
270    /// only after the recorded copy-stream event reports complete.
271    inflight_sources: Vec<(Arc<CudaEvent>, ExpertKeepalive)>,
272    /// Owners for copies whose completion could not be proved. Kept until a whole-stream drain;
273    /// leaked with the GPU slots if teardown cannot establish safety.
274    quarantined_sources: Vec<ExpertKeepalive>,
275    /// Unique owners used by demand/fallback H2D on the compute stream. `stage_expert` receives a
276    /// raw byte slice, so cudarc cannot attach its own source-lifetime event. Retain each backing
277    /// allocation once until cache teardown instead of paying one CUDA event per miss.
278    compute_sources: HashMap<KeepaliveKey, ExpertKeepalive>,
279    /// Opt-in positioned-read backends. Pinned buffers remain owned here until their explicit
280    /// compute-stream completion events fire.
281    pread: Option<PreadPool>,
282    /// Known-next reads submitted to disk workers but not yet consumed by dispatch. They own pinned
283    /// buffers, not GPU slots; all CUDA submission remains on the caller thread.
284    worker_reads: HashMap<BlockId, WorkerRead>,
285    pread_requested: bool,
286    pread_fallbacks: u64,
287    /// Retained so an event-creation failure after copy submission can be drained again during
288    /// teardown. A slot touched by an unprovable copy is quarantined outside every cache queue.
289    copy_stream: Arc<CudaStream>,
290    copy_stream_unknown: bool,
291    compute_stream: Arc<CudaStream>,
292    compute_stream_unknown: bool,
293
294    n: usize,
295    max_block_bytes: usize,
296    size_aware: bool,
297    frequency_evict: bool,
298    frequency_decay: Option<f32>,
299    /// Relative LFU value of a NextN/MTP access. The MTP block is keyed at `u16::MAX` and is
300    /// latency-critical during speculative decode, but contributes only one layer of observations
301    /// versus the full trunk. Keep the neutral default; local fixed-residency profiling may raise
302    /// it after an exact throughput sweep.
303    mtp_frequency_weight: f32,
304    last_forward_layer: Option<u16>,
305    last_forward_t: usize,
306    /// Stable-residency mode for heterogeneous CPU/GPU expert execution. Once frozen, callers may
307    /// still read resident slots, but must stage cache misses through transient scratch instead of
308    /// changing which experts execute on each backend.
309    frozen: bool,
310
311    // --- LAUNCH-STRUCTURE STAGE 3 (2026-07-05): device-side expert-pointer indirection ---
312    /// Resident-block count per LAYER (all 3 projections summed). When a layer reaches
313    /// 3*n_expert every routed block of that layer is cache-resident at a fixed address, so the
314    /// whole layer can dispatch via the DEVICE pointer table with ZERO host routing (no router
315    /// DtoH, no per-layer stream sync — the round-trip stall the decode profile measured at
316    /// ~36us x 40 layers/token). Maintained by admit/evict.
317    per_layer: HashMap<u16, u32>,
318    /// Per-layer device pointer row [3, n_expert] of slot base addresses (u64), uploaded lazily
319    /// when the layer first reads as fully resident. Slots are fixed-address for the cache's
320    /// lifetime, so a row stays valid until an eviction touches that layer (which drops the row
321    /// -> re-upload on next full residency).
322    dev_rows: HashMap<u16, CudaSlice<u64>>,
323    /// Layers whose one-shot prewarm was already attempted (success or not) — spill rigs whose
324    /// free slots can't hold a full layer must not re-scan 3*n_expert blocks every token.
325    prewarm_tried: HashSet<u16>,
326
327    // --- §D.4 instrumentation ---
328    pub hits: u64,
329    pub misses: u64,
330    pub staged_bytes: u64, // total H2D bytes the cache caused (admit + first-miss transient)
331}
332
333struct PendingBlock {
334    slot: usize,
335    ready: Arc<CudaEvent>,
336    keepalive: Option<ExpertKeepalive>,
337}
338
339#[derive(Clone, Copy)]
340struct WorkerRead {
341    ticket: ReadTicket,
342    len: usize,
343}
344
345#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
346enum KeepaliveKey {
347    Pinned(usize),
348    Buffer(usize),
349    Mmap(usize),
350}
351
352impl KeepaliveKey {
353    fn from_owner(owner: &ExpertKeepalive) -> Self {
354        match owner {
355            ExpertKeepalive::Pinned(value) => Self::Pinned(Arc::as_ptr(value) as usize),
356            ExpertKeepalive::Buffer(value) => Self::Buffer(Arc::as_ptr(value) as usize),
357            ExpertKeepalive::Mmap(value) => Self::Mmap(Arc::as_ptr(value) as usize),
358        }
359    }
360}
361
362/// Exact-length view over one CUDA-pinned pool allocation. cudarc's raw `&[u8]` HostSlice waits
363/// for the whole stream before returning, while passing `PinnedHostSlice` would copy its full
364/// capacity. This wrapper submits exactly the expert prefix; the caller records and retains the
365/// completion event before the backing allocation can be reused.
366struct ExactPinnedPrefix<'a>(&'a [u8]);
367
368impl HostSlice<u8> for ExactPinnedPrefix<'_> {
369    fn len(&self) -> usize {
370        self.0.len()
371    }
372
373    unsafe fn stream_synced_slice<'a>(
374        &'a self,
375        _stream: &'a CudaStream,
376    ) -> (&'a [u8], SyncOnDrop<'a>) {
377        // SAFETY: the pread staging helpers record an explicit event immediately after the async
378        // memcpy and PreadPool retains both allocation and event until it completes.
379        (self.0, SyncOnDrop::Record(None))
380    }
381
382    unsafe fn stream_synced_mut_slice<'a>(
383        &'a mut self,
384        _stream: &'a CudaStream,
385    ) -> (&'a mut [u8], SyncOnDrop<'a>) {
386        panic!("ExactPinnedPrefix is a source-only HostSlice")
387    }
388}
389
390fn stage_on_copy_stream(
391    e: &Engine,
392    host_bytes: &[u8],
393    slot: &mut CudaSlice<u8>,
394) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
395    // Protect all earlier compute-stream users of a reused slot before the copy stream overwrites it.
396    let prior = match e.stream().record_event(None) {
397        Ok(prior) => prior,
398        Err(err) => return Err((err.into(), true)),
399    };
400    if let Err(err) = e.copy_stream.wait(&prior) {
401        return Err((err.into(), true));
402    }
403    match e.stage_expert_async(host_bytes, slot, 0) {
404        Ok(ready) => Ok(Arc::new(ready)),
405        Err(err) => {
406            // The H2D may have been submitted before event creation failed. Never release either the
407            // destination slot or pinned source until the copy stream has drained.
408            match e.copy_stream.synchronize() {
409                Ok(()) => Err((err, true)),
410                Err(sync_err) => Err((std::io::Error::other(format!(
411                    "copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
412                )).into(), false)),
413            }
414        }
415    }
416}
417
418fn stage_pread_on_compute_stream(
419    e: &Engine,
420    host_bytes: &[u8],
421    slot: &mut CudaSlice<u8>,
422) -> Result<Arc<CudaEvent>, Box<dyn std::error::Error>> {
423    let ready = Arc::new(e.ctx().new_event(None)?);
424    let source = ExactPinnedPrefix(host_bytes);
425    let mut dst = slot.slice_mut(0..host_bytes.len());
426    e.stream().memcpy_htod(&source, &mut dst)?;
427    ready.record(&e.stream())?;
428    Ok(ready)
429}
430
431fn stage_pread_prefetch_on_copy_stream(
432    e: &Engine,
433    host_bytes: &[u8],
434    slot: &mut CudaSlice<u8>,
435) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
436    let ready = match e.ctx().new_event(None) {
437        Ok(ready) => Arc::new(ready),
438        Err(err) => return Err((err.into(), true)),
439    };
440    let source = ExactPinnedPrefix(host_bytes);
441    let mut dst = slot.slice_mut(0..host_bytes.len());
442    let submitted = e
443        .copy_stream
444        .memcpy_htod(&source, &mut dst)
445        .and_then(|()| ready.record(&e.copy_stream));
446    match submitted {
447        Ok(()) => Ok(ready),
448        Err(err) => match e.copy_stream.synchronize() {
449            Ok(()) => Err((err.into(), true)),
450            Err(sync_err) => Err((
451                std::io::Error::other(format!(
452                "pread copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
453            ))
454                .into(),
455                false,
456            )),
457        },
458    }
459}
460
461/// Allocate the same fraction of every exact block-size class under one byte budget. This avoids
462/// biasing residency toward either low-bit or high-bit tiers while eliminating max-slot padding.
463fn size_class_plan(block_bytes: &[usize], budget_bytes: usize) -> Vec<(usize, usize)> {
464    let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
465    for &bytes in block_bytes.iter().filter(|&&bytes| bytes > 0) {
466        *counts.entry(bytes).or_insert(0) += 1;
467    }
468    if counts.is_empty() || budget_bytes == 0 {
469        return Vec::new();
470    }
471    let total_bytes: u128 = counts
472        .iter()
473        .map(|(&bytes, &count)| (bytes as u128 + 8) * count as u128)
474        .sum();
475    let budget = budget_bytes as u128;
476    let mut plan: Vec<(usize, usize, u128)> = counts
477        .iter()
478        .map(|(&bytes, &count)| {
479            let scaled = count as u128 * budget;
480            (
481                bytes,
482                (scaled / total_bytes).min(count as u128) as usize,
483                scaled % total_bytes,
484            )
485        })
486        .collect();
487    let mut used: u128 = plan
488        .iter()
489        .map(|(bytes, count, _)| (*bytes as u128 + 8) * *count as u128)
490        .sum();
491
492    // Hamilton-style remainder pass keeps class proportions close after flooring. There are only
493    // a handful of layout classes, so one additional slot per class covers all rounding loss.
494    let mut order: Vec<usize> = (0..plan.len()).collect();
495    order.sort_by(|&a, &b| plan[b].2.cmp(&plan[a].2).then(a.cmp(&b)));
496    for index in order {
497        let (bytes, count, _) = plan[index];
498        let available = counts[&bytes];
499        let required = bytes as u128 + 8;
500        if count < available && used + required <= budget {
501            plan[index].1 += 1;
502            used += required;
503        }
504    }
505    plan.into_iter()
506        .filter_map(|(bytes, count, _)| (count > 0).then_some((bytes, count)))
507        .collect()
508}
509
510impl MoeSlotCache {
511    /// Build the cache sizing N from free VRAM (MOE-SLRU-PLAN §B.4): probe free VRAM AFTER residents
512    /// are loaded; N is shared across ALL layers so it must hold the WHOLE-MODEL hot set, not one
513    /// layer's. The 35B-A3B keeps its 256 experts HOST-resident, so the GPU has ~20+ GB free at
514    /// decode — empirically a 256-slot cache thrashes (~2-7% hit) while a few-thousand-slot cache
515    /// reaches ~85%+ steady-state. So the DEFAULT auto-sizes N to fill `MEMRA_MOE_VRAM_FRAC` (default
516    /// 0.85) of free VRAM, clamped to [256, ~hot-set]. `MEMRA_MOE_SLOTS` forces an exact N.
517    pub fn new(e: &Engine, max_block_bytes: usize) -> Result<Self, Box<dyn std::error::Error>> {
518        let (free, _total) = e.ctx().mem_get_info()?;
519        // Keep two blocks of slack after the machine-specific hard ceiling. The default remains
520        // 80%; tightly provisioned spill rigs may raise it only after an OOM-gated local sweep.
521        let hard_frac = cache_hard_vram_frac();
522        let hard_bytes =
523            ((free as f64 * hard_frac) as usize).saturating_sub(2 * (max_block_bytes + 8));
524        let forced_slots = std::env::var("MEMRA_MOE_SLOTS")
525            .ok()
526            .and_then(|s| s.parse::<usize>().ok());
527        let requested_bytes = if let Some(n) = forced_slots {
528            n.saturating_mul(max_block_bytes + 8)
529        } else {
530            // auto: fill MEMRA_MOE_VRAM_FRAC of free VRAM with slots (default 85%).
531            // DEFAULT 0.85 (2026-07-06 local sweep: 0.40=25.0, 0.60=28.0, 0.85=28.5 tok/s on the
532            // spill-regime 35B — hit-rate 87.8% -> 99.2%, PCIe 55 -> 3.8 MB/tok; the 0.80
533            // hard-headroom cap below still bounds the true allocation, so 0.85 requests the max).
534            // Rigs co-running other GPU work should set MEMRA_MOE_VRAM_FRAC lower.
535            let frac = std::env::var("MEMRA_MOE_VRAM_FRAC")
536                .ok()
537                .and_then(|s| s.parse::<f64>().ok())
538                .unwrap_or(0.85);
539            (free as f64 * frac) as usize
540        };
541        let budget_bytes = requested_bytes.min(hard_bytes);
542        let layout = e.moe_cache_layout().unwrap_or_default();
543        let size_aware = forced_slots.is_none()
544            && std::env::var("MEMRA_MOE_SIZE_AWARE").as_deref() == Ok("1")
545            && !layout.is_empty();
546        let frequency_evict = std::env::var("MEMRA_MOE_LFU").as_deref() == Ok("1");
547        let frequency_decay = if frequency_evict {
548            cache_lfu_decay()
549        } else {
550            None
551        };
552        let mtp_frequency_weight = cache_lfu_mtp_weight();
553        let mut class_plan = if size_aware {
554            size_class_plan(&layout, budget_bytes)
555        } else {
556            Vec::new()
557        };
558        if class_plan.iter().map(|(_, count)| count).sum::<usize>() < 8 {
559            let n = (budget_bytes / (max_block_bytes + 8)).max(8);
560            class_plan = vec![(max_block_bytes, n)];
561        }
562        let n: usize = class_plan.iter().map(|(_, count)| count).sum();
563
564        let mut slots = Vec::with_capacity(n);
565        let mut slot_class = Vec::with_capacity(n);
566        let mut classes = Vec::with_capacity(class_plan.len());
567        let mut occupant = Vec::with_capacity(n);
568        for (class_index, &(capacity, count)) in class_plan.iter().enumerate() {
569            let start = slots.len();
570            for _ in 0..count {
571                // +8 tail pad: wide expert dots may issue an aligned read past the final block.
572                slots.push(e.alloc_u8(capacity + 8)?);
573                slot_class.push(class_index);
574                occupant.push(None);
575            }
576            let free_slots = (start..start + count).rev().collect();
577            classes.push(SlotClass {
578                capacity,
579                probation: SlruList::new(),
580                protected: SlruList::new(),
581                free: free_slots,
582                protected_cap: ((count as f64 * 0.8) as usize).max(1),
583            });
584        }
585        let links = vec![SlotLink::none(); n];
586        if size_aware {
587            let allocated: usize = class_plan
588                .iter()
589                .map(|(bytes, count)| (bytes + 8) * count)
590                .sum();
591            eprintln!(
592                "[moe-cache] size-aware fixed slots: {n} slots in {} classes, {:.2} GB / {:.2} GB budget",
593                class_plan.len(),
594                allocated as f64 / 1e9,
595                budget_bytes as f64 / 1e9
596            );
597        }
598        let pread_mode = crate::spill_pread::configured_mode();
599        let pread_requested = pread_mode != SpillIoMode::Mmap;
600        let pread = if pread_requested {
601            match PreadPool::try_new(e, max_block_bytes, pread_mode) {
602                Ok(pool) => Some(pool),
603                Err(err) => {
604                    eprintln!(
605                        "[spill-pread] pinned-buffer initialization failed ({err}); using mmap"
606                    );
607                    None
608                }
609            }
610        } else {
611            None
612        };
613
614        Ok(MoeSlotCache {
615            slots,
616            slot_class,
617            classes,
618            links,
619            occupant,
620            table: HashMap::with_capacity(n * 2),
621            frequencies: HashMap::with_capacity(layout.len().max(n * 2)),
622            pending: HashMap::new(),
623            inflight_sources: Vec::new(),
624            quarantined_sources: Vec::new(),
625            compute_sources: HashMap::new(),
626            pread,
627            worker_reads: HashMap::new(),
628            pread_requested,
629            pread_fallbacks: 0,
630            copy_stream: e.copy_stream.clone(),
631            copy_stream_unknown: false,
632            compute_stream: e.stream().clone(),
633            compute_stream_unknown: false,
634            n,
635            max_block_bytes,
636            size_aware,
637            frequency_evict,
638            frequency_decay,
639            mtp_frequency_weight,
640            last_forward_layer: None,
641            last_forward_t: 0,
642            frozen: false,
643            per_layer: HashMap::new(),
644            dev_rows: HashMap::new(),
645            prewarm_tried: HashSet::new(),
646            hits: 0,
647            misses: 0,
648            staged_bytes: 0,
649        })
650    }
651
652    #[inline]
653    pub fn n_slots(&self) -> usize {
654        self.n
655    }
656    #[inline]
657    pub fn is_frozen(&self) -> bool {
658        self.frozen
659    }
660    pub fn freeze(&mut self) {
661        if !self.frozen {
662            self.frozen = true;
663            let (_, complete, one_projection, two_projections, stranded_blocks) =
664                self.expert_residency_shape();
665            eprintln!(
666                "[moe-cache] residency frozen: {} slots, {} resident blocks; \
667                 {complete} complete experts, {one_projection} one-projection fragments, \
668                 {two_projections} two-projection fragments ({stranded_blocks} stranded blocks)",
669                self.n,
670                self.table.len()
671            );
672            let mut mtp_masks = HashMap::<u16, u8>::new();
673            for id in self.table.keys().filter(|id| id.layer == u16::MAX) {
674                *mtp_masks.entry(id.ex).or_insert(0) |= 1u8 << id.proj;
675            }
676            if !mtp_masks.is_empty() {
677                let complete = mtp_masks.values().filter(|&&mask| mask == 0b111).count();
678                eprintln!(
679                    "[moe-cache] frozen MTP residency: {} blocks, {complete} complete experts",
680                    mtp_masks
681                        .values()
682                        .map(|mask| mask.count_ones() as usize)
683                        .sum::<usize>()
684                );
685            }
686        }
687    }
688
689    pub(crate) fn expert_residency_shape(&self) -> (usize, usize, usize, usize, usize) {
690        let mut masks = HashMap::<(u16, u16), u8>::new();
691        for id in self.table.keys() {
692            *masks.entry((id.layer, id.ex)).or_insert(0) |= 1u8 << id.proj;
693        }
694        let complete = masks.values().filter(|&&mask| mask == 0b111).count();
695        let one_projection = masks
696            .values()
697            .filter(|&&mask| mask.count_ones() == 1)
698            .count();
699        let two_projections = masks
700            .values()
701            .filter(|&&mask| mask.count_ones() == 2)
702            .count();
703        let stranded_blocks = one_projection + 2 * two_projections;
704        (
705            masks.len(),
706            complete,
707            one_projection,
708            two_projections,
709            stranded_blocks,
710        )
711    }
712
713    #[inline]
714    pub fn max_block_bytes(&self) -> usize {
715        self.max_block_bytes
716    }
717
718    /// O(1) residency check (the ktransformers `generate_gpu_experts_masks` analog).
719    #[inline]
720    pub fn resident(&self, id: BlockId) -> Option<usize> {
721        self.table.get(&id).copied()
722    }
723
724    #[inline]
725    fn frequency_increment(&self, id: BlockId) -> f32 {
726        if id.layer == u16::MAX {
727            self.mtp_frequency_weight
728        } else {
729            1.0
730        }
731    }
732
733    /// Record a routed block that a fused all-hit path consumed without going through dispatch.
734    /// Warmup-only callers use this to make the LFU profile reflect actual grouped GPU traffic;
735    /// frozen serving skips it because residency can no longer change.
736    pub(crate) fn note_profile_hit(&mut self, id: BlockId) {
737        if self.frozen || !self.table.contains_key(&id) {
738            return;
739        }
740        let increment = self.frequency_increment(id);
741        *self.frequencies.entry(id).or_insert(0.0) += increment;
742    }
743
744    /// HIT promotion (SLRU): on a probation hit promote to protected; on a protected hit bump to MRU.
745    ///
746    /// O(1) EARLY-OUT (STAGING-ELISION stage, 2026-07-04): while FREE slots remain, `admit` pops
747    /// `free` and `evict_one` is unreachable — recency order is dead state until the cache fills.
748    /// Kept even though promotion is now O(1) either way (audit-fix Q5, 2026-08-06): skipping it
749    /// preserves the not-yet-full ordering behavior BYTE-FOR-BYTE with the pre-fix policy (slots
750    /// stay in admission order until the class fills), and on 96GB rigs (slots >= whole-model
751    /// block count) every HIT stays a pure table lookup forever.
752    ///
753    /// FULL-CLASS promotion was the Q5 audit item (research/sweep-audits-20260805/AUDIT.md
754    /// item 2): the old VecDeque `position()+remove()` was O(n_slots) PER HIT — at ~46k slots x
755    /// ~850 hits/token ~40M host ops/token, measured as the fast-admit A/B regression
756    /// 48.5 -> 46.0 tok/s on the 35B rtx6000 decode (the 2026-07-04 fix only DEFERRED the scan to
757    /// the spill regime, where the cache is permanently full). The intrusive-list rewrite makes
758    /// every arm O(1) with IDENTICAL eviction decisions for the same access pattern (list order
759    /// == the old VecDeque order at every step; unit-pinned by `slru_intrusive_tests`).
760    /// Bookkeeping-only: the dispatched bytes are identical either way (the D.2 gate pins it).
761    fn on_hit(&mut self, slot: usize) {
762        let class_index = self.slot_class[slot];
763        let class = &mut self.classes[class_index];
764        if !class.free.is_empty() {
765            return;
766        }
767        class.on_hit_full(slot, &mut self.links);
768    }
769
770    fn remove_occupant(&mut self, slot: usize) {
771        if let Some(old) = self.occupant[slot].take() {
772            self.table.remove(&old);
773            self.on_block_evicted(old.layer);
774        }
775    }
776
777    /// Lowest cumulative-frequency resident in one class; ties keep ordinary LRU order. A cold
778    /// admission therefore becomes the sacrificial slot on the next miss instead of displacing a
779    /// prompt-proven hot expert. `keep` protects the expert whose kernels are currently queued.
780    /// (Deliberately still O(n_slots) PER EVICTION — the opt-in LFU policy is a full-scan argmin
781    /// by definition; the Q5 fix targeted the per-HIT scan. Eviction order over the linked lists
782    /// == the old probation-then-protected VecDeque order.)
783    fn frequency_victim_in_class(&mut self, class_index: usize, keep: &[BlockId]) -> Option<usize> {
784        let class = &self.classes[class_index];
785        let candidate = class
786            .probation
787            .iter(&self.links)
788            .chain(class.protected.iter(&self.links))
789            .enumerate()
790            .filter_map(|(position, slot)| {
791                let id = self.occupant[slot]?;
792                (!keep.contains(&id)).then_some((
793                    self.frequencies.get(&id).copied().unwrap_or(0.0),
794                    position,
795                    slot,
796                ))
797            })
798            .min_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
799        let (_, _, slot) = candidate?;
800        self.classes[class_index].unlink_from_segment(slot, &mut self.links);
801        Some(slot)
802    }
803
804    /// Pick the LRU victim from the smallest class that can hold `required` bytes.
805    fn evict_one(&mut self, required: usize) -> Option<usize> {
806        for class_index in 0..self.classes.len() {
807            if self.classes[class_index].capacity < required {
808                continue;
809            }
810            let slot = if self.frequency_evict {
811                self.frequency_victim_in_class(class_index, &[])
812            } else {
813                self.classes[class_index].pop_lru(&mut self.links)
814            };
815            if let Some(slot) = slot {
816                self.remove_occupant(slot);
817                return Some(slot);
818            }
819        }
820        None
821    }
822
823    /// Pick a resident victim that is not needed by the expert currently being computed. Pending
824    /// slots never enter the SLRU queues, so they are excluded automatically. Returns `None` rather
825    /// than evicting a protected block; the caller then leaves this block to the synchronous path.
826    /// (LRU-to-MRU scan skipping `keep` members — same visit order as the old VecDeque scan;
827    /// O(n) worst per PREFETCH eviction only, unchanged from the pre-fix shape.)
828    fn evict_one_excluding(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
829        fn take(
830            q: &mut SlruList,
831            links: &mut [SlotLink],
832            occupant: &[Option<BlockId>],
833            keep: &[BlockId],
834        ) -> Option<usize> {
835            let slot = q
836                .iter(links)
837                .find(|&s| occupant[s].is_some_and(|id| !keep.contains(&id)))?;
838            q.unlink(slot, links);
839            Some(slot)
840        }
841        for class_index in 0..self.classes.len() {
842            if self.classes[class_index].capacity < required {
843                continue;
844            }
845            let slot = if self.frequency_evict {
846                self.frequency_victim_in_class(class_index, keep)
847            } else {
848                let class = &mut self.classes[class_index];
849                take(&mut class.probation, &mut self.links, &self.occupant, keep)
850                    .or_else(|| take(&mut class.protected, &mut self.links, &self.occupant, keep))
851            };
852            if let Some(slot) = slot {
853                self.remove_occupant(slot);
854                return Some(slot);
855            }
856        }
857        None
858    }
859
860    /// STAGE 3 bookkeeping: a resident block of `layer` was evicted — the layer is no longer fully
861    /// resident, so its device pointer row (if uploaded) must be invalidated. NOTE: the row's device
862    /// buffer is dropped here, which is safe because the fully-resident fast path is only taken when
863    /// `dev_rows` contains the layer at DISPATCH time and all launches consuming the row were
864    /// enqueued BEFORE this eviction's staging memcpy on the same stream (single-stream ordering).
865    fn on_block_evicted(&mut self, layer: u16) {
866        if let Some(c) = self.per_layer.get_mut(&layer) {
867            *c -= 1;
868        }
869        self.dev_rows.remove(&layer);
870    }
871
872    fn reserve_slot(&mut self, required: usize) -> Option<usize> {
873        for class in &mut self.classes {
874            if class.capacity >= required
875                && let Some(slot) = class.free.pop()
876            {
877                return Some(slot);
878            }
879        }
880        self.evict_one(required)
881    }
882
883    fn release_reserved_slot(&mut self, slot: usize) {
884        debug_assert!(self.occupant[slot].is_none());
885        self.classes[self.slot_class[slot]].free.push(slot);
886    }
887
888    fn publish(&mut self, id: BlockId, slot: usize) {
889        self.occupant[slot] = Some(id);
890        self.table.insert(id, slot);
891        self.classes[self.slot_class[slot]].probation.push_back(
892            slot,
893            SEG_PROBATION,
894            &mut self.links,
895        );
896        *self.per_layer.entry(id.layer).or_insert(0) += 1;
897    }
898
899    fn reap_copy_sources(&mut self) {
900        self.inflight_sources
901            .retain(|(ready, _)| !ready.is_complete());
902    }
903
904    fn retain_compute_source(&mut self, owner: Option<ExpertKeepalive>) {
905        if let Some(owner) = owner {
906            let key = KeepaliveKey::from_owner(&owner);
907            self.compute_sources.entry(key).or_insert(owner);
908        }
909    }
910
911    /// Admit a block: evict a victim, stage `host_bytes` into its slot, register residency, place in
912    /// probation (new admissions enter probation — they earn promotion on a later hit).
913    fn admit(
914        &mut self,
915        id: BlockId,
916        host_bytes: &[u8],
917        e: &Engine,
918    ) -> Result<usize, Box<dyn std::error::Error>> {
919        let slot = self.reserve_slot(host_bytes.len()).ok_or_else(|| {
920            std::io::Error::other(format!(
921                "no MoE cache slot can hold {} bytes (max class {})",
922                host_bytes.len(),
923                self.classes.last().map(|class| class.capacity).unwrap_or(0)
924            ))
925        })?;
926        // Pending copy-stream admissions are not in either SLRU queue, so `evict_one` cannot return
927        // an in-flight slot. This synchronous copy and its consumer remain ordered on gpu.stream.
928        if let Err(err) = e.stage_expert(host_bytes, &mut self.slots[slot], 0) {
929            return match e.stream().synchronize() {
930                Ok(()) => {
931                    self.release_reserved_slot(slot);
932                    Err(err)
933                }
934                Err(sync_err) => {
935                    // Keep the slot outside free/table/SLRU. Drop retries the stream drain and
936                    // leaks every slot if CUDA never provides a completion proof.
937                    self.compute_stream_unknown = true;
938                    Err(std::io::Error::other(format!(
939                        "compute-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
940                    )).into())
941                }
942            };
943        }
944        self.staged_bytes += host_bytes.len() as u64;
945        self.publish(id, slot);
946        Ok(slot)
947    }
948
949    fn note_pread_fallback(&mut self, reason: &dyn std::fmt::Display) {
950        self.pread_fallbacks += 1;
951        if let Some(pool) = self.pread.as_mut() {
952            pool.note_fallback();
953        }
954        if self.pread_fallbacks <= 3 {
955            eprintln!("[spill-pread] falling back to mmap: {reason}");
956        }
957    }
958
959    /// Start one MoE forward's worker-I/O scope. Any ticket left by an earlier error/early return is
960    /// no longer a valid lookahead target; cancel it before this scope submits its own known-next
961    /// reads. In-flight CPU reads keep their buffers until completion restores them safely.
962    pub(crate) fn begin_worker_scope(&mut self) {
963        if self.worker_reads.is_empty() {
964            return;
965        }
966        let tickets: Vec<_> = self
967            .worker_reads
968            .drain()
969            .map(|(_, read)| read.ticket)
970            .collect();
971        if let Some(pool) = self.pread.as_mut().filter(|pool| pool.is_worker()) {
972            for ticket in tickets {
973                let _ = pool.cancel_worker(ticket);
974            }
975        }
976    }
977
978    /// Age cumulative LFU at decode-token boundaries. A batched prompt may touch one block many
979    /// times before decode begins; treating those touches as permanent future-use votes poisons a
980    /// spill cache. The first T=1 sweep starts a fresh frequency epoch while preserving populated
981    /// GPU slots. Later decode sweeps exponentially age history so recent cross-token reuse can
982    /// displace stale prompt-specific experts.
983    ///
984    /// MoE layers are visited in ascending order and the cache is model-global, so
985    /// `layer <= previous_layer` marks a new model forward. This changes victim selection only;
986    /// every hit and miss still feeds identical expert bytes to the same GPU kernel.
987    pub(crate) fn begin_forward_epoch(&mut self, layer: u16, t: usize) {
988        let Some(decay) = self.frequency_decay else {
989            self.last_forward_layer = Some(layer);
990            self.last_forward_t = t;
991            return;
992        };
993        let new_sweep = self
994            .last_forward_layer
995            .is_some_and(|previous| layer <= previous);
996        if new_sweep && t == 1 {
997            if self.last_forward_t != 1 {
998                self.frequencies.clear();
999            } else {
1000                self.frequencies.retain(|_, score| {
1001                    *score *= decay;
1002                    *score >= 1.0e-3
1003                });
1004            }
1005        }
1006        self.last_forward_layer = Some(layer);
1007        self.last_forward_t = t;
1008    }
1009
1010    /// Turn already-submitted disk reads into copy-stream GPU admissions at a host-routing
1011    /// boundary. The caller must invoke this only after the router's DtoH synchronization has
1012    /// completed all earlier-layer compute, and `keep` must contain every block selected in the
1013    /// current layer. A reserved victim is therefore neither in use nor about to be used, so its
1014    /// H2D can start immediately while the CPU workers finish later reads. Consumers still insert
1015    /// an explicit compute-stream wait through the ordinary `pending` dispatch path.
1016    pub(crate) fn promote_worker_reads_at_safe_boundary(
1017        &mut self,
1018        order: &[BlockId],
1019        keep: &[BlockId],
1020        e: &Engine,
1021    ) -> Result<usize, Box<dyn std::error::Error>> {
1022        if !crate::spill_pread::copy_h2d_enabled() {
1023            return Ok(0);
1024        }
1025        let mut promoted = 0usize;
1026        for &id in order {
1027            if self.table.contains_key(&id) || self.pending.contains_key(&id) {
1028                if let Some(read) = self.worker_reads.remove(&id)
1029                    && let Some(pool) = self.pread.as_mut()
1030                {
1031                    let _ = pool.cancel_worker(read.ticket);
1032                }
1033                continue;
1034            }
1035            let Some(read) = self.worker_reads.get(&id).copied() else {
1036                continue;
1037            };
1038            let Some(slot) = self.reserve_prefetch_slot(read.len, keep) else {
1039                continue;
1040            };
1041            self.worker_reads.remove(&id);
1042
1043            let index = match self.pread.as_mut().unwrap().wait_worker(read.ticket) {
1044                Ok(index) => index,
1045                Err(err) => {
1046                    let _ = self.pread.as_mut().unwrap().cancel_worker(read.ticket);
1047                    self.release_reserved_slot(slot);
1048                    self.note_pread_fallback(err.as_ref());
1049                    continue;
1050                }
1051            };
1052            let ready = {
1053                let bytes = match self.pread.as_ref().unwrap().bytes(index, read.len) {
1054                    Ok(bytes) => bytes,
1055                    Err(err) => {
1056                        self.pread.as_mut().unwrap().abort_read(index);
1057                        self.release_reserved_slot(slot);
1058                        self.note_pread_fallback(err.as_ref());
1059                        continue;
1060                    }
1061                };
1062                stage_pread_prefetch_on_copy_stream(e, bytes, &mut self.slots[slot])
1063            };
1064            let ready = match ready {
1065                Ok(ready) => ready,
1066                Err((err, reusable)) => {
1067                    if reusable {
1068                        self.pread.as_mut().unwrap().abort_read(index);
1069                        self.release_reserved_slot(slot);
1070                        self.note_pread_fallback(err.as_ref());
1071                        continue;
1072                    }
1073                    self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1074                    self.copy_stream_unknown = true;
1075                    return Err(err);
1076                }
1077            };
1078            self.pread.as_mut().unwrap().mark_h2d(index, ready.clone());
1079            self.occupant[slot] = Some(id);
1080            self.pending.insert(
1081                id,
1082                PendingBlock {
1083                    slot,
1084                    ready,
1085                    keepalive: None,
1086                },
1087            );
1088            self.staged_bytes += read.len as u64;
1089            promoted += 1;
1090        }
1091        Ok(promoted)
1092    }
1093
1094    fn dispatch_disk(
1095        &mut self,
1096        id: BlockId,
1097        file: &Arc<std::fs::File>,
1098        offset: u64,
1099        len: usize,
1100        fallback: &[u8],
1101        e: &Engine,
1102    ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1103        if self.pread.is_none() {
1104            if self.pread_requested {
1105                self.note_pread_fallback(&"pinned-buffer backend unavailable");
1106            }
1107            return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1108        }
1109
1110        let pending = self.worker_reads.remove(&id);
1111        let pool = self.pread.as_mut().unwrap();
1112        let read = if pool.is_worker() {
1113            let ticket = match pending {
1114                Some(read) => Ok(Some(read.ticket)),
1115                None => pool.submit_worker(file.clone(), offset, len),
1116            };
1117            match ticket {
1118                Ok(Some(ticket)) => match pool.wait_worker(ticket) {
1119                    Ok(index) => Ok(index),
1120                    Err(err) => {
1121                        // Read errors normally release in wait_worker. A worker/channel failure may
1122                        // return earlier; cancel defensively so the next scope cannot lose the slot.
1123                        let _ = pool.cancel_worker(ticket);
1124                        Err(err)
1125                    }
1126                },
1127                Ok(None) => Err(std::io::Error::other("worker read ring is busy").into()),
1128                Err(err) => Err(err),
1129            }
1130        } else {
1131            debug_assert!(pending.is_none());
1132            pool.read(file.as_ref(), offset, len)
1133        };
1134        let index = match read {
1135            Ok(index) => index,
1136            Err(err) => {
1137                self.note_pread_fallback(err.as_ref());
1138                return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1139            }
1140        };
1141
1142        // The blocking read happens before eviction, so an I/O failure leaves cache residency
1143        // untouched and can safely use the mmap oracle.
1144        let slot = self.reserve_slot(len).ok_or_else(|| {
1145            std::io::Error::other(format!(
1146                "no MoE cache slot can hold {len} bytes (max class {})",
1147                self.classes.last().map(|class| class.capacity).unwrap_or(0)
1148            ))
1149        })?;
1150        let ready = {
1151            let bytes = match self.pread.as_ref().unwrap().bytes(index, len) {
1152                Ok(bytes) => bytes,
1153                Err(err) => {
1154                    self.pread.as_mut().unwrap().abort_read(index);
1155                    self.release_reserved_slot(slot);
1156                    self.note_pread_fallback(err.as_ref());
1157                    return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1158                }
1159            };
1160            stage_pread_on_compute_stream(e, bytes, &mut self.slots[slot])
1161        };
1162        let ready = match ready {
1163            Ok(ready) => ready,
1164            Err(err) => {
1165                // A memcpy or event-record failure can occur after submission. Synchronize the
1166                // retained compute stream before either source or destination is reused. If CUDA
1167                // cannot prove completion, quarantine both and fail instead of risking UAF.
1168                match e.stream().synchronize() {
1169                    Ok(()) => {
1170                        self.pread.as_mut().unwrap().abort_read(index);
1171                        self.release_reserved_slot(slot);
1172                        self.note_pread_fallback(err.as_ref());
1173                        return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1174                    }
1175                    Err(sync_err) => {
1176                        self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1177                        return Err(std::io::Error::other(format!(
1178                            "pread H2D setup failed ({err}); CUDA stream drain also failed ({sync_err})"
1179                        )).into());
1180                    }
1181                }
1182            }
1183        };
1184        self.pread.as_mut().unwrap().mark_h2d(index, ready);
1185        // Copy and dependent GEMM share the compute stream, so stream order is the consumer fence.
1186        // Publish only after both memcpy submission and explicit completion-event recording.
1187        self.staged_bytes += len as u64;
1188        self.publish(id, slot);
1189        Ok(DispatchSlot::Resident(slot))
1190    }
1191
1192    /// The dispatch decision for one (BlockId, host_bytes). Returns where the block landed; resolve
1193    /// the device buffer with `buf()`. On the bit-identity-critical path the buffer holds EXACTLY
1194    /// `host_bytes` either way (a HIT skipped the copy; the prior stage wrote the same bytes).
1195    ///
1196    /// Policy (MOE-SLRU-PLAN §B.2, first-miss admit since 2026-07-06):
1197    /// - HIT  (table[id] = s): promote, return s. ZERO PCIe.
1198    /// - MISS: admit (stage into a retained slot, evicting an SLRU victim when full).
1199    pub fn dispatch(
1200        &mut self,
1201        id: BlockId,
1202        host_bytes: &[u8],
1203        e: &Engine,
1204    ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1205        self.dispatch_source(
1206            id,
1207            ExpertSource::Memory {
1208                bytes: host_bytes,
1209                keepalive: None,
1210            },
1211            e,
1212        )
1213    }
1214
1215    pub(crate) fn dispatch_source(
1216        &mut self,
1217        id: BlockId,
1218        source: ExpertSource<'_>,
1219        e: &Engine,
1220    ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1221        self.reap_copy_sources();
1222        let increment = self.frequency_increment(id);
1223        *self.frequencies.entry(id).or_insert(0.0) += increment;
1224        if let Some(s) = self.table.get(&id).copied() {
1225            self.hits += 1;
1226            self.on_hit(s);
1227            return Ok(DispatchSlot::Resident(s));
1228        }
1229        if let Some(pending) = self.pending.remove(&id) {
1230            if let Err(err) = e.compute_wait(pending.ready.as_ref()) {
1231                self.pending.insert(id, pending);
1232                return Err(err);
1233            }
1234            self.misses += 1;
1235            let slot = pending.slot;
1236            if let Some(keepalive) = pending.keepalive {
1237                self.inflight_sources.push((pending.ready, keepalive));
1238            }
1239            self.publish(id, slot);
1240            return Ok(DispatchSlot::Resident(slot));
1241        }
1242        self.misses += 1;
1243        // FIRST-MISS ADMIT (the only policy since 2026-07-08; the second-miss "ghost" filter and
1244        // its seams MEMRA_MOE_GHOST / MEMRA_MOE_FAST_ADMIT are gone). Measured record: while FREE
1245        // slots remain, admission evicts nothing — filtering only delayed residency (96GB: 83.7%
1246        // steady hit-rate instead of ~100%, 74 MB/token avoidable PCIe; 2026-07-04). In the SPILL
1247        // regime (cache permanently full, local 35B) the filter made every cold block pay TWO H2D
1248        // copies — ~6% of token PCIe, measured ABOVE its eviction-protection benefit (24.2 -> 25.0
1249        // tok/s with it off, 2026-07-06). First-miss admit evicts an SLRU victim when full; the
1250        // SLRU probation segment still protects the protected set. Bit-identity unchanged: the
1251        // slot holds byte-for-byte the same GGUF block (D.2 gate).
1252        match source {
1253            ExpertSource::Memory { bytes, keepalive } => {
1254                // Retain before H2D submission so even the setup-error path cannot release a
1255                // pinned/mapped source while CUDA may still be reading it.
1256                self.retain_compute_source(keepalive);
1257                let slot = self.admit(id, bytes, e)?;
1258                Ok(DispatchSlot::Resident(slot))
1259            }
1260            ExpertSource::Disk {
1261                file,
1262                offset,
1263                len,
1264                fallback,
1265                keepalive,
1266            } => {
1267                // The owner is only needed when dispatch_disk falls back to mmap, but retaining
1268                // the usually shared mmap Arc once keeps every fallback branch simple and safe.
1269                self.retain_compute_source(Some(keepalive));
1270                self.dispatch_disk(id, file, offset, len, fallback, e)
1271            }
1272        }
1273    }
1274
1275    /// Deterministically stage a known-future block on the copy stream. The slot is reserved but is
1276    /// not considered resident until `dispatch` inserts a compute-stream wait for the returned copy
1277    /// event. Before overwriting a reused slot, the copy stream waits for all compute work already
1278    /// queued at this call site; the caller issues prefetch before the current expert's kernels, so
1279    /// the transfer can overlap those kernels without racing any earlier consumer of the victim.
1280    ///
1281    /// `keep` is the current expert's gate/up/down ids. If no safe victim exists, return `false` and
1282    /// let the normal synchronous miss path handle the block.
1283    pub fn prefetch(
1284        &mut self,
1285        id: BlockId,
1286        host_bytes: &[u8],
1287        keep: &[BlockId],
1288        e: &Engine,
1289    ) -> Result<bool, Box<dyn std::error::Error>> {
1290        self.prefetch_source(
1291            id,
1292            ExpertSource::Memory {
1293                bytes: host_bytes,
1294                keepalive: None,
1295            },
1296            keep,
1297            e,
1298        )
1299    }
1300
1301    fn reserve_prefetch_slot(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
1302        for class in &mut self.classes {
1303            if class.capacity >= required
1304                && let Some(slot) = class.free.pop()
1305            {
1306                return Some(slot);
1307            }
1308        }
1309        self.evict_one_excluding(required, keep)
1310    }
1311
1312    fn prefetch_bytes(
1313        &mut self,
1314        id: BlockId,
1315        host_bytes: &[u8],
1316        keepalive: Option<ExpertKeepalive>,
1317        keep: &[BlockId],
1318        e: &Engine,
1319    ) -> Result<bool, Box<dyn std::error::Error>> {
1320        let Some(slot) = self.reserve_prefetch_slot(host_bytes.len(), keep) else {
1321            return Ok(false);
1322        };
1323        let ready = match stage_on_copy_stream(e, host_bytes, &mut self.slots[slot]) {
1324            Ok(ready) => ready,
1325            Err((err, reusable)) => {
1326                if reusable {
1327                    self.release_reserved_slot(slot);
1328                } else {
1329                    // The slot is absent from free/table/SLRU and cannot be reused. Drop retries a
1330                    // whole copy-stream drain and leaks all slots if CUDA still cannot prove safety.
1331                    self.copy_stream_unknown = true;
1332                    if let Some(keepalive) = keepalive {
1333                        self.quarantined_sources.push(keepalive);
1334                    }
1335                    eprintln!(
1336                        "[moe-cache] quarantining slot {slot} after unprovable copy completion"
1337                    );
1338                }
1339                return Err(err);
1340            }
1341        };
1342        self.occupant[slot] = Some(id);
1343        self.pending.insert(
1344            id,
1345            PendingBlock {
1346                slot,
1347                ready,
1348                keepalive,
1349            },
1350        );
1351        self.staged_bytes += host_bytes.len() as u64;
1352        Ok(true)
1353    }
1354
1355    pub(crate) fn prefetch_source(
1356        &mut self,
1357        id: BlockId,
1358        source: ExpertSource<'_>,
1359        keep: &[BlockId],
1360        e: &Engine,
1361    ) -> Result<bool, Box<dyn std::error::Error>> {
1362        self.reap_copy_sources();
1363        if self.table.contains_key(&id)
1364            || self.pending.contains_key(&id)
1365            || self.worker_reads.contains_key(&id)
1366        {
1367            return Ok(false);
1368        }
1369        match source {
1370            ExpertSource::Memory { bytes, keepalive } => {
1371                self.prefetch_bytes(id, bytes, keepalive, keep, e)
1372            }
1373            ExpertSource::Disk {
1374                file,
1375                offset,
1376                len,
1377                fallback,
1378                keepalive,
1379            } => {
1380                if self.pread.as_ref().is_some_and(PreadPool::is_worker) {
1381                    match self.pread.as_mut().unwrap().submit_worker_speculative(
1382                        file.clone(),
1383                        offset,
1384                        len,
1385                    ) {
1386                        Ok(Some(ticket)) => {
1387                            self.worker_reads.insert(id, WorkerRead { ticket, len });
1388                            Ok(true)
1389                        }
1390                        Ok(None) => Ok(false),
1391                        Err(err) => {
1392                            self.note_pread_fallback(err.as_ref());
1393                            Ok(false)
1394                        }
1395                    }
1396                } else if self.pread.is_some() {
1397                    // Blocking `pread` remains demand-only so it cannot delay current compute.
1398                    Ok(false)
1399                } else {
1400                    self.prefetch_bytes(id, fallback, Some(keepalive), keep, e)
1401                }
1402            }
1403        }
1404    }
1405
1406    /// Pre-warm: force-admit a block (used by the §D.2 bit-identity gate to make all blocks resident).
1407    pub fn force_admit(
1408        &mut self,
1409        id: BlockId,
1410        host_bytes: &[u8],
1411        e: &Engine,
1412    ) -> Result<usize, Box<dyn std::error::Error>> {
1413        if let Some(s) = self.table.get(&id).copied() {
1414            return Ok(s);
1415        }
1416        self.admit(id, host_bytes, e)
1417    }
1418
1419    /// STAGE 3 one-shot PREWARM: force-admit every block of `layer` while FREE slots can hold it
1420    /// (never evicts — a spill rig whose cache can't fit the layer just skips; organic residency
1421    /// still applies). Runs at most once per layer (success or not). The H2D copies are the SAME
1422    /// stage_expert bytes the miss path would issue — bit-identity unchanged; this only front-loads
1423    /// them so the device-dispatch fast path fires from token 0 instead of after the SLRU fill.
1424    /// Frozen residency as (layer, proj, ex) triples in slot order, for the freeze-profile
1425    /// sidecar. Slot order keeps the restage admit sequence close to the original placement.
1426    pub fn export_residency(&self) -> Vec<(u16, u8, u16)> {
1427        self.occupant
1428            .iter()
1429            .flatten()
1430            .map(|id| (id.layer, id.proj, id.ex))
1431            .collect()
1432    }
1433
1434    /// Admit one specific block from a saved freeze profile, reading through the layer's
1435    /// established expert source (the same recipe as `prewarm_layer`, but id-targeted so a
1436    /// persisted residency set restages without a profiling warmup). Returns false for ids
1437    /// that no longer resolve (changed plan, pruned expert) — the caller counts and reports.
1438    pub fn restage_block(
1439        &mut self,
1440        id: BlockId,
1441        m: &crate::hybrid::MoeWeights,
1442        e: &Engine,
1443    ) -> Result<bool, Box<dyn std::error::Error>> {
1444        if self.table.contains_key(&id) {
1445            return Ok(true);
1446        }
1447        let exps = match id.proj {
1448            PROJ_GATE => &m.gate_exps,
1449            PROJ_UP => &m.up_exps,
1450            PROJ_DOWN => &m.down_exps,
1451            _ => return Ok(false),
1452        };
1453        if id.ex as usize >= exps.n_expert {
1454            return Ok(false);
1455        }
1456        if m.active_experts
1457            .as_ref()
1458            .is_some_and(|active| !active[id.ex as usize])
1459        {
1460            return Ok(false);
1461        }
1462        if exps.expert_layout(id.ex as usize).len == 0 {
1463            return Ok(false);
1464        }
1465        match exps.expert_source(id.ex as usize) {
1466            ExpertSource::Memory { bytes, keepalive } => {
1467                self.retain_compute_source(keepalive);
1468                self.admit(id, bytes, e)?;
1469            }
1470            ExpertSource::Disk {
1471                fallback,
1472                keepalive,
1473                ..
1474            } => {
1475                self.retain_compute_source(Some(keepalive));
1476                self.admit(id, fallback, e)?;
1477            }
1478        }
1479        Ok(true)
1480    }
1481
1482    pub fn prewarm_layer(
1483        &mut self,
1484        layer: u16,
1485        m: &crate::hybrid::MoeWeights,
1486        e: &Engine,
1487    ) -> Result<(), Box<dyn std::error::Error>> {
1488        if !self.prewarm_tried.insert(layer) {
1489            return Ok(());
1490        }
1491        let n_expert = m.gate_exps.n_expert;
1492        if self.pread.is_some()
1493            && (0..n_expert).any(|ex| {
1494                matches!(m.gate_exps.expert_source(ex), ExpertSource::Disk { .. })
1495                    || matches!(m.up_exps.expert_source(ex), ExpertSource::Disk { .. })
1496                    || matches!(m.down_exps.expert_source(ex), ExpertSource::Disk { .. })
1497            })
1498        {
1499            // Prewarm is a whole-layer scan. It must not silently turn explicit demand I/O back
1500            // into an mmap walk; organic misses will populate the cache through dispatch_source.
1501            return Ok(());
1502        }
1503        let resident = self.per_layer.get(&layer).copied().unwrap_or(0) as usize;
1504        let missing = 3 * n_expert - resident;
1505        if self.size_aware {
1506            return Ok(());
1507        } // heterogeneous prewarm needs a per-class fit proof
1508        if self
1509            .classes
1510            .iter()
1511            .map(|class| class.free.len())
1512            .sum::<usize>()
1513            < missing
1514        {
1515            return Ok(()); // won't evict for a prewarm
1516        }
1517        for ex in 0..n_expert {
1518            for (proj, exps) in [
1519                (PROJ_GATE, &m.gate_exps),
1520                (PROJ_UP, &m.up_exps),
1521                (PROJ_DOWN, &m.down_exps),
1522            ] {
1523                let id = BlockId::new(layer, proj, ex as u16);
1524                if self.table.contains_key(&id) {
1525                    continue;
1526                }
1527                match exps.expert_source(ex) {
1528                    ExpertSource::Memory { bytes, keepalive } => {
1529                        self.retain_compute_source(keepalive);
1530                        self.admit(id, bytes, e)?;
1531                    }
1532                    ExpertSource::Disk {
1533                        fallback,
1534                        keepalive,
1535                        ..
1536                    } => {
1537                        self.retain_compute_source(Some(keepalive));
1538                        self.admit(id, fallback, e)?;
1539                    }
1540                }
1541            }
1542        }
1543        Ok(())
1544    }
1545
1546    /// STAGE 3: device pointer row for a FULLY-RESIDENT layer. Returns the [3, n_expert] u64 slot
1547    /// base-address table (proj-major: gate row, up row, down row) if EVERY block of `layer` is
1548    /// cache-resident, else None (caller falls back to host routing). The row is built+uploaded on
1549    /// first full residency and reused until an eviction touches the layer. `n_expert` is the
1550    /// layer's expert count (the full-residency threshold is 3*n_expert blocks).
1551    pub fn layer_dev_row(
1552        &mut self,
1553        layer: u16,
1554        n_expert: usize,
1555        e: &Engine,
1556    ) -> Result<Option<&CudaSlice<u64>>, Box<dyn std::error::Error>> {
1557        if self.per_layer.get(&layer).copied().unwrap_or(0) as usize != 3 * n_expert {
1558            return Ok(None);
1559        }
1560        if !self.dev_rows.contains_key(&layer) {
1561            use cudarc::driver::DevicePtr;
1562            let mut host = vec![0u64; 3 * n_expert];
1563            for proj in 0..3u8 {
1564                for ex in 0..n_expert {
1565                    let Some(&s) = self.table.get(&BlockId::new(layer, proj, ex as u16)) else {
1566                        // count said fully resident but a block is missing — inconsistent; bail safe.
1567                        return Ok(None);
1568                    };
1569                    let __s_ev = e.stream();
1570                    let (p, _ev) = self.slots[s].device_ptr(&__s_ev);
1571                    host[proj as usize * n_expert + ex] = p;
1572                }
1573            }
1574            let row = e.stream().clone_htod(&host)?;
1575            self.dev_rows.insert(layer, row);
1576        }
1577        Ok(self.dev_rows.get(&layer))
1578    }
1579
1580    /// Resolve a `DispatchSlot` to the device buffer to feed `qmatvec_view`.
1581    #[inline]
1582    pub fn buf(&self, d: DispatchSlot) -> &CudaSlice<u8> {
1583        match d {
1584            DispatchSlot::Resident(s) => &self.slots[s],
1585        }
1586    }
1587
1588    /// Read-only access to a slot's device buffer (the `qmatvec_view` source on a HIT).
1589    #[inline]
1590    pub fn slot(&self, s: usize) -> &CudaSlice<u8> {
1591        &self.slots[s]
1592    }
1593
1594    /// Hit rate over this cache's lifetime (for the §D.4 print).
1595    pub fn hit_rate(&self) -> f64 {
1596        let tot = self.hits + self.misses;
1597        if tot == 0 {
1598            0.0
1599        } else {
1600            self.hits as f64 / tot as f64
1601        }
1602    }
1603
1604    /// Reset the per-window perf counters (lets the run print steady-state vs warmup separately).
1605    pub fn reset_counters(&mut self) {
1606        self.hits = 0;
1607        self.misses = 0;
1608        self.staged_bytes = 0;
1609    }
1610
1611    pub(crate) fn pread_stats(&self) -> Option<PreadStats> {
1612        if !self.pread_requested {
1613            return None;
1614        }
1615        let mut stats = self
1616            .pread
1617            .as_ref()
1618            .map(PreadPool::stats)
1619            .unwrap_or_default();
1620        stats.fallbacks = self.pread_fallbacks;
1621        Some(stats)
1622    }
1623}
1624
1625fn cache_lfu_decay() -> Option<f32> {
1626    let raw = std::env::var("MEMRA_MOE_LFU_DECAY").ok()?;
1627    match parse_cache_lfu_decay(Some(&raw)) {
1628        Ok(value) => value,
1629        Err(reason) => {
1630            eprintln!(
1631                "[moe-cache] invalid MEMRA_MOE_LFU_DECAY={raw:?} ({reason}); disabling LFU decay"
1632            );
1633            None
1634        }
1635    }
1636}
1637
1638fn cache_lfu_mtp_weight() -> f32 {
1639    const DEFAULT: f32 = 1.0;
1640    let raw = std::env::var("MEMRA_MOE_LFU_MTP_WEIGHT").ok();
1641    match parse_cache_lfu_mtp_weight(raw.as_deref()) {
1642        Ok(value) => value,
1643        Err(reason) => {
1644            eprintln!(
1645                "[moe-cache] invalid MEMRA_MOE_LFU_MTP_WEIGHT={:?} ({reason}); using {DEFAULT}",
1646                raw.as_deref().unwrap_or("")
1647            );
1648            DEFAULT
1649        }
1650    }
1651}
1652
1653fn parse_cache_lfu_mtp_weight(raw: Option<&str>) -> Result<f32, &'static str> {
1654    let value = raw
1655        .unwrap_or("1")
1656        .parse::<f32>()
1657        .map_err(|_| "expected a number")?;
1658    if value.is_finite() && (0.25..=64.0).contains(&value) {
1659        Ok(value)
1660    } else {
1661        Err("expected a finite multiplier from 0.25 through 64")
1662    }
1663}
1664
1665fn parse_cache_lfu_decay(raw: Option<&str>) -> Result<Option<f32>, &'static str> {
1666    let Some(raw) = raw else { return Ok(None) };
1667    let value = raw.parse::<f32>().map_err(|_| "expected a number")?;
1668    if value.is_finite() && value > 0.0 && value <= 1.0 {
1669        Ok(Some(value))
1670    } else {
1671        Err("expected a finite fraction greater than 0 and at most 1")
1672    }
1673}
1674
1675fn cache_hard_vram_frac() -> f64 {
1676    const DEFAULT: f64 = 0.80;
1677    let raw = std::env::var("MEMRA_MOE_HARD_VRAM_FRAC").ok();
1678    match parse_cache_hard_vram_frac(raw.as_deref()) {
1679        Ok(value) => value,
1680        Err(reason) => {
1681            eprintln!(
1682                "[moe-cache] invalid MEMRA_MOE_HARD_VRAM_FRAC={:?} ({reason}); using {DEFAULT}",
1683                raw.as_deref().unwrap_or("")
1684            );
1685            DEFAULT
1686        }
1687    }
1688}
1689
1690fn parse_cache_hard_vram_frac(raw: Option<&str>) -> Result<f64, &'static str> {
1691    let value = raw
1692        .unwrap_or("0.80")
1693        .parse::<f64>()
1694        .map_err(|_| "expected a number")?;
1695    if value.is_finite() && (0.10..=0.95).contains(&value) {
1696        Ok(value)
1697    } else {
1698        Err("expected a finite fraction from 0.10 through 0.95")
1699    }
1700}
1701
1702#[cfg(test)]
1703mod slru_intrusive_tests {
1704    //! Q5 policy-equivalence proof (research/audit-fixes2-20260805): the intrusive-list SLRU
1705    //! must make the SAME eviction decisions for the same access pattern as the pre-fix
1706    //! VecDeque SLRU. `OldSlru` below is the pre-fix implementation transcribed verbatim
1707    //! (position()+remove() promotion, probation-then-protected pop_front eviction,
1708    //! protected_cap demotion loop); both are driven with identical randomized op sequences
1709    //! and their full segment orders compared after EVERY op.
1710    use super::{SEG_PROBATION, SEG_PROTECTED, SlotClass, SlotLink, SlruList};
1711    use std::collections::VecDeque;
1712
1713    /// The pre-fix policy, verbatim (moe_cache.rs @ 61953206 lines 540-571).
1714    struct OldSlru {
1715        probation: VecDeque<usize>,
1716        protected: VecDeque<usize>,
1717        protected_cap: usize,
1718    }
1719    impl OldSlru {
1720        fn on_hit_full(&mut self, slot: usize) {
1721            if let Some(pos) = self.probation.iter().position(|&x| x == slot) {
1722                self.probation.remove(pos);
1723                self.push_protected(slot);
1724            } else if let Some(pos) = self.protected.iter().position(|&x| x == slot) {
1725                self.protected.remove(pos);
1726                self.protected.push_back(slot); // MRU
1727            } else {
1728                self.push_protected(slot);
1729            }
1730        }
1731        fn push_protected(&mut self, slot: usize) {
1732            self.protected.push_back(slot);
1733            while self.protected.len() > self.protected_cap {
1734                if let Some(demoted) = self.protected.pop_front() {
1735                    self.probation.push_back(demoted);
1736                } else {
1737                    break;
1738                }
1739            }
1740        }
1741        fn pop_lru(&mut self) -> Option<usize> {
1742            self.probation
1743                .pop_front()
1744                .or_else(|| self.protected.pop_front())
1745        }
1746        fn take_excluding(&mut self, banned: &[usize]) -> Option<usize> {
1747            let take = |q: &mut VecDeque<usize>| {
1748                q.iter()
1749                    .position(|&s| !banned.contains(&s))
1750                    .and_then(|pos| q.remove(pos))
1751            };
1752            take(&mut self.probation).or_else(|| take(&mut self.protected))
1753        }
1754    }
1755
1756    fn new_pair(n: usize, protected_cap: usize) -> (SlotClass, Vec<SlotLink>, OldSlru) {
1757        let class = SlotClass {
1758            capacity: 1,
1759            probation: SlruList::new(),
1760            protected: SlruList::new(),
1761            free: Vec::new(),
1762            protected_cap,
1763        };
1764        let links = vec![SlotLink::none(); n];
1765        let old = OldSlru {
1766            probation: VecDeque::new(),
1767            protected: VecDeque::new(),
1768            protected_cap,
1769        };
1770        (class, links, old)
1771    }
1772
1773    fn orders_match(class: &SlotClass, links: &[SlotLink], old: &OldSlru) -> bool {
1774        let np: Vec<usize> = class.probation.iter(links).collect();
1775        let nt: Vec<usize> = class.protected.iter(links).collect();
1776        let op: Vec<usize> = old.probation.iter().copied().collect();
1777        let ot: Vec<usize> = old.protected.iter().copied().collect();
1778        np == op && nt == ot
1779    }
1780
1781    /// Deterministic PRNG (SplitMix64) — no dev-dependencies.
1782    struct Rng(u64);
1783    impl Rng {
1784        fn next(&mut self) -> u64 {
1785            self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
1786            let mut z = self.0;
1787            z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1788            z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1789            z ^ (z >> 31)
1790        }
1791        fn below(&mut self, n: usize) -> usize {
1792            (self.next() % n as u64) as usize
1793        }
1794    }
1795
1796    #[test]
1797    fn same_eviction_decisions_randomized_soak() {
1798        // Sweep several (n_slots, protected_cap) shapes including cap=1 edge and the 0.8 default.
1799        for &(n, cap) in &[(8usize, 1usize), (16, 12), (64, 51), (128, 102)] {
1800            let (mut class, mut links, mut old) = new_pair(n, cap);
1801            let mut rng = Rng(0xC0FFEE ^ (n as u64) << 8 ^ cap as u64);
1802            let mut resident: Vec<usize> = Vec::new();
1803            let mut free: Vec<usize> = (0..n).rev().collect();
1804            for step in 0..200_000 {
1805                let op = rng.below(100);
1806                if op < 55 && !resident.is_empty() {
1807                    // HIT on a random resident slot (full-class path — free handled below).
1808                    let slot = resident[rng.below(resident.len())];
1809                    class.on_hit_full(slot, &mut links);
1810                    old.on_hit_full(slot);
1811                } else if op < 80 {
1812                    // ADMIT: free slot first (publish -> probation MRU), else evict LRU + reuse.
1813                    let slot = if let Some(s) = free.pop() {
1814                        s
1815                    } else {
1816                        let v_new = class.pop_lru(&mut links);
1817                        let v_old = old.pop_lru();
1818                        assert_eq!(
1819                            v_new, v_old,
1820                            "victim diverged at step {step} (n={n} cap={cap})"
1821                        );
1822                        let v = v_new.unwrap();
1823                        resident.retain(|&s| s != v);
1824                        v
1825                    };
1826                    class.probation.push_back(slot, SEG_PROBATION, &mut links);
1827                    old.probation.push_back(slot);
1828                    resident.push(slot);
1829                } else if op < 92 && resident.len() > 2 {
1830                    // PREFETCH eviction: skip up to 3 "keep" slots (evict_one_excluding shape).
1831                    let banned: Vec<usize> = (0..3.min(resident.len()))
1832                        .map(|_| resident[rng.below(resident.len())])
1833                        .collect();
1834                    let take_new = {
1835                        let q = &mut class.probation;
1836                        let found = q.iter(&links).find(|s| !banned.contains(s));
1837                        match found {
1838                            Some(s) => {
1839                                q.unlink(s, &mut links);
1840                                Some(s)
1841                            }
1842                            None => {
1843                                let q = &mut class.protected;
1844                                q.iter(&links).find(|s| !banned.contains(s)).inspect(|&s| {
1845                                    q.unlink(s, &mut links);
1846                                })
1847                            }
1848                        }
1849                    };
1850                    let take_old = old.take_excluding(&banned);
1851                    assert_eq!(
1852                        take_new, take_old,
1853                        "excluding-victim diverged at step {step}"
1854                    );
1855                    if let Some(v) = take_new {
1856                        resident.retain(|&s| s != v);
1857                        free.push(v);
1858                    }
1859                } else if !resident.is_empty() {
1860                    // Defensive arm: hit on a slot in NEITHER segment (unlink first, then hit).
1861                    let slot = resident[rng.below(resident.len())];
1862                    match links[slot].seg {
1863                        SEG_PROBATION => class.probation.unlink(slot, &mut links),
1864                        SEG_PROTECTED => class.protected.unlink(slot, &mut links),
1865                        _ => {}
1866                    }
1867                    if let Some(pos) = old.probation.iter().position(|&x| x == slot) {
1868                        old.probation.remove(pos);
1869                    } else if let Some(pos) = old.protected.iter().position(|&x| x == slot) {
1870                        old.protected.remove(pos);
1871                    }
1872                    class.on_hit_full(slot, &mut links);
1873                    old.on_hit_full(slot);
1874                }
1875                assert!(
1876                    orders_match(&class, &links, &old),
1877                    "segment order diverged at step {step} (n={n} cap={cap})"
1878                );
1879            }
1880        }
1881    }
1882
1883    #[test]
1884    fn hit_promotion_is_o1_not_on() {
1885        // Op-count proof: time-per-hit must not grow with n_slots. 46k slots x 850 hits — the
1886        // audit's spill shape — as a wall-clock microbench: the old structure walked ~n/2 per
1887        // hit (~20M steps); the intrusive list does constant work. Assert the per-hit cost at
1888        // 46k slots stays within 8x of the 1k-slot cost (an O(n) scan would be ~46x+).
1889        fn bench(n: usize, hits: usize) -> std::time::Duration {
1890            let (mut class, mut links, _) = new_pair(n, (n as f64 * 0.8) as usize);
1891            for s in 0..n {
1892                class.probation.push_back(s, SEG_PROBATION, &mut links);
1893            }
1894            let mut rng = Rng(0xBEEF);
1895            let t0 = std::time::Instant::now();
1896            for _ in 0..hits {
1897                class.on_hit_full(rng.below(n), &mut links);
1898            }
1899            t0.elapsed()
1900        }
1901        // Warm both shapes once (alloc noise), then measure.
1902        bench(1_000, 10_000);
1903        bench(46_000, 10_000);
1904        let small = bench(1_000, 850_000).as_secs_f64() / 850_000.0;
1905        let large = bench(46_000, 850_000).as_secs_f64() / 850_000.0;
1906        assert!(
1907            large < small * 8.0,
1908            "per-hit cost scaled with n_slots: {:.1}ns @1k vs {:.1}ns @46k",
1909            small * 1e9,
1910            large * 1e9
1911        );
1912    }
1913
1914    #[test]
1915    fn slru_list_basic_invariants() {
1916        let mut links = vec![SlotLink::none(); 4];
1917        let mut l = SlruList::new();
1918        assert_eq!(l.pop_front(&mut links), None);
1919        l.push_back(2, SEG_PROBATION, &mut links);
1920        l.push_back(0, SEG_PROBATION, &mut links);
1921        l.push_back(3, SEG_PROBATION, &mut links);
1922        assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 0, 3]);
1923        assert_eq!(l.len, 3);
1924        l.unlink(0, &mut links); // middle
1925        assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 3]);
1926        l.unlink(3, &mut links); // tail
1927        assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2]);
1928        assert_eq!(l.pop_front(&mut links), Some(2)); // head
1929        assert_eq!(l.len, 0);
1930        assert_eq!(l.head, super::NIL);
1931        assert_eq!(l.tail, super::NIL);
1932        assert!(links.iter().all(|k| k.seg == super::SEG_NONE));
1933    }
1934}
1935
1936#[cfg(test)]
1937mod vram_fraction_tests {
1938    use super::{
1939        parse_cache_hard_vram_frac, parse_cache_lfu_decay, parse_cache_lfu_mtp_weight,
1940        size_class_plan,
1941    };
1942
1943    #[test]
1944    fn hard_vram_fraction_defaults_and_rejects_unsafe_values() {
1945        assert_eq!(parse_cache_hard_vram_frac(None), Ok(0.80));
1946        assert_eq!(parse_cache_hard_vram_frac(Some("0.82")), Ok(0.82));
1947        assert!(parse_cache_hard_vram_frac(Some("NaN")).is_err());
1948        assert_eq!(parse_cache_hard_vram_frac(Some("0.95")), Ok(0.95));
1949        assert!(parse_cache_hard_vram_frac(Some("0.96")).is_err());
1950        assert!(parse_cache_hard_vram_frac(Some("1.0")).is_err());
1951        assert!(parse_cache_hard_vram_frac(Some("bad")).is_err());
1952    }
1953
1954    #[test]
1955    fn lfu_decay_is_opt_in_and_bounded() {
1956        assert_eq!(parse_cache_lfu_decay(None), Ok(None));
1957        assert_eq!(parse_cache_lfu_decay(Some("0.8")), Ok(Some(0.8)));
1958        assert_eq!(parse_cache_lfu_decay(Some("1")), Ok(Some(1.0)));
1959        for value in ["0", "-0.1", "1.1", "NaN", "bad"] {
1960            assert!(
1961                parse_cache_lfu_decay(Some(value)).is_err(),
1962                "accepted {value}"
1963            );
1964        }
1965    }
1966
1967    #[test]
1968    fn lfu_mtp_weight_defaults_and_is_bounded() {
1969        assert_eq!(parse_cache_lfu_mtp_weight(None), Ok(1.0));
1970        assert_eq!(parse_cache_lfu_mtp_weight(Some("4")), Ok(4.0));
1971        for value in ["0", "0.1", "65", "NaN", "bad"] {
1972            assert!(
1973                parse_cache_lfu_mtp_weight(Some(value)).is_err(),
1974                "accepted {value}"
1975            );
1976        }
1977    }
1978
1979    #[test]
1980    fn size_class_plan_preserves_classes_and_never_exceeds_budget() {
1981        let blocks = [100usize, 100, 100, 200, 200, 400];
1982        let budget = (108 * 2) + 208 + 408;
1983        let plan = size_class_plan(&blocks, budget);
1984        assert!(plan.iter().all(|(_, count)| *count > 0));
1985        assert!(
1986            plan.iter()
1987                .map(|(bytes, count)| (bytes + 8) * count)
1988                .sum::<usize>()
1989                <= budget
1990        );
1991        assert!(plan.iter().all(|(bytes, count)| {
1992            *count <= blocks.iter().filter(|block| **block == *bytes).count()
1993        }));
1994    }
1995
1996    #[test]
1997    fn size_class_plan_returns_full_inventory_when_it_fits() {
1998        let blocks = [100usize, 100, 200, 400];
1999        let budget: usize = blocks.iter().map(|bytes| bytes + 8).sum();
2000        assert_eq!(
2001            size_class_plan(&blocks, budget),
2002            vec![(100, 2), (200, 1), (400, 1)]
2003        );
2004    }
2005
2006    #[test]
2007    fn size_class_plan_does_not_overflow_on_pathological_sizes() {
2008        let plan = size_class_plan(&[usize::MAX, usize::MAX], usize::MAX);
2009        assert!(plan.is_empty());
2010    }
2011}
2012
2013impl Drop for MoeSlotCache {
2014    fn drop(&mut self) {
2015        // Event tracking is intentionally disabled in Engine. Drain explicit copy-stream handoffs
2016        // before either the destination slots or pinned read buffers begin field destruction.
2017        let mut safe_to_drop_slots = true;
2018        if self.compute_stream_unknown || !self.compute_sources.is_empty() {
2019            if let Err(err) = self.compute_stream.synchronize() {
2020                safe_to_drop_slots = false;
2021                eprintln!(
2022                    "[moe-cache] unknown compute-stream drain failed ({err}); leaking GPU slots for safety"
2023                );
2024                for (_, keepalive) in self.compute_sources.drain() {
2025                    std::mem::forget(keepalive);
2026                }
2027            } else {
2028                self.compute_stream_unknown = false;
2029                self.compute_sources.clear();
2030            }
2031        }
2032        let need_copy_drain = self.copy_stream_unknown
2033            || !self.pending.is_empty()
2034            || !self.inflight_sources.is_empty()
2035            || !self.quarantined_sources.is_empty();
2036        if need_copy_drain {
2037            if let Err(err) = self.copy_stream.synchronize() {
2038                safe_to_drop_slots = false;
2039                eprintln!(
2040                    "[moe-cache] unknown copy-stream drain failed ({err}); leaking GPU slots for safety"
2041                );
2042                for (_, keepalive) in self.inflight_sources.drain(..) {
2043                    std::mem::forget(keepalive);
2044                }
2045                for keepalive in self.quarantined_sources.drain(..) {
2046                    std::mem::forget(keepalive);
2047                }
2048                for (_, pending) in self.pending.drain() {
2049                    if let Some(keepalive) = pending.keepalive {
2050                        std::mem::forget(keepalive);
2051                    }
2052                }
2053            } else {
2054                self.copy_stream_unknown = false;
2055                self.inflight_sources.clear();
2056                self.quarantined_sources.clear();
2057                self.pending.clear();
2058            }
2059        }
2060        if let Some(pool) = self.pread.as_mut() {
2061            safe_to_drop_slots &= pool.drain();
2062        } else if self.pread_requested && self.pread_fallbacks != 0 {
2063            eprintln!(
2064                "[spill-pread] backend unavailable; mmap_fallbacks={}",
2065                self.pread_fallbacks
2066            );
2067        }
2068        if !safe_to_drop_slots {
2069            for slot in self.slots.drain(..) {
2070                std::mem::forget(slot);
2071            }
2072        }
2073    }
2074}