1use std::collections::{BTreeMap, HashMap, HashSet};
18use std::sync::Arc;
19use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, HostSlice, SyncOnDrop};
20use crate::Engine;
21use crate::model::{ExpertKeepalive, ExpertSource};
22use crate::spill_pread::{PreadPool, PreadStats, ReadTicket, SpillIoMode};
23
24pub const PROJ_GATE: u8 = 0;
26pub const PROJ_UP: u8 = 1;
27pub const PROJ_DOWN: u8 = 2;
28
29#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
31pub struct BlockId {
32 pub layer: u16,
33 pub proj: u8,
34 pub ex: u16,
35}
36impl BlockId {
37 #[inline]
38 pub fn new(layer: u16, proj: u8, ex: u16) -> Self { BlockId { layer, proj, ex } }
39}
40
41#[derive(Clone, Copy, Debug)]
44pub enum DispatchSlot {
45 Resident(usize),
46}
47
48const NIL: u32 = u32::MAX;
50const SEG_NONE: u8 = 0;
51const SEG_PROBATION: u8 = 1;
52const SEG_PROTECTED: u8 = 2;
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
57struct SlotLink {
58 prev: u32,
59 next: u32,
60 seg: u8,
61}
62impl SlotLink {
63 const fn none() -> Self { SlotLink { prev: NIL, next: NIL, seg: SEG_NONE } }
64}
65
66#[derive(Debug)]
75struct SlruList {
76 head: u32,
77 tail: u32,
78 len: usize,
79}
80impl SlruList {
81 const fn new() -> Self { SlruList { head: NIL, tail: NIL, len: 0 } }
82
83 fn push_back(&mut self, slot: usize, seg: u8, links: &mut [SlotLink]) {
85 debug_assert_eq!(links[slot].seg, SEG_NONE, "slot {slot} already in a segment");
86 let s = slot as u32;
87 links[slot] = SlotLink { prev: self.tail, next: NIL, seg };
88 if self.tail != NIL {
89 links[self.tail as usize].next = s;
90 } else {
91 self.head = s;
92 }
93 self.tail = s;
94 self.len += 1;
95 }
96
97 fn pop_front(&mut self, links: &mut [SlotLink]) -> Option<usize> {
99 if self.head == NIL {
100 return None;
101 }
102 let s = self.head as usize;
103 self.unlink(s, links);
104 Some(s)
105 }
106
107 fn unlink(&mut self, slot: usize, links: &mut [SlotLink]) {
109 let l = links[slot];
110 debug_assert_ne!(l.seg, SEG_NONE, "unlink of slot {slot} not in a segment");
111 if l.prev != NIL {
112 links[l.prev as usize].next = l.next;
113 } else {
114 debug_assert_eq!(self.head, slot as u32);
115 self.head = l.next;
116 }
117 if l.next != NIL {
118 links[l.next as usize].prev = l.prev;
119 } else {
120 debug_assert_eq!(self.tail, slot as u32);
121 self.tail = l.prev;
122 }
123 links[slot] = SlotLink::none();
124 self.len -= 1;
125 }
126
127 fn iter<'a>(&self, links: &'a [SlotLink]) -> SlruIter<'a> {
129 SlruIter { links, cur: self.head }
130 }
131}
132
133struct SlruIter<'a> {
134 links: &'a [SlotLink],
135 cur: u32,
136}
137impl Iterator for SlruIter<'_> {
138 type Item = usize;
139 fn next(&mut self) -> Option<usize> {
140 if self.cur == NIL {
141 return None;
142 }
143 let s = self.cur as usize;
144 self.cur = self.links[s].next;
145 Some(s)
146 }
147}
148
149struct SlotClass {
152 capacity: usize,
153 probation: SlruList,
154 protected: SlruList,
155 free: Vec<usize>,
156 protected_cap: usize,
157}
158
159impl SlotClass {
160 fn on_hit_full(&mut self, slot: usize, links: &mut [SlotLink]) {
164 match links[slot].seg {
165 SEG_PROBATION => {
166 self.probation.unlink(slot, links);
167 self.push_protected(slot, links);
168 }
169 SEG_PROTECTED => {
170 self.protected.unlink(slot, links);
171 self.protected.push_back(slot, SEG_PROTECTED, links); }
173 _ => self.push_protected(slot, links),
175 }
176 }
177
178 fn push_protected(&mut self, slot: usize, links: &mut [SlotLink]) {
180 self.protected.push_back(slot, SEG_PROTECTED, links);
181 while self.protected.len > self.protected_cap {
182 if let Some(demoted) = self.protected.pop_front(links) {
183 self.probation.push_back(demoted, SEG_PROBATION, links);
184 } else {
185 break;
186 }
187 }
188 }
189
190 fn pop_lru(&mut self, links: &mut [SlotLink]) -> Option<usize> {
192 self.probation.pop_front(links).or_else(|| self.protected.pop_front(links))
193 }
194
195 fn unlink_from_segment(&mut self, slot: usize, links: &mut [SlotLink]) {
197 match links[slot].seg {
198 SEG_PROBATION => self.probation.unlink(slot, links),
199 SEG_PROTECTED => self.protected.unlink(slot, links),
200 _ => {}
201 }
202 }
203}
204
205pub struct MoeSlotCache {
208 slots: Vec<CudaSlice<u8>>, slot_class: Vec<usize>, classes: Vec<SlotClass>,
211 links: Vec<SlotLink>,
214 occupant: Vec<Option<BlockId>>, table: HashMap<BlockId, usize>, frequencies: HashMap<BlockId, f32>,
220 pending: HashMap<BlockId, PendingBlock>,
225 inflight_sources: Vec<(Arc<CudaEvent>, ExpertKeepalive)>,
228 quarantined_sources: Vec<ExpertKeepalive>,
231 compute_sources: HashMap<KeepaliveKey, ExpertKeepalive>,
235 pread: Option<PreadPool>,
238 worker_reads: HashMap<BlockId, WorkerRead>,
241 pread_requested: bool,
242 pread_fallbacks: u64,
243 copy_stream: Arc<CudaStream>,
246 copy_stream_unknown: bool,
247 compute_stream: Arc<CudaStream>,
248 compute_stream_unknown: bool,
249
250 n: usize,
251 max_block_bytes: usize,
252 size_aware: bool,
253 frequency_evict: bool,
254 frequency_decay: Option<f32>,
255 mtp_frequency_weight: f32,
260 last_forward_layer: Option<u16>,
261 last_forward_t: usize,
262 frozen: bool,
266
267 per_layer: HashMap<u16, u32>,
274 dev_rows: HashMap<u16, CudaSlice<u64>>,
279 prewarm_tried: HashSet<u16>,
282
283 pub hits: u64,
285 pub misses: u64,
286 pub staged_bytes: u64, }
288
289struct PendingBlock {
290 slot: usize,
291 ready: Arc<CudaEvent>,
292 keepalive: Option<ExpertKeepalive>,
293}
294
295#[derive(Clone, Copy)]
296struct WorkerRead {
297 ticket: ReadTicket,
298 len: usize,
299}
300
301#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
302enum KeepaliveKey {
303 Pinned(usize),
304 Buffer(usize),
305 Mmap(usize),
306}
307
308impl KeepaliveKey {
309 fn from_owner(owner: &ExpertKeepalive) -> Self {
310 match owner {
311 ExpertKeepalive::Pinned(value) => Self::Pinned(Arc::as_ptr(value) as usize),
312 ExpertKeepalive::Buffer(value) => Self::Buffer(Arc::as_ptr(value) as usize),
313 ExpertKeepalive::Mmap(value) => Self::Mmap(Arc::as_ptr(value) as usize),
314 }
315 }
316}
317
318struct ExactPinnedPrefix<'a>(&'a [u8]);
323
324impl HostSlice<u8> for ExactPinnedPrefix<'_> {
325 fn len(&self) -> usize { self.0.len() }
326
327 unsafe fn stream_synced_slice<'a>(
328 &'a self,
329 _stream: &'a CudaStream,
330 ) -> (&'a [u8], SyncOnDrop<'a>) {
331 (self.0, SyncOnDrop::Record(None))
334 }
335
336
337
338 unsafe fn stream_synced_mut_slice<'a>(
339 &'a mut self,
340 _stream: &'a CudaStream,
341 ) -> (&'a mut [u8], SyncOnDrop<'a>) {
342 panic!("ExactPinnedPrefix is a source-only HostSlice")
343 }
344}
345
346fn stage_on_copy_stream(
347 e: &Engine,
348 host_bytes: &[u8],
349 slot: &mut CudaSlice<u8>,
350) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
351 let prior = match e.stream().record_event(None) {
353 Ok(prior) => prior,
354 Err(err) => return Err((err.into(), true)),
355 };
356 if let Err(err) = e.copy_stream.wait(&prior) {
357 return Err((err.into(), true));
358 }
359 match e.stage_expert_async(host_bytes, slot, 0) {
360 Ok(ready) => Ok(Arc::new(ready)),
361 Err(err) => {
362 match e.copy_stream.synchronize() {
365 Ok(()) => Err((err, true)),
366 Err(sync_err) => Err((std::io::Error::other(format!(
367 "copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
368 )).into(), false)),
369 }
370 }
371 }
372}
373
374fn stage_pread_on_compute_stream(
375 e: &Engine,
376 host_bytes: &[u8],
377 slot: &mut CudaSlice<u8>,
378) -> Result<Arc<CudaEvent>, Box<dyn std::error::Error>> {
379 let ready = Arc::new(e.ctx().new_event(None)?);
380 let source = ExactPinnedPrefix(host_bytes);
381 let mut dst = slot.slice_mut(0..host_bytes.len());
382 e.stream().memcpy_htod(&source, &mut dst)?;
383 ready.record(&e.stream())?;
384 Ok(ready)
385}
386
387fn stage_pread_prefetch_on_copy_stream(
388 e: &Engine,
389 host_bytes: &[u8],
390 slot: &mut CudaSlice<u8>,
391) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
392 let ready = match e.ctx().new_event(None) {
393 Ok(ready) => Arc::new(ready),
394 Err(err) => return Err((err.into(), true)),
395 };
396 let source = ExactPinnedPrefix(host_bytes);
397 let mut dst = slot.slice_mut(0..host_bytes.len());
398 let submitted = e
399 .copy_stream
400 .memcpy_htod(&source, &mut dst)
401 .and_then(|()| ready.record(&e.copy_stream));
402 match submitted {
403 Ok(()) => Ok(ready),
404 Err(err) => match e.copy_stream.synchronize() {
405 Ok(()) => Err((err.into(), true)),
406 Err(sync_err) => Err((
407 std::io::Error::other(format!(
408 "pread copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
409 ))
410 .into(),
411 false,
412 )),
413 },
414 }
415}
416
417fn size_class_plan(block_bytes: &[usize], budget_bytes: usize) -> Vec<(usize, usize)> {
420 let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
421 for &bytes in block_bytes.iter().filter(|&&bytes| bytes > 0) {
422 *counts.entry(bytes).or_insert(0) += 1;
423 }
424 if counts.is_empty() || budget_bytes == 0 {
425 return Vec::new();
426 }
427 let total_bytes: u128 = counts
428 .iter()
429 .map(|(&bytes, &count)| (bytes as u128 + 8) * count as u128)
430 .sum();
431 let budget = budget_bytes as u128;
432 let mut plan: Vec<(usize, usize, u128)> = counts
433 .iter()
434 .map(|(&bytes, &count)| {
435 let scaled = count as u128 * budget;
436 (
437 bytes,
438 (scaled / total_bytes).min(count as u128) as usize,
439 scaled % total_bytes,
440 )
441 })
442 .collect();
443 let mut used: u128 = plan
444 .iter()
445 .map(|(bytes, count, _)| (*bytes as u128 + 8) * *count as u128)
446 .sum();
447
448 let mut order: Vec<usize> = (0..plan.len()).collect();
451 order.sort_by(|&a, &b| plan[b].2.cmp(&plan[a].2).then(a.cmp(&b)));
452 for index in order {
453 let (bytes, count, _) = plan[index];
454 let available = counts[&bytes];
455 let required = bytes as u128 + 8;
456 if count < available && used + required <= budget {
457 plan[index].1 += 1;
458 used += required;
459 }
460 }
461 plan.into_iter()
462 .filter_map(|(bytes, count, _)| (count > 0).then_some((bytes, count)))
463 .collect()
464}
465
466impl MoeSlotCache {
467 pub fn new(e: &Engine, max_block_bytes: usize) -> Result<Self, Box<dyn std::error::Error>> {
474 let (free, _total) = e.ctx().mem_get_info()?;
475 let hard_frac = cache_hard_vram_frac();
478 let hard_bytes =
479 ((free as f64 * hard_frac) as usize).saturating_sub(2 * (max_block_bytes + 8));
480 let forced_slots = std::env::var("MEMRA_MOE_SLOTS")
481 .ok()
482 .and_then(|s| s.parse::<usize>().ok());
483 let requested_bytes = if let Some(n) = forced_slots {
484 n.saturating_mul(max_block_bytes + 8)
485 } else {
486 let frac = std::env::var("MEMRA_MOE_VRAM_FRAC") .ok()
492 .and_then(|s| s.parse::<f64>().ok())
493 .unwrap_or(0.85);
494 (free as f64 * frac) as usize
495 };
496 let budget_bytes = requested_bytes.min(hard_bytes);
497 let layout = e.moe_cache_layout().unwrap_or_default();
498 let size_aware = forced_slots.is_none()
499 && std::env::var("MEMRA_MOE_SIZE_AWARE").as_deref() == Ok("1")
500 && !layout.is_empty();
501 let frequency_evict = std::env::var("MEMRA_MOE_LFU").as_deref() == Ok("1");
502 let frequency_decay = if frequency_evict {
503 cache_lfu_decay()
504 } else {
505 None
506 };
507 let mtp_frequency_weight = cache_lfu_mtp_weight();
508 let mut class_plan = if size_aware {
509 size_class_plan(&layout, budget_bytes)
510 } else {
511 Vec::new()
512 };
513 if class_plan.iter().map(|(_, count)| count).sum::<usize>() < 8 {
514 let n = (budget_bytes / (max_block_bytes + 8)).max(8);
515 class_plan = vec![(max_block_bytes, n)];
516 }
517 let n: usize = class_plan.iter().map(|(_, count)| count).sum();
518
519 let mut slots = Vec::with_capacity(n);
520 let mut slot_class = Vec::with_capacity(n);
521 let mut classes = Vec::with_capacity(class_plan.len());
522 let mut occupant = Vec::with_capacity(n);
523 for (class_index, &(capacity, count)) in class_plan.iter().enumerate() {
524 let start = slots.len();
525 for _ in 0..count {
526 slots.push(e.alloc_u8(capacity + 8)?);
528 slot_class.push(class_index);
529 occupant.push(None);
530 }
531 let free_slots = (start..start + count).rev().collect();
532 classes.push(SlotClass {
533 capacity,
534 probation: SlruList::new(),
535 protected: SlruList::new(),
536 free: free_slots,
537 protected_cap: ((count as f64 * 0.8) as usize).max(1),
538 });
539 }
540 let links = vec![SlotLink::none(); n];
541 if size_aware {
542 let allocated: usize = class_plan
543 .iter()
544 .map(|(bytes, count)| (bytes + 8) * count)
545 .sum();
546 eprintln!(
547 "[moe-cache] size-aware fixed slots: {n} slots in {} classes, {:.2} GB / {:.2} GB budget",
548 class_plan.len(),
549 allocated as f64 / 1e9,
550 budget_bytes as f64 / 1e9
551 );
552 }
553 let pread_mode = crate::spill_pread::configured_mode();
554 let pread_requested = pread_mode != SpillIoMode::Mmap;
555 let pread = if pread_requested {
556
557 match PreadPool::try_new(e, max_block_bytes, pread_mode) {
558 Ok(pool) => Some(pool),
559 Err(err) => {
560 eprintln!("[spill-pread] pinned-buffer initialization failed ({err}); using mmap");
561 None
562 }
563 }
564 } else { None };
565
566
567 Ok(MoeSlotCache {
568 slots,
569 slot_class,
570 classes,
571 links,
572 occupant,
573 table: HashMap::with_capacity(n * 2),
574 frequencies: HashMap::with_capacity(layout.len().max(n * 2)),
575 pending: HashMap::new(),
576 inflight_sources: Vec::new(),
577 quarantined_sources: Vec::new(),
578 compute_sources: HashMap::new(),
579 pread, worker_reads: HashMap::new(), pread_requested, pread_fallbacks: 0,
580 copy_stream: e.copy_stream.clone(), copy_stream_unknown: false,
581 compute_stream: e.stream().clone(),
582 compute_stream_unknown: false,
583 n,
584 max_block_bytes,
585 size_aware,
586 frequency_evict,
587 frequency_decay,
588 mtp_frequency_weight,
589 last_forward_layer: None,
590 last_forward_t: 0,
591 frozen: false,
592 per_layer: HashMap::new(),
593 dev_rows: HashMap::new(),
594 prewarm_tried: HashSet::new(),
595 hits: 0, misses: 0, staged_bytes: 0,
596 })
597 }
598
599 #[inline]
600 pub fn n_slots(&self) -> usize { self.n
601 }
602 #[inline]
603 pub fn is_frozen(&self) -> bool {
604 self.frozen
605 }
606 pub fn freeze(&mut self) {
607 if !self.frozen {
608 self.frozen = true;
609 let (_, complete, one_projection, two_projections, stranded_blocks) =
610 self.expert_residency_shape();
611 eprintln!(
612 "[moe-cache] residency frozen: {} slots, {} resident blocks; \
613 {complete} complete experts, {one_projection} one-projection fragments, \
614 {two_projections} two-projection fragments ({stranded_blocks} stranded blocks)",
615 self.n,
616 self.table.len()
617 );
618 let mut mtp_masks = HashMap::<u16, u8>::new();
619 for id in self.table.keys().filter(|id| id.layer == u16::MAX) {
620 *mtp_masks.entry(id.ex).or_insert(0) |= 1u8 << id.proj;
621 }
622 if !mtp_masks.is_empty() {
623 let complete = mtp_masks.values().filter(|&&mask| mask == 0b111).count();
624 eprintln!(
625 "[moe-cache] frozen MTP residency: {} blocks, {complete} complete experts",
626 mtp_masks.values().map(|mask| mask.count_ones() as usize).sum::<usize>()
627 );
628 }
629 }
630 }
631
632 pub(crate) fn expert_residency_shape(&self) -> (usize, usize, usize, usize, usize) {
633 let mut masks = HashMap::<(u16, u16), u8>::new();
634 for id in self.table.keys() {
635 *masks.entry((id.layer, id.ex)).or_insert(0) |= 1u8 << id.proj;
636 }
637 let complete = masks.values().filter(|&&mask| mask == 0b111).count();
638 let one_projection = masks
639 .values()
640 .filter(|&&mask| mask.count_ones() == 1)
641 .count();
642 let two_projections = masks
643 .values()
644 .filter(|&&mask| mask.count_ones() == 2)
645 .count();
646 let stranded_blocks = one_projection + 2 * two_projections;
647 (
648 masks.len(),
649 complete,
650 one_projection,
651 two_projections,
652 stranded_blocks,
653 )
654 }
655
656 #[inline]
657 pub fn max_block_bytes(&self) -> usize {
658 self.max_block_bytes
659 }
660
661
662 #[inline]
664 pub fn resident(&self, id: BlockId) -> Option<usize> { self.table.get(&id).copied() }
665
666 #[inline]
667 fn frequency_increment(&self, id: BlockId) -> f32 {
668 if id.layer == u16::MAX { self.mtp_frequency_weight } else { 1.0 }
669 }
670
671 pub(crate) fn note_profile_hit(&mut self, id: BlockId) {
675 if self.frozen || !self.table.contains_key(&id) {
676 return;
677 }
678 let increment = self.frequency_increment(id);
679 *self.frequencies.entry(id).or_insert(0.0) += increment;
680 }
681
682 fn on_hit(&mut self, slot: usize) {
700 let class_index = self.slot_class[slot];
701 let class = &mut self.classes[class_index];
702 if !class.free.is_empty() {
703 return;
704 }
705 class.on_hit_full(slot, &mut self.links);
706 }
707
708 fn remove_occupant(&mut self, slot: usize) {
709 if let Some(old) = self.occupant[slot].take() {
710 self.table.remove(&old);
711 self.on_block_evicted(old.layer);
712 }
713 }
714
715 fn frequency_victim_in_class(&mut self, class_index: usize, keep: &[BlockId]) -> Option<usize> {
722 let class = &self.classes[class_index];
723 let candidate = class
724 .probation
725 .iter(&self.links)
726 .chain(class.protected.iter(&self.links))
727 .enumerate()
728 .filter_map(|(position, slot)| {
729 let id = self.occupant[slot]?;
730 (!keep.contains(&id)).then_some((
731 self.frequencies.get(&id).copied().unwrap_or(0.0),
732 position,
733 slot,
734 ))
735 })
736 .min_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
737 let (_, _, slot) = candidate?;
738 self.classes[class_index].unlink_from_segment(slot, &mut self.links);
739 Some(slot)
740 }
741
742 fn evict_one(&mut self, required: usize) -> Option<usize> {
744 for class_index in 0..self.classes.len() {
745 if self.classes[class_index].capacity < required {
746 continue;
747 }
748 let slot = if self.frequency_evict {
749 self.frequency_victim_in_class(class_index, &[])
750 } else {
751 self.classes[class_index].pop_lru(&mut self.links)
752 };
753 if let Some(slot) = slot {
754 self.remove_occupant(slot);
755 return Some(slot);
756 }
757 }
758 None
759 }
760
761 fn evict_one_excluding(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
767 fn take(
768 q: &mut SlruList,
769 links: &mut [SlotLink],
770 occupant: &[Option<BlockId>],
771 keep: &[BlockId],
772 ) -> Option<usize> {
773 let slot = q
774 .iter(links)
775 .find(|&s| occupant[s].is_some_and(|id| !keep.contains(&id)))?;
776 q.unlink(slot, links);
777 Some(slot)
778 }
779 for class_index in 0..self.classes.len() {
780 if self.classes[class_index].capacity < required {
781 continue;
782 }
783 let slot = if self.frequency_evict {
784 self.frequency_victim_in_class(class_index, keep)
785 } else {
786 let class = &mut self.classes[class_index];
787 take(&mut class.probation, &mut self.links, &self.occupant, keep)
788 .or_else(|| take(&mut class.protected, &mut self.links, &self.occupant, keep))
789 };
790 if let Some(slot) = slot {
791 self.remove_occupant(slot);
792 return Some(slot);
793 }
794 }
795 None
796 }
797
798 fn on_block_evicted(&mut self, layer: u16) {
804 if let Some(c) = self.per_layer.get_mut(&layer) { *c -= 1; }
805 self.dev_rows.remove(&layer);
806 }
807
808 fn reserve_slot(&mut self, required: usize) -> Option<usize> {
809 for class in &mut self.classes {
810 if class.capacity >= required {
811 if let Some(slot) = class.free.pop() {
812 return Some(slot);
813 }
814 }
815 }
816 self.evict_one(required)
817 }
818
819 fn release_reserved_slot(&mut self, slot: usize) {
820 debug_assert!(self.occupant[slot].is_none());
821 self.classes[self.slot_class[slot]].free.push(slot);
822 }
823
824 fn publish(&mut self, id: BlockId, slot: usize) {
825 self.occupant[slot] = Some(id);
826 self.table.insert(id, slot);
827 self.classes[self.slot_class[slot]]
828 .probation
829 .push_back(slot, SEG_PROBATION, &mut self.links);
830 *self.per_layer.entry(id.layer).or_insert(0) += 1;
831 }
832
833
834
835 fn reap_copy_sources(&mut self) {
836 self.inflight_sources.retain(|(ready, _)| !ready.is_complete());
837 }
838
839 fn retain_compute_source(&mut self, owner: Option<ExpertKeepalive>) {
840 if let Some(owner) = owner {
841 let key = KeepaliveKey::from_owner(&owner);
842 self.compute_sources.entry(key).or_insert(owner);
843 }
844 }
845
846 fn admit(&mut self, id: BlockId, host_bytes: &[u8], e: &Engine)
849 -> Result<usize, Box<dyn std::error::Error>> {
850 let slot = self.reserve_slot(host_bytes.len()).ok_or_else(|| {
851 std::io::Error::other(format!(
852 "no MoE cache slot can hold {} bytes (max class {})",
853 host_bytes.len(),
854 self.classes.last().map(|class| class.capacity).unwrap_or(0)
855 ))
856 })?;
857 if let Err(err) = e.stage_expert(host_bytes, &mut self.slots[slot], 0) {
860 return match e.stream().synchronize() {
861 Ok(()) => {
862 self.release_reserved_slot(slot);
863 Err(err)
864 }
865 Err(sync_err) => {
866 self.compute_stream_unknown = true;
869 Err(std::io::Error::other(format!(
870 "compute-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
871 )).into())
872 }
873 };
874 }
875 self.staged_bytes += host_bytes.len() as u64;
876 self.publish(id, slot);
877 Ok(slot)
878 }
879
880 fn note_pread_fallback(&mut self, reason: &dyn std::fmt::Display) {
881 self.pread_fallbacks += 1;
882 if let Some(pool) = self.pread.as_mut() {
883 pool.note_fallback();
884 }
885 if self.pread_fallbacks <= 3 {
886 eprintln!("[spill-pread] falling back to mmap: {reason}");
887 }
888 }
889
890 pub(crate) fn begin_worker_scope(&mut self) {
894 if self.worker_reads.is_empty() { return; }
895 let tickets: Vec<_> = self
896 .worker_reads
897 .drain()
898 .map(|(_, read)| read.ticket)
899 .collect();
900 if let Some(pool) = self.pread.as_mut().filter(|pool| pool.is_worker()) {
901 for ticket in tickets {
902 let _ = pool.cancel_worker(ticket); }
903 }
904 }
905
906 pub(crate) fn begin_forward_epoch(&mut self, layer: u16, t: usize) {
916 let Some(decay) = self.frequency_decay else {
917 self.last_forward_layer = Some(layer);
918 self.last_forward_t = t;
919 return;
920 };
921 let new_sweep = self
922 .last_forward_layer
923 .is_some_and(|previous| layer <= previous);
924 if new_sweep && t == 1 {
925 if self.last_forward_t != 1 {
926 self.frequencies.clear();
927 } else {
928 self.frequencies.retain(|_, score| {
929 *score *= decay;
930 *score >= 1.0e-3
931 });
932 }
933 }
934 self.last_forward_layer = Some(layer);
935 self.last_forward_t = t;
936 }
937
938 pub(crate) fn promote_worker_reads_at_safe_boundary(
945 &mut self,
946 order: &[BlockId],
947 keep: &[BlockId],
948 e: &Engine,
949 ) -> Result<usize, Box<dyn std::error::Error>> {
950 if !crate::spill_pread::copy_h2d_enabled() {
951 return Ok(0);
952 }
953 let mut promoted = 0usize;
954 for &id in order {
955 if self.table.contains_key(&id) || self.pending.contains_key(&id) {
956 if let Some(read) = self.worker_reads.remove(&id) {
957 if let Some(pool) = self.pread.as_mut() {
958 let _ = pool.cancel_worker(read.ticket);
959 }
960 }
961 continue;
962 }
963 let Some(read) = self.worker_reads.get(&id).copied() else {
964 continue;
965 };
966 let Some(slot) = self.reserve_prefetch_slot(read.len, keep) else {
967 continue;
968 };
969 self.worker_reads.remove(&id);
970
971 let index = match self.pread.as_mut().unwrap().wait_worker(read.ticket) {
972 Ok(index) => index,
973 Err(err) => {
974 let _ = self.pread.as_mut().unwrap().cancel_worker(read.ticket);
975 self.release_reserved_slot(slot);
976 self.note_pread_fallback(err.as_ref());
977 continue;
978 }
979 };
980 let ready = {
981 let bytes = match self.pread.as_ref().unwrap().bytes(index, read.len) {
982 Ok(bytes) => bytes,
983 Err(err) => {
984 self.pread.as_mut().unwrap().abort_read(index);
985 self.release_reserved_slot(slot);
986 self.note_pread_fallback(err.as_ref());
987 continue;
988 }
989 };
990 stage_pread_prefetch_on_copy_stream(e, bytes, &mut self.slots[slot])
991 };
992 let ready = match ready {
993 Ok(ready) => ready,
994 Err((err, reusable)) => {
995 if reusable {
996 self.pread.as_mut().unwrap().abort_read(index);
997 self.release_reserved_slot(slot);
998 self.note_pread_fallback(err.as_ref());
999 continue;
1000 }
1001 self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1002 self.copy_stream_unknown = true;
1003 return Err(err);
1004 }
1005 };
1006 self.pread.as_mut().unwrap().mark_h2d(index, ready.clone());
1007 self.occupant[slot] = Some(id);
1008 self.pending.insert(
1009 id,
1010 PendingBlock {
1011 slot,
1012 ready,
1013 keepalive: None,
1014 },
1015 );
1016 self.staged_bytes += read.len as u64;
1017 promoted += 1;
1018 }
1019 Ok(promoted)
1020 }
1021
1022 fn dispatch_disk(
1023 &mut self,
1024 id: BlockId,
1025 file: &Arc<std::fs::File>,
1026 offset: u64,
1027 len: usize,
1028 fallback: &[u8],
1029 e: &Engine,
1030 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1031 if self.pread.is_none() {
1032 if self.pread_requested {
1033 self.note_pread_fallback(&"pinned-buffer backend unavailable");
1034 }
1035 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1036 }
1037
1038 let pending = self.worker_reads.remove(&id);
1039 let pool = self.pread.as_mut().unwrap();
1040 let read = if pool.is_worker() {
1041 let ticket = match pending {
1042 Some(read) => Ok(Some(read.ticket)),
1043 None => pool.submit_worker(file.clone(), offset, len),
1044 };
1045 match ticket {
1046
1047 Ok(Some(ticket)) => match pool.wait_worker(ticket) {
1048 Ok(index) => Ok(index),
1049 Err(err) => {
1050 let _ = pool.cancel_worker(ticket);
1053 Err(err)
1054 }
1055 },
1056 Ok(None) => Err(std::io::Error::other("worker read ring is busy").into()),
1057 Err(err) => Err(err),
1058 }
1059 } else {
1060 debug_assert!(pending.is_none());
1061 pool.read(file.as_ref(), offset, len)
1062 };
1063 let index = match read {
1064 Ok(index) => index,
1065 Err(err) => {
1066 self.note_pread_fallback(err.as_ref());
1067 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1068 }
1069 };
1070
1071
1072 let slot = self.reserve_slot(len).ok_or_else(|| {
1075 std::io::Error::other(format!(
1076 "no MoE cache slot can hold {len} bytes (max class {})",
1077 self.classes.last().map(|class| class.capacity).unwrap_or(0)
1078 ))
1079 })?;
1080 let ready = {
1081 let bytes = match self.pread.as_ref().unwrap().bytes(index, len) {
1082 Ok(bytes) => bytes,
1083 Err(err) => {
1084 self.pread.as_mut().unwrap().abort_read(index);
1085 self.release_reserved_slot(slot);
1086 self.note_pread_fallback(err.as_ref());
1087 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1088 }
1089 };
1090 stage_pread_on_compute_stream(e, bytes, &mut self.slots[slot])
1091 };
1092 let ready = match ready {
1093 Ok(ready) => ready,
1094 Err(err) => {
1095 match e.stream().synchronize() {
1099 Ok(()) => {
1100 self.pread.as_mut().unwrap().abort_read(index);
1101 self.release_reserved_slot(slot);
1102 self.note_pread_fallback(err.as_ref());
1103 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1104 }
1105 Err(sync_err) => {
1106 self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1107 return Err(std::io::Error::other(format!(
1108 "pread H2D setup failed ({err}); CUDA stream drain also failed ({sync_err})"
1109 )).into());
1110 }
1111 }
1112 }
1113 };
1114 self.pread.as_mut().unwrap().mark_h2d(index, ready);
1115 self.staged_bytes += len as u64;
1118 self.publish(id, slot);
1119 Ok(DispatchSlot::Resident(slot))
1120 }
1121
1122 pub fn dispatch(&mut self, id: BlockId, host_bytes: &[u8], e: &Engine)
1130 -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1131 self.dispatch_source(id, ExpertSource::Memory { bytes: host_bytes, keepalive: None }, e)
1132 }
1133
1134 pub(crate) fn dispatch_source(&mut self, id: BlockId, source: ExpertSource<'_>, e: &Engine)
1135 -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1136 self.reap_copy_sources();
1137 let increment = self.frequency_increment(id);
1138 *self.frequencies.entry(id).or_insert(0.0) += increment;
1139 if let Some(s) = self.table.get(&id).copied() {
1140 self.hits += 1;
1141 self.on_hit(s);
1142 return Ok(DispatchSlot::Resident(s));
1143 }
1144 if let Some(pending) = self.pending.remove(&id) {
1145 if let Err(err) = e.compute_wait(pending.ready.as_ref()) {
1146 self.pending.insert(id, pending);
1147 return Err(err);
1148 }
1149 self.misses += 1;
1150 let slot = pending.slot;
1151 if let Some(keepalive) = pending.keepalive {
1152 self.inflight_sources.push((pending.ready, keepalive));
1153 }
1154 self.publish(id, slot);
1155 return Ok(DispatchSlot::Resident(slot));
1156 }
1157 self.misses += 1;
1158 match source {
1168 ExpertSource::Memory { bytes, keepalive } => {
1169 self.retain_compute_source(keepalive);
1172 let slot = self.admit(id, bytes, e)?;
1173 Ok(DispatchSlot::Resident(slot))
1174 }
1175 ExpertSource::Disk { file, offset, len, fallback, keepalive } => {
1176 self.retain_compute_source(Some(keepalive));
1179 self.dispatch_disk(id, file, offset, len, fallback, e)
1180 }
1181 }
1182 }
1183
1184 pub fn prefetch(&mut self, id: BlockId, host_bytes: &[u8], keep: &[BlockId], e: &Engine)
1193 -> Result<bool, Box<dyn std::error::Error>> {
1194 self.prefetch_source(
1195 id,
1196 ExpertSource::Memory { bytes: host_bytes, keepalive: None },
1197 keep,
1198 e,
1199 )
1200 }
1201
1202 fn reserve_prefetch_slot(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
1203 for class in &mut self.classes {
1204 if class.capacity >= required {
1205 if let Some(slot) = class.free.pop() {
1206 return Some(slot);
1207 }
1208 }
1209 }
1210 self.evict_one_excluding(required, keep)
1211 }
1212
1213 fn prefetch_bytes(
1214&mut self, id: BlockId, host_bytes: &[u8],
1215 keepalive: Option<ExpertKeepalive>, keep: &[BlockId], e: &Engine)
1216 -> Result<bool, Box<dyn std::error::Error>> {
1217 let Some(slot) = self.reserve_prefetch_slot(host_bytes.len(), keep) else { return Ok(false) };
1218 let ready = match stage_on_copy_stream(e, host_bytes, &mut self.slots[slot]) {
1219 Ok(ready) => ready,
1220 Err((err, reusable)) => {
1221 if reusable {
1222 self.release_reserved_slot(slot);
1223 } else {
1224 self.copy_stream_unknown = true;
1227 if let Some(keepalive) = keepalive {
1228 self.quarantined_sources.push(keepalive);
1229 }
1230 eprintln!("[moe-cache] quarantining slot {slot} after unprovable copy completion");
1231 }
1232 return Err(err);
1233 }
1234 };
1235 self.occupant[slot] = Some(id);
1236 self.pending.insert(id, PendingBlock { slot, ready, keepalive });
1237 self.staged_bytes += host_bytes.len() as u64;
1238 Ok(true)
1239 }
1240
1241 pub(crate) fn prefetch_source(
1242 &mut self,
1243 id: BlockId,
1244 source: ExpertSource<'_>,
1245 keep: &[BlockId],
1246 e: &Engine,
1247 ) -> Result<bool, Box<dyn std::error::Error>> {
1248 self.reap_copy_sources();
1249 if self.table.contains_key(&id) || self.pending.contains_key(&id)
1250 || self.worker_reads.contains_key(&id) {
1251 return Ok(false);
1252 }
1253 match source {
1254 ExpertSource::Memory { bytes, keepalive } => {
1255 self.prefetch_bytes(id, bytes, keepalive, keep, e)
1256 }
1257 ExpertSource::Disk { file, offset, len, fallback, keepalive } => {
1258 if self.pread.as_ref().is_some_and(PreadPool::is_worker) {
1259 match self.pread.as_mut().unwrap()
1260 .submit_worker_speculative(file.clone(), offset, len) {
1261 Ok(Some(ticket)) => {
1262 self.worker_reads.insert(id, WorkerRead { ticket, len });
1263 Ok(true)
1264 }
1265 Ok(None) => Ok(false),
1266 Err(err) => {
1267 self.note_pread_fallback(err.as_ref());
1268 Ok(false)
1269 }
1270 }
1271 } else if self.pread.is_some() {
1272 Ok(false)
1274 } else {
1275 self.prefetch_bytes(id, fallback, Some(keepalive), keep, e)
1276 }
1277 }
1278 }
1279 }
1280
1281 pub fn force_admit(&mut self, id: BlockId, host_bytes: &[u8], e: &Engine)
1283 -> Result<usize, Box<dyn std::error::Error>> {
1284 if let Some(s) = self.table.get(&id).copied() { return Ok(s); }
1285 self.admit(id, host_bytes, e)
1286 }
1287
1288 pub fn export_residency(&self) -> Vec<(u16, u8, u16)> {
1296 self.occupant
1297 .iter()
1298 .flatten()
1299 .map(|id| (id.layer, id.proj, id.ex))
1300 .collect()
1301 }
1302
1303 pub fn restage_block(&mut self, id: BlockId, m: &crate::hybrid::MoeWeights, e: &Engine)
1308 -> Result<bool, Box<dyn std::error::Error>> {
1309 if self.table.contains_key(&id) {
1310 return Ok(true);
1311 }
1312 let exps = match id.proj {
1313 PROJ_GATE => &m.gate_exps,
1314 PROJ_UP => &m.up_exps,
1315 PROJ_DOWN => &m.down_exps,
1316 _ => return Ok(false),
1317 };
1318 if id.ex as usize >= exps.n_expert {
1319 return Ok(false);
1320 }
1321 if m.active_experts.as_ref().is_some_and(|active| !active[id.ex as usize]) {
1322 return Ok(false);
1323 }
1324 if exps.expert_layout(id.ex as usize).len == 0 {
1325 return Ok(false);
1326 }
1327 match exps.expert_source(id.ex as usize) {
1328 ExpertSource::Memory { bytes, keepalive } => {
1329 self.retain_compute_source(keepalive);
1330 self.admit(id, bytes, e)?;
1331 }
1332 ExpertSource::Disk { fallback, keepalive, .. } => {
1333 self.retain_compute_source(Some(keepalive));
1334 self.admit(id, fallback, e)?;
1335 }
1336 }
1337 Ok(true)
1338 }
1339
1340 pub fn prewarm_layer(&mut self, layer: u16, m: &crate::hybrid::MoeWeights, e: &Engine)
1341 -> Result<(), Box<dyn std::error::Error>> {
1342 if !self.prewarm_tried.insert(layer) { return Ok(()); }
1343 let n_expert = m.gate_exps.n_expert;
1344 if self.pread.is_some() && (0..n_expert).any(|ex| {
1345 matches!(m.gate_exps.expert_source(ex), ExpertSource::Disk { .. })
1346 || matches!(m.up_exps.expert_source(ex), ExpertSource::Disk { .. })
1347 || matches!(m.down_exps.expert_source(ex), ExpertSource::Disk { .. })
1348 }) {
1349 return Ok(());
1352 }
1353 let resident = self.per_layer.get(&layer).copied().unwrap_or(0) as usize;
1354 let missing = 3 * n_expert - resident;
1355 if self.size_aware {
1356 return Ok(());
1357 } if self
1359 .classes
1360 .iter()
1361 .map(|class| class.free.len())
1362 .sum::<usize>()
1363 < missing
1364 {
1365 return Ok(()); }
1367 for ex in 0..n_expert {
1368 for (proj, exps) in [
1369 (PROJ_GATE, &m.gate_exps),
1370 (PROJ_UP, &m.up_exps),
1371 (PROJ_DOWN, &m.down_exps)] {
1372 let id = BlockId::new(layer, proj, ex as u16);
1373 if self.table.contains_key(&id) { continue; }
1374 match exps.expert_source(ex) {
1375 ExpertSource::Memory { bytes, keepalive } => {
1376 self.retain_compute_source(keepalive);
1377 self.admit(id, bytes, e)?;
1378 }
1379 ExpertSource::Disk { fallback, keepalive, .. } => {
1380 self.retain_compute_source(Some(keepalive));
1381 self.admit(id, fallback, e)?;
1382 }
1383 }
1384 }
1385 }
1386 Ok(())
1387 }
1388
1389 pub fn layer_dev_row(&mut self, layer: u16, n_expert: usize, e: &Engine)
1395 -> Result<Option<&CudaSlice<u64>>, Box<dyn std::error::Error>> {
1396 if self.per_layer.get(&layer).copied().unwrap_or(0) as usize != 3 * n_expert {
1397 return Ok(None);
1398 }
1399 if !self.dev_rows.contains_key(&layer) {
1400 use cudarc::driver::DevicePtr;
1401 let mut host = vec![0u64; 3 * n_expert];
1402 for proj in 0..3u8 {
1403 for ex in 0..n_expert {
1404 let Some(&s) = self.table.get(&BlockId::new(layer, proj, ex as u16)) else {
1405 return Ok(None);
1407 };
1408 let __s_ev = e.stream();
1409 let (p, _ev) = self.slots[s].device_ptr(&__s_ev);
1410 host[proj as usize * n_expert + ex] = p as u64;
1411 }
1412 }
1413 let row = e.stream().clone_htod(&host)?;
1414 self.dev_rows.insert(layer, row);
1415 }
1416 Ok(self.dev_rows.get(&layer))
1417 }
1418
1419 #[inline]
1421 pub fn buf(&self, d: DispatchSlot) -> &CudaSlice<u8> {
1422 match d {
1423 DispatchSlot::Resident(s) => &self.slots[s],
1424 }
1425 }
1426
1427 #[inline]
1429 pub fn slot(&self, s: usize) -> &CudaSlice<u8> { &self.slots[s] }
1430
1431 pub fn hit_rate(&self) -> f64 {
1433 let tot = self.hits + self.misses;
1434 if tot == 0 { 0.0 } else { self.hits as f64 / tot as f64 }
1435 }
1436
1437 pub fn reset_counters(&mut self) {
1439 self.hits = 0; self.misses = 0; self.staged_bytes = 0;
1440 }
1441
1442 pub(crate) fn pread_stats(&self) -> Option<PreadStats> {
1443 if !self.pread_requested { return None; }
1444 let mut stats = self.pread.as_ref().map(PreadPool::stats).unwrap_or_default();
1445 stats.fallbacks = self.pread_fallbacks;
1446 Some(stats)
1447 }
1448}
1449
1450fn cache_lfu_decay() -> Option<f32> {
1451 let raw = std::env::var("MEMRA_MOE_LFU_DECAY").ok()?;
1452 match parse_cache_lfu_decay(Some(&raw)) {
1453 Ok(value) => value,
1454 Err(reason) => {
1455 eprintln!(
1456 "[moe-cache] invalid MEMRA_MOE_LFU_DECAY={raw:?} ({reason}); disabling LFU decay"
1457 );
1458 None
1459 }
1460 }
1461}
1462
1463fn cache_lfu_mtp_weight() -> f32 {
1464 const DEFAULT: f32 = 1.0;
1465 let raw = std::env::var("MEMRA_MOE_LFU_MTP_WEIGHT").ok();
1466 match parse_cache_lfu_mtp_weight(raw.as_deref()) {
1467 Ok(value) => value,
1468 Err(reason) => {
1469 eprintln!(
1470 "[moe-cache] invalid MEMRA_MOE_LFU_MTP_WEIGHT={:?} ({reason}); using {DEFAULT}",
1471 raw.as_deref().unwrap_or("")
1472 );
1473 DEFAULT
1474 }
1475 }
1476}
1477
1478fn parse_cache_lfu_mtp_weight(raw: Option<&str>) -> Result<f32, &'static str> {
1479 let value = raw
1480 .unwrap_or("1")
1481 .parse::<f32>()
1482 .map_err(|_| "expected a number")?;
1483 if value.is_finite() && (0.25..=64.0).contains(&value) {
1484 Ok(value)
1485 } else {
1486 Err("expected a finite multiplier from 0.25 through 64")
1487 }
1488}
1489
1490fn parse_cache_lfu_decay(raw: Option<&str>) -> Result<Option<f32>, &'static str> {
1491 let Some(raw) = raw else { return Ok(None) };
1492 let value = raw.parse::<f32>().map_err(|_| "expected a number")?;
1493 if value.is_finite() && value > 0.0 && value <= 1.0 {
1494 Ok(Some(value))
1495 } else {
1496 Err("expected a finite fraction greater than 0 and at most 1")
1497 }
1498}
1499
1500fn cache_hard_vram_frac() -> f64 {
1501 const DEFAULT: f64 = 0.80;
1502 let raw = std::env::var("MEMRA_MOE_HARD_VRAM_FRAC").ok();
1503 match parse_cache_hard_vram_frac(raw.as_deref()) {
1504 Ok(value) => value,
1505 Err(reason) => {
1506 eprintln!(
1507 "[moe-cache] invalid MEMRA_MOE_HARD_VRAM_FRAC={:?} ({reason}); using {DEFAULT}",
1508 raw.as_deref().unwrap_or("")
1509 );
1510 DEFAULT
1511 }
1512 }
1513}
1514
1515fn parse_cache_hard_vram_frac(raw: Option<&str>) -> Result<f64, &'static str> {
1516 let value = raw
1517 .unwrap_or("0.80")
1518 .parse::<f64>()
1519 .map_err(|_| "expected a number")?;
1520 if value.is_finite() && (0.10..=0.95).contains(&value) {
1521 Ok(value)
1522 } else {
1523 Err("expected a finite fraction from 0.10 through 0.95")
1524 }
1525}
1526
1527#[cfg(test)]
1528mod slru_intrusive_tests {
1529 use super::{SlotClass, SlotLink, SlruList, SEG_PROBATION, SEG_PROTECTED};
1536 use std::collections::VecDeque;
1537
1538 struct OldSlru {
1540 probation: VecDeque<usize>,
1541 protected: VecDeque<usize>,
1542 protected_cap: usize,
1543 }
1544 impl OldSlru {
1545 fn on_hit_full(&mut self, slot: usize) {
1546 if let Some(pos) = self.probation.iter().position(|&x| x == slot) {
1547 self.probation.remove(pos);
1548 self.push_protected(slot);
1549 } else if let Some(pos) = self.protected.iter().position(|&x| x == slot) {
1550 self.protected.remove(pos);
1551 self.protected.push_back(slot); } else {
1553 self.push_protected(slot);
1554 }
1555 }
1556 fn push_protected(&mut self, slot: usize) {
1557 self.protected.push_back(slot);
1558 while self.protected.len() > self.protected_cap {
1559 if let Some(demoted) = self.protected.pop_front() {
1560 self.probation.push_back(demoted);
1561 } else {
1562 break;
1563 }
1564 }
1565 }
1566 fn pop_lru(&mut self) -> Option<usize> {
1567 self.probation.pop_front().or_else(|| self.protected.pop_front())
1568 }
1569 fn take_excluding(&mut self, banned: &[usize]) -> Option<usize> {
1570 let take = |q: &mut VecDeque<usize>| {
1571 q.iter()
1572 .position(|&s| !banned.contains(&s))
1573 .and_then(|pos| q.remove(pos))
1574 };
1575 take(&mut self.probation).or_else(|| take(&mut self.protected))
1576 }
1577 }
1578
1579 fn new_pair(n: usize, protected_cap: usize) -> (SlotClass, Vec<SlotLink>, OldSlru) {
1580 let class = SlotClass {
1581 capacity: 1,
1582 probation: SlruList::new(),
1583 protected: SlruList::new(),
1584 free: Vec::new(),
1585 protected_cap,
1586 };
1587 let links = vec![SlotLink::none(); n];
1588 let old = OldSlru {
1589 probation: VecDeque::new(),
1590 protected: VecDeque::new(),
1591 protected_cap,
1592 };
1593 (class, links, old)
1594 }
1595
1596 fn orders_match(class: &SlotClass, links: &[SlotLink], old: &OldSlru) -> bool {
1597 let np: Vec<usize> = class.probation.iter(links).collect();
1598 let nt: Vec<usize> = class.protected.iter(links).collect();
1599 let op: Vec<usize> = old.probation.iter().copied().collect();
1600 let ot: Vec<usize> = old.protected.iter().copied().collect();
1601 np == op && nt == ot
1602 }
1603
1604 struct Rng(u64);
1606 impl Rng {
1607 fn next(&mut self) -> u64 {
1608 self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
1609 let mut z = self.0;
1610 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1611 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1612 z ^ (z >> 31)
1613 }
1614 fn below(&mut self, n: usize) -> usize {
1615 (self.next() % n as u64) as usize
1616 }
1617 }
1618
1619 #[test]
1620 fn same_eviction_decisions_randomized_soak() {
1621 for &(n, cap) in &[(8usize, 1usize), (16, 12), (64, 51), (128, 102)] {
1623 let (mut class, mut links, mut old) = new_pair(n, cap);
1624 let mut rng = Rng(0xC0FFEE ^ (n as u64) << 8 ^ cap as u64);
1625 let mut resident: Vec<usize> = Vec::new();
1626 let mut free: Vec<usize> = (0..n).rev().collect();
1627 for step in 0..200_000 {
1628 let op = rng.below(100);
1629 if op < 55 && !resident.is_empty() {
1630 let slot = resident[rng.below(resident.len())];
1632 class.on_hit_full(slot, &mut links);
1633 old.on_hit_full(slot);
1634 } else if op < 80 {
1635 let slot = if let Some(s) = free.pop() {
1637 s
1638 } else {
1639 let v_new = class.pop_lru(&mut links);
1640 let v_old = old.pop_lru();
1641 assert_eq!(v_new, v_old, "victim diverged at step {step} (n={n} cap={cap})");
1642 let v = v_new.unwrap();
1643 resident.retain(|&s| s != v);
1644 v
1645 };
1646 class.probation.push_back(slot, SEG_PROBATION, &mut links);
1647 old.probation.push_back(slot);
1648 resident.push(slot);
1649 } else if op < 92 && resident.len() > 2 {
1650 let banned: Vec<usize> = (0..3.min(resident.len()))
1652 .map(|_| resident[rng.below(resident.len())])
1653 .collect();
1654 let take_new = {
1655 let q = &mut class.probation;
1656 let found = q.iter(&links).find(|s| !banned.contains(s));
1657 match found {
1658 Some(s) => {
1659 q.unlink(s, &mut links);
1660 Some(s)
1661 }
1662 None => {
1663 let q = &mut class.protected;
1664 q.iter(&links).find(|s| !banned.contains(s)).inspect(|&s| {
1665 q.unlink(s, &mut links);
1666 })
1667 }
1668 }
1669 };
1670 let take_old = old.take_excluding(&banned);
1671 assert_eq!(take_new, take_old, "excluding-victim diverged at step {step}");
1672 if let Some(v) = take_new {
1673 resident.retain(|&s| s != v);
1674 free.push(v);
1675 }
1676 } else if !resident.is_empty() {
1677 let slot = resident[rng.below(resident.len())];
1679 match links[slot].seg {
1680 SEG_PROBATION => class.probation.unlink(slot, &mut links),
1681 SEG_PROTECTED => class.protected.unlink(slot, &mut links),
1682 _ => {}
1683 }
1684 if let Some(pos) = old.probation.iter().position(|&x| x == slot) {
1685 old.probation.remove(pos);
1686 } else if let Some(pos) = old.protected.iter().position(|&x| x == slot) {
1687 old.protected.remove(pos);
1688 }
1689 class.on_hit_full(slot, &mut links);
1690 old.on_hit_full(slot);
1691 }
1692 assert!(
1693 orders_match(&class, &links, &old),
1694 "segment order diverged at step {step} (n={n} cap={cap})"
1695 );
1696 }
1697 }
1698 }
1699
1700 #[test]
1701 fn hit_promotion_is_o1_not_on() {
1702 fn bench(n: usize, hits: usize) -> std::time::Duration {
1707 let (mut class, mut links, _) = new_pair(n, (n as f64 * 0.8) as usize);
1708 for s in 0..n {
1709 class.probation.push_back(s, SEG_PROBATION, &mut links);
1710 }
1711 let mut rng = Rng(0xBEEF);
1712 let t0 = std::time::Instant::now();
1713 for _ in 0..hits {
1714 class.on_hit_full(rng.below(n), &mut links);
1715 }
1716 t0.elapsed()
1717 }
1718 bench(1_000, 10_000);
1720 bench(46_000, 10_000);
1721 let small = bench(1_000, 850_000).as_secs_f64() / 850_000.0;
1722 let large = bench(46_000, 850_000).as_secs_f64() / 850_000.0;
1723 assert!(
1724 large < small * 8.0,
1725 "per-hit cost scaled with n_slots: {:.1}ns @1k vs {:.1}ns @46k",
1726 small * 1e9,
1727 large * 1e9
1728 );
1729 }
1730
1731 #[test]
1732 fn slru_list_basic_invariants() {
1733 let mut links = vec![SlotLink::none(); 4];
1734 let mut l = SlruList::new();
1735 assert_eq!(l.pop_front(&mut links), None);
1736 l.push_back(2, SEG_PROBATION, &mut links);
1737 l.push_back(0, SEG_PROBATION, &mut links);
1738 l.push_back(3, SEG_PROBATION, &mut links);
1739 assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 0, 3]);
1740 assert_eq!(l.len, 3);
1741 l.unlink(0, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 3]);
1743 l.unlink(3, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2]);
1745 assert_eq!(l.pop_front(&mut links), Some(2)); assert_eq!(l.len, 0);
1747 assert_eq!(l.head, super::NIL);
1748 assert_eq!(l.tail, super::NIL);
1749 assert!(links.iter().all(|k| k.seg == super::SEG_NONE));
1750 }
1751}
1752
1753#[cfg(test)]
1754mod vram_fraction_tests {
1755 use super::{
1756 parse_cache_hard_vram_frac, parse_cache_lfu_decay, parse_cache_lfu_mtp_weight,
1757 size_class_plan,
1758 };
1759
1760 #[test]
1761 fn hard_vram_fraction_defaults_and_rejects_unsafe_values() {
1762 assert_eq!(parse_cache_hard_vram_frac(None), Ok(0.80));
1763 assert_eq!(parse_cache_hard_vram_frac(Some("0.82")), Ok(0.82));
1764 assert!(parse_cache_hard_vram_frac(Some("NaN")).is_err());
1765 assert_eq!(parse_cache_hard_vram_frac(Some("0.95")), Ok(0.95));
1766 assert!(parse_cache_hard_vram_frac(Some("0.96")).is_err());
1767 assert!(parse_cache_hard_vram_frac(Some("1.0")).is_err());
1768 assert!(parse_cache_hard_vram_frac(Some("bad")).is_err());
1769 }
1770
1771 #[test]
1772 fn lfu_decay_is_opt_in_and_bounded() {
1773 assert_eq!(parse_cache_lfu_decay(None), Ok(None));
1774 assert_eq!(parse_cache_lfu_decay(Some("0.8")), Ok(Some(0.8)));
1775 assert_eq!(parse_cache_lfu_decay(Some("1")), Ok(Some(1.0)));
1776 for value in ["0", "-0.1", "1.1", "NaN", "bad"] {
1777 assert!(
1778 parse_cache_lfu_decay(Some(value)).is_err(),
1779 "accepted {value}"
1780 );
1781 }
1782 }
1783
1784 #[test]
1785 fn lfu_mtp_weight_defaults_and_is_bounded() {
1786 assert_eq!(parse_cache_lfu_mtp_weight(None), Ok(1.0));
1787 assert_eq!(parse_cache_lfu_mtp_weight(Some("4")), Ok(4.0));
1788 for value in ["0", "0.1", "65", "NaN", "bad"] {
1789 assert!(
1790 parse_cache_lfu_mtp_weight(Some(value)).is_err(),
1791 "accepted {value}"
1792 );
1793 }
1794 }
1795
1796 #[test]
1797 fn size_class_plan_preserves_classes_and_never_exceeds_budget() {
1798 let blocks = [100usize, 100, 100, 200, 200, 400];
1799 let budget = (108 * 2) + 208 + 408;
1800 let plan = size_class_plan(&blocks, budget);
1801 assert!(plan.iter().all(|(_, count)| *count > 0));
1802 assert!(
1803 plan.iter()
1804 .map(|(bytes, count)| (bytes + 8) * count)
1805 .sum::<usize>()
1806 <= budget
1807 );
1808 assert!(plan.iter().all(|(bytes, count)| {
1809 *count <= blocks.iter().filter(|block| **block == *bytes).count()
1810 }));
1811 }
1812
1813 #[test]
1814 fn size_class_plan_returns_full_inventory_when_it_fits() {
1815 let blocks = [100usize, 100, 200, 400];
1816 let budget: usize = blocks.iter().map(|bytes| bytes + 8).sum();
1817 assert_eq!(
1818 size_class_plan(&blocks, budget),
1819 vec![(100, 2), (200, 1), (400, 1)]
1820 );
1821 }
1822
1823 #[test]
1824 fn size_class_plan_does_not_overflow_on_pathological_sizes() {
1825 let plan = size_class_plan(&[usize::MAX, usize::MAX], usize::MAX);
1826 assert!(plan.is_empty());
1827 }
1828
1829}
1830
1831impl Drop for MoeSlotCache {
1832 fn drop(&mut self) {
1833 let mut safe_to_drop_slots = true;
1836 if self.compute_stream_unknown || !self.compute_sources.is_empty() {
1837 if let Err(err) = self.compute_stream.synchronize() {
1838 safe_to_drop_slots = false;
1839 eprintln!("[moe-cache] unknown compute-stream drain failed ({err}); leaking GPU slots for safety");
1840 for (_, keepalive) in self.compute_sources.drain() {
1841 std::mem::forget(keepalive);
1842 }
1843 } else {
1844 self.compute_stream_unknown = false;
1845 self.compute_sources.clear();
1846 }
1847 }
1848 let need_copy_drain = self.copy_stream_unknown || !self.pending.is_empty()
1849 || !self.inflight_sources.is_empty() || !self.quarantined_sources.is_empty();
1850 if need_copy_drain {
1851 if let Err(err) = self.copy_stream.synchronize() {
1852 safe_to_drop_slots = false;
1853 eprintln!("[moe-cache] unknown copy-stream drain failed ({err}); leaking GPU slots for safety");
1854 for (_, keepalive) in self.inflight_sources.drain(..) {
1855 std::mem::forget(keepalive);
1856 }
1857 for keepalive in self.quarantined_sources.drain(..) {
1858 std::mem::forget(keepalive);
1859 }
1860 for (_, pending) in self.pending.drain() {
1861 if let Some(keepalive) = pending.keepalive {
1862 std::mem::forget(keepalive);
1863 }
1864 }
1865 } else {
1866 self.copy_stream_unknown = false;
1867 self.inflight_sources.clear();
1868 self.quarantined_sources.clear();
1869 self.pending.clear();
1870 }
1871 }
1872 if let Some(pool) = self.pread.as_mut() {
1873 safe_to_drop_slots &= pool.drain();
1874 } else if self.pread_requested && self.pread_fallbacks != 0 {
1875 eprintln!(
1876 "[spill-pread] backend unavailable; mmap_fallbacks={}",
1877 self.pread_fallbacks
1878 );
1879 }
1880 if !safe_to_drop_slots {
1881 for slot in self.slots.drain(..) {
1882 std::mem::forget(slot);
1883 }
1884 }
1885 }
1886}