1use 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
43pub const PROJ_GATE: u8 = 0;
45pub const PROJ_UP: u8 = 1;
46pub const PROJ_DOWN: u8 = 2;
47
48#[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#[derive(Clone, Copy, Debug)]
65pub enum DispatchSlot {
66 Resident(usize),
67}
68
69const NIL: u32 = u32::MAX;
71const SEG_NONE: u8 = 0;
72const SEG_PROBATION: u8 = 1;
73const SEG_PROTECTED: u8 = 2;
74
75#[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#[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 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 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 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 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
192struct SlotClass {
195 capacity: usize,
196 probation: SlruList,
197 protected: SlruList,
198 free: Vec<usize>,
199 protected_cap: usize,
200}
201
202impl SlotClass {
203 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); }
216 _ => self.push_protected(slot, links),
218 }
219 }
220
221 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 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 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
250pub struct MoeSlotCache {
253 slots: Vec<CudaSlice<u8>>, slot_class: Vec<usize>, classes: Vec<SlotClass>,
256 links: Vec<SlotLink>,
259 occupant: Vec<Option<BlockId>>, table: HashMap<BlockId, usize>, frequencies: HashMap<BlockId, f32>,
265 pending: HashMap<BlockId, PendingBlock>,
269 inflight_sources: Vec<(Arc<CudaEvent>, ExpertKeepalive)>,
272 quarantined_sources: Vec<ExpertKeepalive>,
275 compute_sources: HashMap<KeepaliveKey, ExpertKeepalive>,
279 pread: Option<PreadPool>,
282 worker_reads: HashMap<BlockId, WorkerRead>,
285 pread_requested: bool,
286 pread_fallbacks: u64,
287 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 mtp_frequency_weight: f32,
304 last_forward_layer: Option<u16>,
305 last_forward_t: usize,
306 frozen: bool,
310
311 per_layer: HashMap<u16, u32>,
318 dev_rows: HashMap<u16, CudaSlice<u64>>,
323 prewarm_tried: HashSet<u16>,
326
327 pub hits: u64,
329 pub misses: u64,
330 pub staged_bytes: u64, }
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}
343
344#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
345enum KeepaliveKey {
346 Pinned(usize),
347 Buffer(usize),
348 Mmap(usize),
349}
350
351impl KeepaliveKey {
352 fn from_owner(owner: &ExpertKeepalive) -> Self {
353 match owner {
354 ExpertKeepalive::Pinned(value) => Self::Pinned(Arc::as_ptr(value) as usize),
355 ExpertKeepalive::Buffer(value) => Self::Buffer(Arc::as_ptr(value) as usize),
356 ExpertKeepalive::Mmap(value) => Self::Mmap(Arc::as_ptr(value) as usize),
357 }
358 }
359}
360
361struct ExactPinnedPrefix<'a>(&'a [u8]);
366
367impl HostSlice<u8> for ExactPinnedPrefix<'_> {
368 fn len(&self) -> usize {
369 self.0.len()
370 }
371
372 unsafe fn stream_synced_slice<'a>(
373 &'a self,
374 _stream: &'a CudaStream,
375 ) -> (&'a [u8], SyncOnDrop<'a>) {
376 (self.0, SyncOnDrop::Record(None))
379 }
380
381 unsafe fn stream_synced_mut_slice<'a>(
382 &'a mut self,
383 _stream: &'a CudaStream,
384 ) -> (&'a mut [u8], SyncOnDrop<'a>) {
385 panic!("ExactPinnedPrefix is a source-only HostSlice")
386 }
387}
388
389fn stage_on_copy_stream(
390 e: &Engine,
391 host_bytes: &[u8],
392 slot: &mut CudaSlice<u8>,
393) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
394 let prior = match e.stream().record_event(None) {
396 Ok(prior) => prior,
397 Err(err) => return Err((err.into(), true)),
398 };
399 if let Err(err) = e.copy_stream.wait(&prior) {
400 return Err((err.into(), true));
401 }
402 match e.stage_expert_async(host_bytes, slot, 0) {
403 Ok(ready) => Ok(Arc::new(ready)),
404 Err(err) => {
405 match e.copy_stream.synchronize() {
408 Ok(()) => Err((err, true)),
409 Err(sync_err) => Err((std::io::Error::other(format!(
410 "copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
411 )).into(), false)),
412 }
413 }
414 }
415}
416
417fn stage_pread_on_compute_stream(
418 e: &Engine,
419 host_bytes: &[u8],
420 slot: &mut CudaSlice<u8>,
421) -> Result<Arc<CudaEvent>, Box<dyn std::error::Error>> {
422 let ready = Arc::new(e.ctx().new_event(None)?);
423 let source = ExactPinnedPrefix(host_bytes);
424 let mut dst = slot.slice_mut(0..host_bytes.len());
425 e.stream().memcpy_htod(&source, &mut dst)?;
426 ready.record(&e.stream())?;
427 Ok(ready)
428}
429
430fn size_class_plan(block_bytes: &[usize], budget_bytes: usize) -> Vec<(usize, usize)> {
433 let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
434 for &bytes in block_bytes.iter().filter(|&&bytes| bytes > 0) {
435 *counts.entry(bytes).or_insert(0) += 1;
436 }
437 if counts.is_empty() || budget_bytes == 0 {
438 return Vec::new();
439 }
440 let total_bytes: u128 = counts
441 .iter()
442 .map(|(&bytes, &count)| (bytes as u128 + 8) * count as u128)
443 .sum();
444 let budget = budget_bytes as u128;
445 let mut plan: Vec<(usize, usize, u128)> = counts
446 .iter()
447 .map(|(&bytes, &count)| {
448 let scaled = count as u128 * budget;
449 (
450 bytes,
451 (scaled / total_bytes).min(count as u128) as usize,
452 scaled % total_bytes,
453 )
454 })
455 .collect();
456 let mut used: u128 = plan
457 .iter()
458 .map(|(bytes, count, _)| (*bytes as u128 + 8) * *count as u128)
459 .sum();
460
461 let mut order: Vec<usize> = (0..plan.len()).collect();
464 order.sort_by(|&a, &b| plan[b].2.cmp(&plan[a].2).then(a.cmp(&b)));
465 for index in order {
466 let (bytes, count, _) = plan[index];
467 let available = counts[&bytes];
468 let required = bytes as u128 + 8;
469 if count < available && used + required <= budget {
470 plan[index].1 += 1;
471 used += required;
472 }
473 }
474 plan.into_iter()
475 .filter_map(|(bytes, count, _)| (count > 0).then_some((bytes, count)))
476 .collect()
477}
478
479impl MoeSlotCache {
480 pub fn new(e: &Engine, max_block_bytes: usize) -> Result<Self, Box<dyn std::error::Error>> {
487 let (free, _total) = e.ctx().mem_get_info()?;
488 let hard_frac = cache_hard_vram_frac();
491 let hard_bytes =
492 ((free as f64 * hard_frac) as usize).saturating_sub(2 * (max_block_bytes + 8));
493 let forced_slots = std::env::var("MEMRA_MOE_SLOTS")
494 .ok()
495 .and_then(|s| s.parse::<usize>().ok());
496 let requested_bytes = if let Some(n) = forced_slots {
497 n.saturating_mul(max_block_bytes + 8)
498 } else {
499 let frac = std::env::var("MEMRA_MOE_VRAM_FRAC")
505 .ok()
506 .and_then(|s| s.parse::<f64>().ok())
507 .unwrap_or(0.85);
508 (free as f64 * frac) as usize
509 };
510 let budget_bytes = requested_bytes.min(hard_bytes);
511 let layout = e.moe_cache_layout().unwrap_or_default();
512 let size_aware = forced_slots.is_none()
513 && std::env::var("MEMRA_MOE_SIZE_AWARE").as_deref() == Ok("1")
514 && !layout.is_empty();
515 let frequency_evict = std::env::var("MEMRA_MOE_LFU").as_deref() == Ok("1");
516 let frequency_decay = if frequency_evict {
517 cache_lfu_decay()
518 } else {
519 None
520 };
521 let mtp_frequency_weight = cache_lfu_mtp_weight();
522 let mut class_plan = if size_aware {
523 size_class_plan(&layout, budget_bytes)
524 } else {
525 Vec::new()
526 };
527 if class_plan.iter().map(|(_, count)| count).sum::<usize>() < 8 {
528 let n = (budget_bytes / (max_block_bytes + 8)).max(8);
529 class_plan = vec![(max_block_bytes, n)];
530 }
531 let n: usize = class_plan.iter().map(|(_, count)| count).sum();
532
533 let mut slots = Vec::with_capacity(n);
534 let mut slot_class = Vec::with_capacity(n);
535 let mut classes = Vec::with_capacity(class_plan.len());
536 let mut occupant = Vec::with_capacity(n);
537 for (class_index, &(capacity, count)) in class_plan.iter().enumerate() {
538 let start = slots.len();
539 for _ in 0..count {
540 slots.push(e.alloc_u8(capacity + 8)?);
542 slot_class.push(class_index);
543 occupant.push(None);
544 }
545 let free_slots = (start..start + count).rev().collect();
546 classes.push(SlotClass {
547 capacity,
548 probation: SlruList::new(),
549 protected: SlruList::new(),
550 free: free_slots,
551 protected_cap: ((count as f64 * 0.8) as usize).max(1),
552 });
553 }
554 let links = vec![SlotLink::none(); n];
555 if size_aware {
556 let allocated: usize = class_plan
557 .iter()
558 .map(|(bytes, count)| (bytes + 8) * count)
559 .sum();
560 eprintln!(
561 "[moe-cache] size-aware fixed slots: {n} slots in {} classes, {:.2} GB / {:.2} GB budget",
562 class_plan.len(),
563 allocated as f64 / 1e9,
564 budget_bytes as f64 / 1e9
565 );
566 }
567 let pread_mode = crate::spill_pread::configured_mode();
568 let pread_requested = pread_mode != SpillIoMode::Mmap;
569 let pread = if pread_requested {
570 match PreadPool::try_new(e, max_block_bytes, pread_mode) {
571 Ok(pool) => Some(pool),
572 Err(err) => {
573 eprintln!(
574 "[spill-pread] pinned-buffer initialization failed ({err}); using mmap"
575 );
576 None
577 }
578 }
579 } else {
580 None
581 };
582
583 Ok(MoeSlotCache {
584 slots,
585 slot_class,
586 classes,
587 links,
588 occupant,
589 table: HashMap::with_capacity(n * 2),
590 frequencies: HashMap::with_capacity(layout.len().max(n * 2)),
591 pending: HashMap::new(),
592 inflight_sources: Vec::new(),
593 quarantined_sources: Vec::new(),
594 compute_sources: HashMap::new(),
595 pread,
596 worker_reads: HashMap::new(),
597 pread_requested,
598 pread_fallbacks: 0,
599 copy_stream: e.copy_stream.clone(),
600 copy_stream_unknown: false,
601 compute_stream: e.stream().clone(),
602 compute_stream_unknown: false,
603 n,
604 max_block_bytes,
605 size_aware,
606 frequency_evict,
607 frequency_decay,
608 mtp_frequency_weight,
609 last_forward_layer: None,
610 last_forward_t: 0,
611 frozen: false,
612 per_layer: HashMap::new(),
613 dev_rows: HashMap::new(),
614 prewarm_tried: HashSet::new(),
615 hits: 0,
616 misses: 0,
617 staged_bytes: 0,
618 })
619 }
620
621 #[inline]
622 pub fn n_slots(&self) -> usize {
623 self.n
624 }
625 #[inline]
626 pub fn is_frozen(&self) -> bool {
627 self.frozen
628 }
629 pub fn freeze(&mut self) {
630 if !self.frozen {
631 self.frozen = true;
632 let (_, complete, one_projection, two_projections, stranded_blocks) =
633 self.expert_residency_shape();
634 eprintln!(
635 "[moe-cache] residency frozen: {} slots, {} resident blocks; \
636 {complete} complete experts, {one_projection} one-projection fragments, \
637 {two_projections} two-projection fragments ({stranded_blocks} stranded blocks)",
638 self.n,
639 self.table.len()
640 );
641 let mut mtp_masks = HashMap::<u16, u8>::new();
642 for id in self.table.keys().filter(|id| id.layer == u16::MAX) {
643 *mtp_masks.entry(id.ex).or_insert(0) |= 1u8 << id.proj;
644 }
645 if !mtp_masks.is_empty() {
646 let complete = mtp_masks.values().filter(|&&mask| mask == 0b111).count();
647 eprintln!(
648 "[moe-cache] frozen MTP residency: {} blocks, {complete} complete experts",
649 mtp_masks
650 .values()
651 .map(|mask| mask.count_ones() as usize)
652 .sum::<usize>()
653 );
654 }
655 }
656 }
657
658 pub(crate) fn expert_residency_shape(&self) -> (usize, usize, usize, usize, usize) {
659 let mut masks = HashMap::<(u16, u16), u8>::new();
660 for id in self.table.keys() {
661 *masks.entry((id.layer, id.ex)).or_insert(0) |= 1u8 << id.proj;
662 }
663 let complete = masks.values().filter(|&&mask| mask == 0b111).count();
664 let one_projection = masks
665 .values()
666 .filter(|&&mask| mask.count_ones() == 1)
667 .count();
668 let two_projections = masks
669 .values()
670 .filter(|&&mask| mask.count_ones() == 2)
671 .count();
672 let stranded_blocks = one_projection + 2 * two_projections;
673 (
674 masks.len(),
675 complete,
676 one_projection,
677 two_projections,
678 stranded_blocks,
679 )
680 }
681
682 #[inline]
683 pub fn max_block_bytes(&self) -> usize {
684 self.max_block_bytes
685 }
686
687 #[inline]
689 pub fn resident(&self, id: BlockId) -> Option<usize> {
690 self.table.get(&id).copied()
691 }
692
693 #[inline]
694 fn frequency_increment(&self, id: BlockId) -> f32 {
695 if id.layer == u16::MAX {
696 self.mtp_frequency_weight
697 } else {
698 1.0
699 }
700 }
701
702 pub(crate) fn note_profile_hit(&mut self, id: BlockId) {
706 if self.frozen || !self.table.contains_key(&id) {
707 return;
708 }
709 let increment = self.frequency_increment(id);
710 *self.frequencies.entry(id).or_insert(0.0) += increment;
711 }
712
713 fn on_hit(&mut self, slot: usize) {
731 let class_index = self.slot_class[slot];
732 let class = &mut self.classes[class_index];
733 if !class.free.is_empty() {
734 return;
735 }
736 class.on_hit_full(slot, &mut self.links);
737 }
738
739 fn remove_occupant(&mut self, slot: usize) {
740 if let Some(old) = self.occupant[slot].take() {
741 self.table.remove(&old);
742 self.on_block_evicted(old.layer);
743 }
744 }
745
746 fn frequency_victim_in_class(&mut self, class_index: usize, keep: &[BlockId]) -> Option<usize> {
753 let class = &self.classes[class_index];
754 let candidate = class
755 .probation
756 .iter(&self.links)
757 .chain(class.protected.iter(&self.links))
758 .enumerate()
759 .filter_map(|(position, slot)| {
760 let id = self.occupant[slot]?;
761 (!keep.contains(&id)).then_some((
762 self.frequencies.get(&id).copied().unwrap_or(0.0),
763 position,
764 slot,
765 ))
766 })
767 .min_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
768 let (_, _, slot) = candidate?;
769 self.classes[class_index].unlink_from_segment(slot, &mut self.links);
770 Some(slot)
771 }
772
773 fn evict_one(&mut self, required: usize) -> Option<usize> {
775 for class_index in 0..self.classes.len() {
776 if self.classes[class_index].capacity < required {
777 continue;
778 }
779 let slot = if self.frequency_evict {
780 self.frequency_victim_in_class(class_index, &[])
781 } else {
782 self.classes[class_index].pop_lru(&mut self.links)
783 };
784 if let Some(slot) = slot {
785 self.remove_occupant(slot);
786 return Some(slot);
787 }
788 }
789 None
790 }
791
792 fn evict_one_excluding(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
798 fn take(
799 q: &mut SlruList,
800 links: &mut [SlotLink],
801 occupant: &[Option<BlockId>],
802 keep: &[BlockId],
803 ) -> Option<usize> {
804 let slot = q
805 .iter(links)
806 .find(|&s| occupant[s].is_some_and(|id| !keep.contains(&id)))?;
807 q.unlink(slot, links);
808 Some(slot)
809 }
810 for class_index in 0..self.classes.len() {
811 if self.classes[class_index].capacity < required {
812 continue;
813 }
814 let slot = if self.frequency_evict {
815 self.frequency_victim_in_class(class_index, keep)
816 } else {
817 let class = &mut self.classes[class_index];
818 take(&mut class.probation, &mut self.links, &self.occupant, keep)
819 .or_else(|| take(&mut class.protected, &mut self.links, &self.occupant, keep))
820 };
821 if let Some(slot) = slot {
822 self.remove_occupant(slot);
823 return Some(slot);
824 }
825 }
826 None
827 }
828
829 fn on_block_evicted(&mut self, layer: u16) {
835 if let Some(c) = self.per_layer.get_mut(&layer) {
836 *c -= 1;
837 }
838 self.dev_rows.remove(&layer);
839 }
840
841 fn reserve_slot(&mut self, required: usize) -> Option<usize> {
842 for class in &mut self.classes {
843 if class.capacity >= required
844 && let Some(slot) = class.free.pop()
845 {
846 return Some(slot);
847 }
848 }
849 self.evict_one(required)
850 }
851
852 fn release_reserved_slot(&mut self, slot: usize) {
853 debug_assert!(self.occupant[slot].is_none());
854 self.classes[self.slot_class[slot]].free.push(slot);
855 }
856
857 fn publish(&mut self, id: BlockId, slot: usize) {
858 self.occupant[slot] = Some(id);
859 self.table.insert(id, slot);
860 self.classes[self.slot_class[slot]].probation.push_back(
861 slot,
862 SEG_PROBATION,
863 &mut self.links,
864 );
865 *self.per_layer.entry(id.layer).or_insert(0) += 1;
866 }
867
868 fn reap_copy_sources(&mut self) {
869 self.inflight_sources
870 .retain(|(ready, _)| !ready.is_complete());
871 }
872
873 fn retain_compute_source(&mut self, owner: Option<ExpertKeepalive>) {
874 if let Some(owner) = owner {
875 let key = KeepaliveKey::from_owner(&owner);
876 self.compute_sources.entry(key).or_insert(owner);
877 }
878 }
879
880 fn admit(
883 &mut self,
884 id: BlockId,
885 host_bytes: &[u8],
886 e: &Engine,
887 ) -> Result<usize, Box<dyn std::error::Error>> {
888 let slot = self.reserve_slot(host_bytes.len()).ok_or_else(|| {
889 std::io::Error::other(format!(
890 "no MoE cache slot can hold {} bytes (max class {})",
891 host_bytes.len(),
892 self.classes.last().map(|class| class.capacity).unwrap_or(0)
893 ))
894 })?;
895 if let Err(err) = e.stage_expert(host_bytes, &mut self.slots[slot], 0) {
898 return match e.stream().synchronize() {
899 Ok(()) => {
900 self.release_reserved_slot(slot);
901 Err(err)
902 }
903 Err(sync_err) => {
904 self.compute_stream_unknown = true;
907 Err(std::io::Error::other(format!(
908 "compute-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
909 )).into())
910 }
911 };
912 }
913 self.staged_bytes += host_bytes.len() as u64;
914 self.publish(id, slot);
915 Ok(slot)
916 }
917
918 fn note_pread_fallback(&mut self, reason: &dyn std::fmt::Display) {
919 self.pread_fallbacks += 1;
920 if let Some(pool) = self.pread.as_mut() {
921 pool.note_fallback();
922 }
923 if self.pread_fallbacks <= 3 {
924 eprintln!("[spill-pread] falling back to mmap: {reason}");
925 }
926 }
927
928 pub(crate) fn begin_worker_scope(&mut self) {
932 if self.worker_reads.is_empty() {
933 return;
934 }
935 let tickets: Vec<_> = self
936 .worker_reads
937 .drain()
938 .map(|(_, read)| read.ticket)
939 .collect();
940 if let Some(pool) = self.pread.as_mut().filter(|pool| pool.is_worker()) {
941 for ticket in tickets {
942 let _ = pool.cancel_worker(ticket);
943 }
944 }
945 }
946
947 pub(crate) fn begin_forward_epoch(&mut self, layer: u16, t: usize) {
957 let Some(decay) = self.frequency_decay else {
958 self.last_forward_layer = Some(layer);
959 self.last_forward_t = t;
960 return;
961 };
962 let new_sweep = self
963 .last_forward_layer
964 .is_some_and(|previous| layer <= previous);
965 if new_sweep && t == 1 {
966 if self.last_forward_t != 1 {
967 self.frequencies.clear();
968 } else {
969 self.frequencies.retain(|_, score| {
970 *score *= decay;
971 *score >= 1.0e-3
972 });
973 }
974 }
975 self.last_forward_layer = Some(layer);
976 self.last_forward_t = t;
977 }
978
979 fn dispatch_disk(
980 &mut self,
981 id: BlockId,
982 file: &Arc<std::fs::File>,
983 offset: u64,
984 len: usize,
985 fallback: &[u8],
986 e: &Engine,
987 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
988 if self.pread.is_none() {
989 if self.pread_requested {
990 self.note_pread_fallback(&"pinned-buffer backend unavailable");
991 }
992 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
993 }
994
995 let pending = self.worker_reads.remove(&id);
996 let pool = self.pread.as_mut().unwrap();
997 let read = if pool.is_worker() {
998 let ticket = match pending {
999 Some(read) => Ok(Some(read.ticket)),
1000 None => pool.submit_worker(file.clone(), offset, len),
1001 };
1002 match ticket {
1003 Ok(Some(ticket)) => match pool.wait_worker(ticket) {
1004 Ok(index) => Ok(index),
1005 Err(err) => {
1006 let _ = pool.cancel_worker(ticket);
1009 Err(err)
1010 }
1011 },
1012 Ok(None) => Err(std::io::Error::other("worker read ring is busy").into()),
1013 Err(err) => Err(err),
1014 }
1015 } else {
1016 debug_assert!(pending.is_none());
1017 pool.read(file.as_ref(), offset, len)
1018 };
1019 let index = match read {
1020 Ok(index) => index,
1021 Err(err) => {
1022 self.note_pread_fallback(err.as_ref());
1023 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1024 }
1025 };
1026
1027 let slot = self.reserve_slot(len).ok_or_else(|| {
1030 std::io::Error::other(format!(
1031 "no MoE cache slot can hold {len} bytes (max class {})",
1032 self.classes.last().map(|class| class.capacity).unwrap_or(0)
1033 ))
1034 })?;
1035 let ready = {
1036 let bytes = match self.pread.as_ref().unwrap().bytes(index, len) {
1037 Ok(bytes) => bytes,
1038 Err(err) => {
1039 self.pread.as_mut().unwrap().abort_read(index);
1040 self.release_reserved_slot(slot);
1041 self.note_pread_fallback(err.as_ref());
1042 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1043 }
1044 };
1045 stage_pread_on_compute_stream(e, bytes, &mut self.slots[slot])
1046 };
1047 let ready = match ready {
1048 Ok(ready) => ready,
1049 Err(err) => {
1050 match e.stream().synchronize() {
1054 Ok(()) => {
1055 self.pread.as_mut().unwrap().abort_read(index);
1056 self.release_reserved_slot(slot);
1057 self.note_pread_fallback(err.as_ref());
1058 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1059 }
1060 Err(sync_err) => {
1061 self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1062 return Err(std::io::Error::other(format!(
1063 "pread H2D setup failed ({err}); CUDA stream drain also failed ({sync_err})"
1064 )).into());
1065 }
1066 }
1067 }
1068 };
1069 self.pread.as_mut().unwrap().mark_h2d(index, ready);
1070 self.staged_bytes += len as u64;
1073 self.publish(id, slot);
1074 Ok(DispatchSlot::Resident(slot))
1075 }
1076
1077 pub fn dispatch(
1085 &mut self,
1086 id: BlockId,
1087 host_bytes: &[u8],
1088 e: &Engine,
1089 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1090 self.dispatch_source(
1091 id,
1092 ExpertSource::Memory {
1093 bytes: host_bytes,
1094 keepalive: None,
1095 },
1096 e,
1097 )
1098 }
1099
1100 pub(crate) fn dispatch_source(
1101 &mut self,
1102 id: BlockId,
1103 source: ExpertSource<'_>,
1104 e: &Engine,
1105 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1106 self.reap_copy_sources();
1107 let increment = self.frequency_increment(id);
1108 *self.frequencies.entry(id).or_insert(0.0) += increment;
1109 if let Some(s) = self.table.get(&id).copied() {
1110 self.hits += 1;
1111 self.on_hit(s);
1112 return Ok(DispatchSlot::Resident(s));
1113 }
1114 if let Some(pending) = self.pending.remove(&id) {
1115 if let Err(err) = e.compute_wait(pending.ready.as_ref()) {
1116 self.pending.insert(id, pending);
1117 return Err(err);
1118 }
1119 self.misses += 1;
1120 let slot = pending.slot;
1121 if let Some(keepalive) = pending.keepalive {
1122 self.inflight_sources.push((pending.ready, keepalive));
1123 }
1124 self.publish(id, slot);
1125 return Ok(DispatchSlot::Resident(slot));
1126 }
1127 self.misses += 1;
1128 match source {
1138 ExpertSource::Memory { bytes, keepalive } => {
1139 self.retain_compute_source(keepalive);
1142 let slot = self.admit(id, bytes, e)?;
1143 Ok(DispatchSlot::Resident(slot))
1144 }
1145 ExpertSource::Disk {
1146 file,
1147 offset,
1148 len,
1149 fallback,
1150 keepalive,
1151 } => {
1152 self.retain_compute_source(Some(keepalive));
1155 self.dispatch_disk(id, file, offset, len, fallback, e)
1156 }
1157 }
1158 }
1159
1160 pub fn prefetch(
1169 &mut self,
1170 id: BlockId,
1171 host_bytes: &[u8],
1172 keep: &[BlockId],
1173 e: &Engine,
1174 ) -> Result<bool, Box<dyn std::error::Error>> {
1175 self.prefetch_source(
1176 id,
1177 ExpertSource::Memory {
1178 bytes: host_bytes,
1179 keepalive: None,
1180 },
1181 keep,
1182 e,
1183 )
1184 }
1185
1186 fn reserve_prefetch_slot(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
1187 for class in &mut self.classes {
1188 if class.capacity >= required
1189 && let Some(slot) = class.free.pop()
1190 {
1191 return Some(slot);
1192 }
1193 }
1194 self.evict_one_excluding(required, keep)
1195 }
1196
1197 fn prefetch_bytes(
1198 &mut self,
1199 id: BlockId,
1200 host_bytes: &[u8],
1201 keepalive: Option<ExpertKeepalive>,
1202 keep: &[BlockId],
1203 e: &Engine,
1204 ) -> Result<bool, Box<dyn std::error::Error>> {
1205 let Some(slot) = self.reserve_prefetch_slot(host_bytes.len(), keep) else {
1206 return Ok(false);
1207 };
1208 let ready = match stage_on_copy_stream(e, host_bytes, &mut self.slots[slot]) {
1209 Ok(ready) => ready,
1210 Err((err, reusable)) => {
1211 if reusable {
1212 self.release_reserved_slot(slot);
1213 } else {
1214 self.copy_stream_unknown = true;
1217 if let Some(keepalive) = keepalive {
1218 self.quarantined_sources.push(keepalive);
1219 }
1220 eprintln!(
1221 "[moe-cache] quarantining slot {slot} after unprovable copy completion"
1222 );
1223 }
1224 return Err(err);
1225 }
1226 };
1227 self.occupant[slot] = Some(id);
1228 self.pending.insert(
1229 id,
1230 PendingBlock {
1231 slot,
1232 ready,
1233 keepalive,
1234 },
1235 );
1236 self.staged_bytes += host_bytes.len() as u64;
1237 Ok(true)
1238 }
1239
1240 pub(crate) fn prefetch_source(
1241 &mut self,
1242 id: BlockId,
1243 source: ExpertSource<'_>,
1244 keep: &[BlockId],
1245 e: &Engine,
1246 ) -> Result<bool, Box<dyn std::error::Error>> {
1247 self.reap_copy_sources();
1248 if self.table.contains_key(&id)
1249 || self.pending.contains_key(&id)
1250 || self.worker_reads.contains_key(&id)
1251 {
1252 return Ok(false);
1253 }
1254 match source {
1255 ExpertSource::Memory { bytes, keepalive } => {
1256 self.prefetch_bytes(id, bytes, keepalive, keep, e)
1257 }
1258 ExpertSource::Disk {
1259 file,
1260 offset,
1261 len,
1262 fallback,
1263 keepalive,
1264 } => {
1265 if self.pread.as_ref().is_some_and(PreadPool::is_worker) {
1266 match self.pread.as_mut().unwrap().submit_worker_speculative(
1267 file.clone(),
1268 offset,
1269 len,
1270 ) {
1271 Ok(Some(ticket)) => {
1272 self.worker_reads.insert(id, WorkerRead { ticket });
1273 Ok(true)
1274 }
1275 Ok(None) => Ok(false),
1276 Err(err) => {
1277 self.note_pread_fallback(err.as_ref());
1278 Ok(false)
1279 }
1280 }
1281 } else if self.pread.is_some() {
1282 Ok(false)
1284 } else {
1285 self.prefetch_bytes(id, fallback, Some(keepalive), keep, e)
1286 }
1287 }
1288 }
1289 }
1290
1291 pub fn force_admit(
1293 &mut self,
1294 id: BlockId,
1295 host_bytes: &[u8],
1296 e: &Engine,
1297 ) -> Result<usize, Box<dyn std::error::Error>> {
1298 if let Some(s) = self.table.get(&id).copied() {
1299 return Ok(s);
1300 }
1301 self.admit(id, host_bytes, e)
1302 }
1303
1304 pub fn export_residency(&self) -> Vec<(u16, u8, u16)> {
1312 self.occupant
1313 .iter()
1314 .flatten()
1315 .map(|id| (id.layer, id.proj, id.ex))
1316 .collect()
1317 }
1318
1319 pub fn restage_block(
1324 &mut self,
1325 id: BlockId,
1326 m: &crate::hybrid::MoeWeights,
1327 e: &Engine,
1328 ) -> Result<bool, Box<dyn std::error::Error>> {
1329 if self.table.contains_key(&id) {
1330 return Ok(true);
1331 }
1332 let exps = match id.proj {
1333 PROJ_GATE => &m.gate_exps,
1334 PROJ_UP => &m.up_exps,
1335 PROJ_DOWN => &m.down_exps,
1336 _ => return Ok(false),
1337 };
1338 if id.ex as usize >= exps.n_expert {
1339 return Ok(false);
1340 }
1341 if m.active_experts
1342 .as_ref()
1343 .is_some_and(|active| !active[id.ex as usize])
1344 {
1345 return Ok(false);
1346 }
1347 if exps.expert_layout(id.ex as usize).len == 0 {
1348 return Ok(false);
1349 }
1350 match exps.expert_source(id.ex as usize) {
1351 ExpertSource::Memory { bytes, keepalive } => {
1352 self.retain_compute_source(keepalive);
1353 self.admit(id, bytes, e)?;
1354 }
1355 ExpertSource::Disk {
1356 fallback,
1357 keepalive,
1358 ..
1359 } => {
1360 self.retain_compute_source(Some(keepalive));
1361 self.admit(id, fallback, e)?;
1362 }
1363 }
1364 Ok(true)
1365 }
1366
1367 pub fn prewarm_layer(
1368 &mut self,
1369 layer: u16,
1370 m: &crate::hybrid::MoeWeights,
1371 e: &Engine,
1372 ) -> Result<(), Box<dyn std::error::Error>> {
1373 if !self.prewarm_tried.insert(layer) {
1374 return Ok(());
1375 }
1376 let n_expert = m.gate_exps.n_expert;
1377 if self.pread.is_some()
1378 && (0..n_expert).any(|ex| {
1379 matches!(m.gate_exps.expert_source(ex), ExpertSource::Disk { .. })
1380 || matches!(m.up_exps.expert_source(ex), ExpertSource::Disk { .. })
1381 || matches!(m.down_exps.expert_source(ex), ExpertSource::Disk { .. })
1382 })
1383 {
1384 return Ok(());
1387 }
1388 let resident = self.per_layer.get(&layer).copied().unwrap_or(0) as usize;
1389 let missing = 3 * n_expert - resident;
1390 if self.size_aware {
1391 return Ok(());
1392 } if self
1394 .classes
1395 .iter()
1396 .map(|class| class.free.len())
1397 .sum::<usize>()
1398 < missing
1399 {
1400 return Ok(()); }
1402 for ex in 0..n_expert {
1403 for (proj, exps) in [
1404 (PROJ_GATE, &m.gate_exps),
1405 (PROJ_UP, &m.up_exps),
1406 (PROJ_DOWN, &m.down_exps),
1407 ] {
1408 let id = BlockId::new(layer, proj, ex as u16);
1409 if self.table.contains_key(&id) {
1410 continue;
1411 }
1412 match exps.expert_source(ex) {
1413 ExpertSource::Memory { bytes, keepalive } => {
1414 self.retain_compute_source(keepalive);
1415 self.admit(id, bytes, e)?;
1416 }
1417 ExpertSource::Disk {
1418 fallback,
1419 keepalive,
1420 ..
1421 } => {
1422 self.retain_compute_source(Some(keepalive));
1423 self.admit(id, fallback, e)?;
1424 }
1425 }
1426 }
1427 }
1428 Ok(())
1429 }
1430
1431 pub fn layer_dev_row(
1437 &mut self,
1438 layer: u16,
1439 n_expert: usize,
1440 e: &Engine,
1441 ) -> Result<Option<&CudaSlice<u64>>, Box<dyn std::error::Error>> {
1442 if self.per_layer.get(&layer).copied().unwrap_or(0) as usize != 3 * n_expert {
1443 return Ok(None);
1444 }
1445 if !self.dev_rows.contains_key(&layer) {
1446 use cudarc::driver::DevicePtr;
1447 let mut host = vec![0u64; 3 * n_expert];
1448 for proj in 0..3u8 {
1449 for ex in 0..n_expert {
1450 let Some(&s) = self.table.get(&BlockId::new(layer, proj, ex as u16)) else {
1451 return Ok(None);
1453 };
1454 let __s_ev = e.stream();
1455 let (p, _ev) = self.slots[s].device_ptr(&__s_ev);
1456 host[proj as usize * n_expert + ex] = p;
1457 }
1458 }
1459 let row = e.stream().clone_htod(&host)?;
1460 self.dev_rows.insert(layer, row);
1461 }
1462 Ok(self.dev_rows.get(&layer))
1463 }
1464
1465 #[inline]
1467 pub fn buf(&self, d: DispatchSlot) -> &CudaSlice<u8> {
1468 match d {
1469 DispatchSlot::Resident(s) => &self.slots[s],
1470 }
1471 }
1472
1473 #[inline]
1475 pub fn slot(&self, s: usize) -> &CudaSlice<u8> {
1476 &self.slots[s]
1477 }
1478
1479 pub fn hit_rate(&self) -> f64 {
1481 let tot = self.hits + self.misses;
1482 if tot == 0 {
1483 0.0
1484 } else {
1485 self.hits as f64 / tot as f64
1486 }
1487 }
1488
1489 pub fn reset_counters(&mut self) {
1491 self.hits = 0;
1492 self.misses = 0;
1493 self.staged_bytes = 0;
1494 }
1495
1496 pub(crate) fn pread_stats(&self) -> Option<PreadStats> {
1497 if !self.pread_requested {
1498 return None;
1499 }
1500 let mut stats = self
1501 .pread
1502 .as_ref()
1503 .map(PreadPool::stats)
1504 .unwrap_or_default();
1505 stats.fallbacks = self.pread_fallbacks;
1506 Some(stats)
1507 }
1508}
1509
1510fn cache_lfu_decay() -> Option<f32> {
1511 let raw = std::env::var("MEMRA_MOE_LFU_DECAY").ok()?;
1512 match parse_cache_lfu_decay(Some(&raw)) {
1513 Ok(value) => value,
1514 Err(reason) => {
1515 eprintln!(
1516 "[moe-cache] invalid MEMRA_MOE_LFU_DECAY={raw:?} ({reason}); disabling LFU decay"
1517 );
1518 None
1519 }
1520 }
1521}
1522
1523fn cache_lfu_mtp_weight() -> f32 {
1524 const DEFAULT: f32 = 1.0;
1525 let raw = std::env::var("MEMRA_MOE_LFU_MTP_WEIGHT").ok();
1526 match parse_cache_lfu_mtp_weight(raw.as_deref()) {
1527 Ok(value) => value,
1528 Err(reason) => {
1529 eprintln!(
1530 "[moe-cache] invalid MEMRA_MOE_LFU_MTP_WEIGHT={:?} ({reason}); using {DEFAULT}",
1531 raw.as_deref().unwrap_or("")
1532 );
1533 DEFAULT
1534 }
1535 }
1536}
1537
1538fn parse_cache_lfu_mtp_weight(raw: Option<&str>) -> Result<f32, &'static str> {
1539 let value = raw
1540 .unwrap_or("1")
1541 .parse::<f32>()
1542 .map_err(|_| "expected a number")?;
1543 if value.is_finite() && (0.25..=64.0).contains(&value) {
1544 Ok(value)
1545 } else {
1546 Err("expected a finite multiplier from 0.25 through 64")
1547 }
1548}
1549
1550fn parse_cache_lfu_decay(raw: Option<&str>) -> Result<Option<f32>, &'static str> {
1551 let Some(raw) = raw else { return Ok(None) };
1552 let value = raw.parse::<f32>().map_err(|_| "expected a number")?;
1553 if value.is_finite() && value > 0.0 && value <= 1.0 {
1554 Ok(Some(value))
1555 } else {
1556 Err("expected a finite fraction greater than 0 and at most 1")
1557 }
1558}
1559
1560fn cache_hard_vram_frac() -> f64 {
1561 const DEFAULT: f64 = 0.80;
1562 let raw = std::env::var("MEMRA_MOE_HARD_VRAM_FRAC").ok();
1563 match parse_cache_hard_vram_frac(raw.as_deref()) {
1564 Ok(value) => value,
1565 Err(reason) => {
1566 eprintln!(
1567 "[moe-cache] invalid MEMRA_MOE_HARD_VRAM_FRAC={:?} ({reason}); using {DEFAULT}",
1568 raw.as_deref().unwrap_or("")
1569 );
1570 DEFAULT
1571 }
1572 }
1573}
1574
1575fn parse_cache_hard_vram_frac(raw: Option<&str>) -> Result<f64, &'static str> {
1576 let value = raw
1577 .unwrap_or("0.80")
1578 .parse::<f64>()
1579 .map_err(|_| "expected a number")?;
1580 if value.is_finite() && (0.10..=0.95).contains(&value) {
1581 Ok(value)
1582 } else {
1583 Err("expected a finite fraction from 0.10 through 0.95")
1584 }
1585}
1586
1587#[cfg(test)]
1588mod slru_intrusive_tests {
1589 use super::{SEG_PROBATION, SEG_PROTECTED, SlotClass, SlotLink, SlruList};
1596 use std::collections::VecDeque;
1597
1598 struct OldSlru {
1600 probation: VecDeque<usize>,
1601 protected: VecDeque<usize>,
1602 protected_cap: usize,
1603 }
1604 impl OldSlru {
1605 fn on_hit_full(&mut self, slot: usize) {
1606 if let Some(pos) = self.probation.iter().position(|&x| x == slot) {
1607 self.probation.remove(pos);
1608 self.push_protected(slot);
1609 } else if let Some(pos) = self.protected.iter().position(|&x| x == slot) {
1610 self.protected.remove(pos);
1611 self.protected.push_back(slot); } else {
1613 self.push_protected(slot);
1614 }
1615 }
1616 fn push_protected(&mut self, slot: usize) {
1617 self.protected.push_back(slot);
1618 while self.protected.len() > self.protected_cap {
1619 if let Some(demoted) = self.protected.pop_front() {
1620 self.probation.push_back(demoted);
1621 } else {
1622 break;
1623 }
1624 }
1625 }
1626 fn pop_lru(&mut self) -> Option<usize> {
1627 self.probation
1628 .pop_front()
1629 .or_else(|| self.protected.pop_front())
1630 }
1631 fn take_excluding(&mut self, banned: &[usize]) -> Option<usize> {
1632 let take = |q: &mut VecDeque<usize>| {
1633 q.iter()
1634 .position(|&s| !banned.contains(&s))
1635 .and_then(|pos| q.remove(pos))
1636 };
1637 take(&mut self.probation).or_else(|| take(&mut self.protected))
1638 }
1639 }
1640
1641 fn new_pair(n: usize, protected_cap: usize) -> (SlotClass, Vec<SlotLink>, OldSlru) {
1642 let class = SlotClass {
1643 capacity: 1,
1644 probation: SlruList::new(),
1645 protected: SlruList::new(),
1646 free: Vec::new(),
1647 protected_cap,
1648 };
1649 let links = vec![SlotLink::none(); n];
1650 let old = OldSlru {
1651 probation: VecDeque::new(),
1652 protected: VecDeque::new(),
1653 protected_cap,
1654 };
1655 (class, links, old)
1656 }
1657
1658 fn orders_match(class: &SlotClass, links: &[SlotLink], old: &OldSlru) -> bool {
1659 let np: Vec<usize> = class.probation.iter(links).collect();
1660 let nt: Vec<usize> = class.protected.iter(links).collect();
1661 let op: Vec<usize> = old.probation.iter().copied().collect();
1662 let ot: Vec<usize> = old.protected.iter().copied().collect();
1663 np == op && nt == ot
1664 }
1665
1666 struct Rng(u64);
1668 impl Rng {
1669 fn next(&mut self) -> u64 {
1670 self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
1671 let mut z = self.0;
1672 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1673 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1674 z ^ (z >> 31)
1675 }
1676 fn below(&mut self, n: usize) -> usize {
1677 (self.next() % n as u64) as usize
1678 }
1679 }
1680
1681 #[test]
1682 fn same_eviction_decisions_randomized_soak() {
1683 for &(n, cap) in &[(8usize, 1usize), (16, 12), (64, 51), (128, 102)] {
1685 let (mut class, mut links, mut old) = new_pair(n, cap);
1686 let mut rng = Rng(0xC0FFEE ^ (n as u64) << 8 ^ cap as u64);
1687 let mut resident: Vec<usize> = Vec::new();
1688 let mut free: Vec<usize> = (0..n).rev().collect();
1689 for step in 0..200_000 {
1690 let op = rng.below(100);
1691 if op < 55 && !resident.is_empty() {
1692 let slot = resident[rng.below(resident.len())];
1694 class.on_hit_full(slot, &mut links);
1695 old.on_hit_full(slot);
1696 } else if op < 80 {
1697 let slot = if let Some(s) = free.pop() {
1699 s
1700 } else {
1701 let v_new = class.pop_lru(&mut links);
1702 let v_old = old.pop_lru();
1703 assert_eq!(
1704 v_new, v_old,
1705 "victim diverged at step {step} (n={n} cap={cap})"
1706 );
1707 let v = v_new.unwrap();
1708 resident.retain(|&s| s != v);
1709 v
1710 };
1711 class.probation.push_back(slot, SEG_PROBATION, &mut links);
1712 old.probation.push_back(slot);
1713 resident.push(slot);
1714 } else if op < 92 && resident.len() > 2 {
1715 let banned: Vec<usize> = (0..3.min(resident.len()))
1717 .map(|_| resident[rng.below(resident.len())])
1718 .collect();
1719 let take_new = {
1720 let q = &mut class.probation;
1721 let found = q.iter(&links).find(|s| !banned.contains(s));
1722 match found {
1723 Some(s) => {
1724 q.unlink(s, &mut links);
1725 Some(s)
1726 }
1727 None => {
1728 let q = &mut class.protected;
1729 q.iter(&links).find(|s| !banned.contains(s)).inspect(|&s| {
1730 q.unlink(s, &mut links);
1731 })
1732 }
1733 }
1734 };
1735 let take_old = old.take_excluding(&banned);
1736 assert_eq!(
1737 take_new, take_old,
1738 "excluding-victim diverged at step {step}"
1739 );
1740 if let Some(v) = take_new {
1741 resident.retain(|&s| s != v);
1742 free.push(v);
1743 }
1744 } else if !resident.is_empty() {
1745 let slot = resident[rng.below(resident.len())];
1747 match links[slot].seg {
1748 SEG_PROBATION => class.probation.unlink(slot, &mut links),
1749 SEG_PROTECTED => class.protected.unlink(slot, &mut links),
1750 _ => {}
1751 }
1752 if let Some(pos) = old.probation.iter().position(|&x| x == slot) {
1753 old.probation.remove(pos);
1754 } else if let Some(pos) = old.protected.iter().position(|&x| x == slot) {
1755 old.protected.remove(pos);
1756 }
1757 class.on_hit_full(slot, &mut links);
1758 old.on_hit_full(slot);
1759 }
1760 assert!(
1761 orders_match(&class, &links, &old),
1762 "segment order diverged at step {step} (n={n} cap={cap})"
1763 );
1764 }
1765 }
1766 }
1767
1768 #[test]
1769 fn hit_promotion_is_o1_not_on() {
1770 fn bench(n: usize, hits: usize) -> std::time::Duration {
1775 let (mut class, mut links, _) = new_pair(n, (n as f64 * 0.8) as usize);
1776 for s in 0..n {
1777 class.probation.push_back(s, SEG_PROBATION, &mut links);
1778 }
1779 let mut rng = Rng(0xBEEF);
1780 let t0 = std::time::Instant::now();
1781 for _ in 0..hits {
1782 class.on_hit_full(rng.below(n), &mut links);
1783 }
1784 t0.elapsed()
1785 }
1786 bench(1_000, 10_000);
1788 bench(46_000, 10_000);
1789 let small = bench(1_000, 850_000).as_secs_f64() / 850_000.0;
1790 let large = bench(46_000, 850_000).as_secs_f64() / 850_000.0;
1791 assert!(
1792 large < small * 8.0,
1793 "per-hit cost scaled with n_slots: {:.1}ns @1k vs {:.1}ns @46k",
1794 small * 1e9,
1795 large * 1e9
1796 );
1797 }
1798
1799 #[test]
1800 fn slru_list_basic_invariants() {
1801 let mut links = vec![SlotLink::none(); 4];
1802 let mut l = SlruList::new();
1803 assert_eq!(l.pop_front(&mut links), None);
1804 l.push_back(2, SEG_PROBATION, &mut links);
1805 l.push_back(0, SEG_PROBATION, &mut links);
1806 l.push_back(3, SEG_PROBATION, &mut links);
1807 assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 0, 3]);
1808 assert_eq!(l.len, 3);
1809 l.unlink(0, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 3]);
1811 l.unlink(3, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2]);
1813 assert_eq!(l.pop_front(&mut links), Some(2)); assert_eq!(l.len, 0);
1815 assert_eq!(l.head, super::NIL);
1816 assert_eq!(l.tail, super::NIL);
1817 assert!(links.iter().all(|k| k.seg == super::SEG_NONE));
1818 }
1819}
1820
1821#[cfg(test)]
1822mod vram_fraction_tests {
1823 use super::{
1824 parse_cache_hard_vram_frac, parse_cache_lfu_decay, parse_cache_lfu_mtp_weight,
1825 size_class_plan,
1826 };
1827
1828 #[test]
1829 fn hard_vram_fraction_defaults_and_rejects_unsafe_values() {
1830 assert_eq!(parse_cache_hard_vram_frac(None), Ok(0.80));
1831 assert_eq!(parse_cache_hard_vram_frac(Some("0.82")), Ok(0.82));
1832 assert!(parse_cache_hard_vram_frac(Some("NaN")).is_err());
1833 assert_eq!(parse_cache_hard_vram_frac(Some("0.95")), Ok(0.95));
1834 assert!(parse_cache_hard_vram_frac(Some("0.96")).is_err());
1835 assert!(parse_cache_hard_vram_frac(Some("1.0")).is_err());
1836 assert!(parse_cache_hard_vram_frac(Some("bad")).is_err());
1837 }
1838
1839 #[test]
1840 fn lfu_decay_is_opt_in_and_bounded() {
1841 assert_eq!(parse_cache_lfu_decay(None), Ok(None));
1842 assert_eq!(parse_cache_lfu_decay(Some("0.8")), Ok(Some(0.8)));
1843 assert_eq!(parse_cache_lfu_decay(Some("1")), Ok(Some(1.0)));
1844 for value in ["0", "-0.1", "1.1", "NaN", "bad"] {
1845 assert!(
1846 parse_cache_lfu_decay(Some(value)).is_err(),
1847 "accepted {value}"
1848 );
1849 }
1850 }
1851
1852 #[test]
1853 fn lfu_mtp_weight_defaults_and_is_bounded() {
1854 assert_eq!(parse_cache_lfu_mtp_weight(None), Ok(1.0));
1855 assert_eq!(parse_cache_lfu_mtp_weight(Some("4")), Ok(4.0));
1856 for value in ["0", "0.1", "65", "NaN", "bad"] {
1857 assert!(
1858 parse_cache_lfu_mtp_weight(Some(value)).is_err(),
1859 "accepted {value}"
1860 );
1861 }
1862 }
1863
1864 #[test]
1865 fn size_class_plan_preserves_classes_and_never_exceeds_budget() {
1866 let blocks = [100usize, 100, 100, 200, 200, 400];
1867 let budget = (108 * 2) + 208 + 408;
1868 let plan = size_class_plan(&blocks, budget);
1869 assert!(plan.iter().all(|(_, count)| *count > 0));
1870 assert!(
1871 plan.iter()
1872 .map(|(bytes, count)| (bytes + 8) * count)
1873 .sum::<usize>()
1874 <= budget
1875 );
1876 assert!(plan.iter().all(|(bytes, count)| {
1877 *count <= blocks.iter().filter(|block| **block == *bytes).count()
1878 }));
1879 }
1880
1881 #[test]
1882 fn size_class_plan_returns_full_inventory_when_it_fits() {
1883 let blocks = [100usize, 100, 200, 400];
1884 let budget: usize = blocks.iter().map(|bytes| bytes + 8).sum();
1885 assert_eq!(
1886 size_class_plan(&blocks, budget),
1887 vec![(100, 2), (200, 1), (400, 1)]
1888 );
1889 }
1890
1891 #[test]
1892 fn size_class_plan_does_not_overflow_on_pathological_sizes() {
1893 let plan = size_class_plan(&[usize::MAX, usize::MAX], usize::MAX);
1894 assert!(plan.is_empty());
1895 }
1896}
1897
1898impl Drop for MoeSlotCache {
1899 fn drop(&mut self) {
1900 let mut safe_to_drop_slots = true;
1903 if self.compute_stream_unknown || !self.compute_sources.is_empty() {
1904 if let Err(err) = self.compute_stream.synchronize() {
1905 safe_to_drop_slots = false;
1906 eprintln!(
1907 "[moe-cache] unknown compute-stream drain failed ({err}); leaking GPU slots for safety"
1908 );
1909 for (_, keepalive) in self.compute_sources.drain() {
1910 std::mem::forget(keepalive);
1911 }
1912 } else {
1913 self.compute_stream_unknown = false;
1914 self.compute_sources.clear();
1915 }
1916 }
1917 let need_copy_drain = self.copy_stream_unknown
1918 || !self.pending.is_empty()
1919 || !self.inflight_sources.is_empty()
1920 || !self.quarantined_sources.is_empty();
1921 if need_copy_drain {
1922 if let Err(err) = self.copy_stream.synchronize() {
1923 safe_to_drop_slots = false;
1924 eprintln!(
1925 "[moe-cache] unknown copy-stream drain failed ({err}); leaking GPU slots for safety"
1926 );
1927 for (_, keepalive) in self.inflight_sources.drain(..) {
1928 std::mem::forget(keepalive);
1929 }
1930 for keepalive in self.quarantined_sources.drain(..) {
1931 std::mem::forget(keepalive);
1932 }
1933 for (_, pending) in self.pending.drain() {
1934 if let Some(keepalive) = pending.keepalive {
1935 std::mem::forget(keepalive);
1936 }
1937 }
1938 } else {
1939 self.copy_stream_unknown = false;
1940 self.inflight_sources.clear();
1941 self.quarantined_sources.clear();
1942 self.pending.clear();
1943 }
1944 }
1945 if let Some(pool) = self.pread.as_mut() {
1946 safe_to_drop_slots &= pool.drain();
1947 } else if self.pread_requested && self.pread_fallbacks != 0 {
1948 eprintln!(
1949 "[spill-pread] backend unavailable; mmap_fallbacks={}",
1950 self.pread_fallbacks
1951 );
1952 }
1953 if !safe_to_drop_slots {
1954 for slot in self.slots.drain(..) {
1955 std::mem::forget(slot);
1956 }
1957 }
1958 }
1959}