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