1use crate::Engine;
18use crate::model::{ExpertKeepalive, ExpertSource};
19use crate::spill_pread::{PreadPool, PreadStats, ReadTicket, SpillIoMode};
20use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, HostSlice, SyncOnDrop};
21use std::collections::{BTreeMap, HashMap, HashSet};
22use std::sync::Arc;
23
24pub const PROJ_GATE: u8 = 0;
26pub const PROJ_UP: u8 = 1;
27pub const PROJ_DOWN: u8 = 2;
28
29#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
31pub struct BlockId {
32 pub layer: u16,
33 pub proj: u8,
34 pub ex: u16,
35}
36impl BlockId {
37 #[inline]
38 pub fn new(layer: u16, proj: u8, ex: u16) -> Self {
39 BlockId { layer, proj, ex }
40 }
41}
42
43#[derive(Clone, Copy, Debug)]
46pub enum DispatchSlot {
47 Resident(usize),
48}
49
50const NIL: u32 = u32::MAX;
52const SEG_NONE: u8 = 0;
53const SEG_PROBATION: u8 = 1;
54const SEG_PROTECTED: u8 = 2;
55
56#[derive(Clone, Copy, Debug, PartialEq, Eq)]
59struct SlotLink {
60 prev: u32,
61 next: u32,
62 seg: u8,
63}
64impl SlotLink {
65 const fn none() -> Self {
66 SlotLink {
67 prev: NIL,
68 next: NIL,
69 seg: SEG_NONE,
70 }
71 }
72}
73
74#[derive(Debug)]
83struct SlruList {
84 head: u32,
85 tail: u32,
86 len: usize,
87}
88impl SlruList {
89 const fn new() -> Self {
90 SlruList {
91 head: NIL,
92 tail: NIL,
93 len: 0,
94 }
95 }
96
97 fn push_back(&mut self, slot: usize, seg: u8, links: &mut [SlotLink]) {
99 debug_assert_eq!(
100 links[slot].seg, SEG_NONE,
101 "slot {slot} already in a segment"
102 );
103 let s = slot as u32;
104 links[slot] = SlotLink {
105 prev: self.tail,
106 next: NIL,
107 seg,
108 };
109 if self.tail != NIL {
110 links[self.tail as usize].next = s;
111 } else {
112 self.head = s;
113 }
114 self.tail = s;
115 self.len += 1;
116 }
117
118 fn pop_front(&mut self, links: &mut [SlotLink]) -> Option<usize> {
120 if self.head == NIL {
121 return None;
122 }
123 let s = self.head as usize;
124 self.unlink(s, links);
125 Some(s)
126 }
127
128 fn unlink(&mut self, slot: usize, links: &mut [SlotLink]) {
130 let l = links[slot];
131 debug_assert_ne!(l.seg, SEG_NONE, "unlink of slot {slot} not in a segment");
132 if l.prev != NIL {
133 links[l.prev as usize].next = l.next;
134 } else {
135 debug_assert_eq!(self.head, slot as u32);
136 self.head = l.next;
137 }
138 if l.next != NIL {
139 links[l.next as usize].prev = l.prev;
140 } else {
141 debug_assert_eq!(self.tail, slot as u32);
142 self.tail = l.prev;
143 }
144 links[slot] = SlotLink::none();
145 self.len -= 1;
146 }
147
148 fn iter<'a>(&self, links: &'a [SlotLink]) -> SlruIter<'a> {
150 SlruIter {
151 links,
152 cur: self.head,
153 }
154 }
155}
156
157struct SlruIter<'a> {
158 links: &'a [SlotLink],
159 cur: u32,
160}
161impl Iterator for SlruIter<'_> {
162 type Item = usize;
163 fn next(&mut self) -> Option<usize> {
164 if self.cur == NIL {
165 return None;
166 }
167 let s = self.cur as usize;
168 self.cur = self.links[s].next;
169 Some(s)
170 }
171}
172
173struct SlotClass {
176 capacity: usize,
177 probation: SlruList,
178 protected: SlruList,
179 free: Vec<usize>,
180 protected_cap: usize,
181}
182
183impl SlotClass {
184 fn on_hit_full(&mut self, slot: usize, links: &mut [SlotLink]) {
188 match links[slot].seg {
189 SEG_PROBATION => {
190 self.probation.unlink(slot, links);
191 self.push_protected(slot, links);
192 }
193 SEG_PROTECTED => {
194 self.protected.unlink(slot, links);
195 self.protected.push_back(slot, SEG_PROTECTED, links); }
197 _ => self.push_protected(slot, links),
199 }
200 }
201
202 fn push_protected(&mut self, slot: usize, links: &mut [SlotLink]) {
204 self.protected.push_back(slot, SEG_PROTECTED, links);
205 while self.protected.len > self.protected_cap {
206 if let Some(demoted) = self.protected.pop_front(links) {
207 self.probation.push_back(demoted, SEG_PROBATION, links);
208 } else {
209 break;
210 }
211 }
212 }
213
214 fn pop_lru(&mut self, links: &mut [SlotLink]) -> Option<usize> {
216 self.probation
217 .pop_front(links)
218 .or_else(|| self.protected.pop_front(links))
219 }
220
221 fn unlink_from_segment(&mut self, slot: usize, links: &mut [SlotLink]) {
223 match links[slot].seg {
224 SEG_PROBATION => self.probation.unlink(slot, links),
225 SEG_PROTECTED => self.protected.unlink(slot, links),
226 _ => {}
227 }
228 }
229}
230
231pub struct MoeSlotCache {
234 slots: Vec<CudaSlice<u8>>, slot_class: Vec<usize>, classes: Vec<SlotClass>,
237 links: Vec<SlotLink>,
240 occupant: Vec<Option<BlockId>>, table: HashMap<BlockId, usize>, frequencies: HashMap<BlockId, f32>,
246 pending: HashMap<BlockId, PendingBlock>,
250 inflight_sources: Vec<(Arc<CudaEvent>, ExpertKeepalive)>,
253 quarantined_sources: Vec<ExpertKeepalive>,
256 compute_sources: HashMap<KeepaliveKey, ExpertKeepalive>,
260 pread: Option<PreadPool>,
263 worker_reads: HashMap<BlockId, WorkerRead>,
266 pread_requested: bool,
267 pread_fallbacks: u64,
268 copy_stream: Arc<CudaStream>,
271 copy_stream_unknown: bool,
272 compute_stream: Arc<CudaStream>,
273 compute_stream_unknown: bool,
274
275 n: usize,
276 max_block_bytes: usize,
277 size_aware: bool,
278 frequency_evict: bool,
279 frequency_decay: Option<f32>,
280 mtp_frequency_weight: f32,
285 last_forward_layer: Option<u16>,
286 last_forward_t: usize,
287 frozen: bool,
291
292 per_layer: HashMap<u16, u32>,
299 dev_rows: HashMap<u16, CudaSlice<u64>>,
304 prewarm_tried: HashSet<u16>,
307
308 pub hits: u64,
310 pub misses: u64,
311 pub staged_bytes: u64, }
313
314struct PendingBlock {
315 slot: usize,
316 ready: Arc<CudaEvent>,
317 keepalive: Option<ExpertKeepalive>,
318}
319
320#[derive(Clone, Copy)]
321struct WorkerRead {
322 ticket: ReadTicket,
323 len: usize,
324}
325
326#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
327enum KeepaliveKey {
328 Pinned(usize),
329 Buffer(usize),
330 Mmap(usize),
331}
332
333impl KeepaliveKey {
334 fn from_owner(owner: &ExpertKeepalive) -> Self {
335 match owner {
336 ExpertKeepalive::Pinned(value) => Self::Pinned(Arc::as_ptr(value) as usize),
337 ExpertKeepalive::Buffer(value) => Self::Buffer(Arc::as_ptr(value) as usize),
338 ExpertKeepalive::Mmap(value) => Self::Mmap(Arc::as_ptr(value) as usize),
339 }
340 }
341}
342
343struct ExactPinnedPrefix<'a>(&'a [u8]);
348
349impl HostSlice<u8> for ExactPinnedPrefix<'_> {
350 fn len(&self) -> usize {
351 self.0.len()
352 }
353
354 unsafe fn stream_synced_slice<'a>(
355 &'a self,
356 _stream: &'a CudaStream,
357 ) -> (&'a [u8], SyncOnDrop<'a>) {
358 (self.0, SyncOnDrop::Record(None))
361 }
362
363 unsafe fn stream_synced_mut_slice<'a>(
364 &'a mut self,
365 _stream: &'a CudaStream,
366 ) -> (&'a mut [u8], SyncOnDrop<'a>) {
367 panic!("ExactPinnedPrefix is a source-only HostSlice")
368 }
369}
370
371fn stage_on_copy_stream(
372 e: &Engine,
373 host_bytes: &[u8],
374 slot: &mut CudaSlice<u8>,
375) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
376 let prior = match e.stream().record_event(None) {
378 Ok(prior) => prior,
379 Err(err) => return Err((err.into(), true)),
380 };
381 if let Err(err) = e.copy_stream.wait(&prior) {
382 return Err((err.into(), true));
383 }
384 match e.stage_expert_async(host_bytes, slot, 0) {
385 Ok(ready) => Ok(Arc::new(ready)),
386 Err(err) => {
387 match e.copy_stream.synchronize() {
390 Ok(()) => Err((err, true)),
391 Err(sync_err) => Err((std::io::Error::other(format!(
392 "copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
393 )).into(), false)),
394 }
395 }
396 }
397}
398
399fn stage_pread_on_compute_stream(
400 e: &Engine,
401 host_bytes: &[u8],
402 slot: &mut CudaSlice<u8>,
403) -> Result<Arc<CudaEvent>, Box<dyn std::error::Error>> {
404 let ready = Arc::new(e.ctx().new_event(None)?);
405 let source = ExactPinnedPrefix(host_bytes);
406 let mut dst = slot.slice_mut(0..host_bytes.len());
407 e.stream().memcpy_htod(&source, &mut dst)?;
408 ready.record(&e.stream())?;
409 Ok(ready)
410}
411
412fn stage_pread_prefetch_on_copy_stream(
413 e: &Engine,
414 host_bytes: &[u8],
415 slot: &mut CudaSlice<u8>,
416) -> Result<Arc<CudaEvent>, (Box<dyn std::error::Error>, bool)> {
417 let ready = match e.ctx().new_event(None) {
418 Ok(ready) => Arc::new(ready),
419 Err(err) => return Err((err.into(), true)),
420 };
421 let source = ExactPinnedPrefix(host_bytes);
422 let mut dst = slot.slice_mut(0..host_bytes.len());
423 let submitted = e
424 .copy_stream
425 .memcpy_htod(&source, &mut dst)
426 .and_then(|()| ready.record(&e.copy_stream));
427 match submitted {
428 Ok(()) => Ok(ready),
429 Err(err) => match e.copy_stream.synchronize() {
430 Ok(()) => Err((err.into(), true)),
431 Err(sync_err) => Err((
432 std::io::Error::other(format!(
433 "pread copy-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
434 ))
435 .into(),
436 false,
437 )),
438 },
439 }
440}
441
442fn size_class_plan(block_bytes: &[usize], budget_bytes: usize) -> Vec<(usize, usize)> {
445 let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
446 for &bytes in block_bytes.iter().filter(|&&bytes| bytes > 0) {
447 *counts.entry(bytes).or_insert(0) += 1;
448 }
449 if counts.is_empty() || budget_bytes == 0 {
450 return Vec::new();
451 }
452 let total_bytes: u128 = counts
453 .iter()
454 .map(|(&bytes, &count)| (bytes as u128 + 8) * count as u128)
455 .sum();
456 let budget = budget_bytes as u128;
457 let mut plan: Vec<(usize, usize, u128)> = counts
458 .iter()
459 .map(|(&bytes, &count)| {
460 let scaled = count as u128 * budget;
461 (
462 bytes,
463 (scaled / total_bytes).min(count as u128) as usize,
464 scaled % total_bytes,
465 )
466 })
467 .collect();
468 let mut used: u128 = plan
469 .iter()
470 .map(|(bytes, count, _)| (*bytes as u128 + 8) * *count as u128)
471 .sum();
472
473 let mut order: Vec<usize> = (0..plan.len()).collect();
476 order.sort_by(|&a, &b| plan[b].2.cmp(&plan[a].2).then(a.cmp(&b)));
477 for index in order {
478 let (bytes, count, _) = plan[index];
479 let available = counts[&bytes];
480 let required = bytes as u128 + 8;
481 if count < available && used + required <= budget {
482 plan[index].1 += 1;
483 used += required;
484 }
485 }
486 plan.into_iter()
487 .filter_map(|(bytes, count, _)| (count > 0).then_some((bytes, count)))
488 .collect()
489}
490
491impl MoeSlotCache {
492 pub fn new(e: &Engine, max_block_bytes: usize) -> Result<Self, Box<dyn std::error::Error>> {
499 let (free, _total) = e.ctx().mem_get_info()?;
500 let hard_frac = cache_hard_vram_frac();
503 let hard_bytes =
504 ((free as f64 * hard_frac) as usize).saturating_sub(2 * (max_block_bytes + 8));
505 let forced_slots = std::env::var("MEMRA_MOE_SLOTS")
506 .ok()
507 .and_then(|s| s.parse::<usize>().ok());
508 let requested_bytes = if let Some(n) = forced_slots {
509 n.saturating_mul(max_block_bytes + 8)
510 } else {
511 let frac = std::env::var("MEMRA_MOE_VRAM_FRAC")
517 .ok()
518 .and_then(|s| s.parse::<f64>().ok())
519 .unwrap_or(0.85);
520 (free as f64 * frac) as usize
521 };
522 let budget_bytes = requested_bytes.min(hard_bytes);
523 let layout = e.moe_cache_layout().unwrap_or_default();
524 let size_aware = forced_slots.is_none()
525 && std::env::var("MEMRA_MOE_SIZE_AWARE").as_deref() == Ok("1")
526 && !layout.is_empty();
527 let frequency_evict = std::env::var("MEMRA_MOE_LFU").as_deref() == Ok("1");
528 let frequency_decay = if frequency_evict {
529 cache_lfu_decay()
530 } else {
531 None
532 };
533 let mtp_frequency_weight = cache_lfu_mtp_weight();
534 let mut class_plan = if size_aware {
535 size_class_plan(&layout, budget_bytes)
536 } else {
537 Vec::new()
538 };
539 if class_plan.iter().map(|(_, count)| count).sum::<usize>() < 8 {
540 let n = (budget_bytes / (max_block_bytes + 8)).max(8);
541 class_plan = vec![(max_block_bytes, n)];
542 }
543 let n: usize = class_plan.iter().map(|(_, count)| count).sum();
544
545 let mut slots = Vec::with_capacity(n);
546 let mut slot_class = Vec::with_capacity(n);
547 let mut classes = Vec::with_capacity(class_plan.len());
548 let mut occupant = Vec::with_capacity(n);
549 for (class_index, &(capacity, count)) in class_plan.iter().enumerate() {
550 let start = slots.len();
551 for _ in 0..count {
552 slots.push(e.alloc_u8(capacity + 8)?);
554 slot_class.push(class_index);
555 occupant.push(None);
556 }
557 let free_slots = (start..start + count).rev().collect();
558 classes.push(SlotClass {
559 capacity,
560 probation: SlruList::new(),
561 protected: SlruList::new(),
562 free: free_slots,
563 protected_cap: ((count as f64 * 0.8) as usize).max(1),
564 });
565 }
566 let links = vec![SlotLink::none(); n];
567 if size_aware {
568 let allocated: usize = class_plan
569 .iter()
570 .map(|(bytes, count)| (bytes + 8) * count)
571 .sum();
572 eprintln!(
573 "[moe-cache] size-aware fixed slots: {n} slots in {} classes, {:.2} GB / {:.2} GB budget",
574 class_plan.len(),
575 allocated as f64 / 1e9,
576 budget_bytes as f64 / 1e9
577 );
578 }
579 let pread_mode = crate::spill_pread::configured_mode();
580 let pread_requested = pread_mode != SpillIoMode::Mmap;
581 let pread = if pread_requested {
582 match PreadPool::try_new(e, max_block_bytes, pread_mode) {
583 Ok(pool) => Some(pool),
584 Err(err) => {
585 eprintln!(
586 "[spill-pread] pinned-buffer initialization failed ({err}); using mmap"
587 );
588 None
589 }
590 }
591 } else {
592 None
593 };
594
595 Ok(MoeSlotCache {
596 slots,
597 slot_class,
598 classes,
599 links,
600 occupant,
601 table: HashMap::with_capacity(n * 2),
602 frequencies: HashMap::with_capacity(layout.len().max(n * 2)),
603 pending: HashMap::new(),
604 inflight_sources: Vec::new(),
605 quarantined_sources: Vec::new(),
606 compute_sources: HashMap::new(),
607 pread,
608 worker_reads: HashMap::new(),
609 pread_requested,
610 pread_fallbacks: 0,
611 copy_stream: e.copy_stream.clone(),
612 copy_stream_unknown: false,
613 compute_stream: e.stream().clone(),
614 compute_stream_unknown: false,
615 n,
616 max_block_bytes,
617 size_aware,
618 frequency_evict,
619 frequency_decay,
620 mtp_frequency_weight,
621 last_forward_layer: None,
622 last_forward_t: 0,
623 frozen: false,
624 per_layer: HashMap::new(),
625 dev_rows: HashMap::new(),
626 prewarm_tried: HashSet::new(),
627 hits: 0,
628 misses: 0,
629 staged_bytes: 0,
630 })
631 }
632
633 #[inline]
634 pub fn n_slots(&self) -> usize {
635 self.n
636 }
637 #[inline]
638 pub fn is_frozen(&self) -> bool {
639 self.frozen
640 }
641 pub fn freeze(&mut self) {
642 if !self.frozen {
643 self.frozen = true;
644 let (_, complete, one_projection, two_projections, stranded_blocks) =
645 self.expert_residency_shape();
646 eprintln!(
647 "[moe-cache] residency frozen: {} slots, {} resident blocks; \
648 {complete} complete experts, {one_projection} one-projection fragments, \
649 {two_projections} two-projection fragments ({stranded_blocks} stranded blocks)",
650 self.n,
651 self.table.len()
652 );
653 let mut mtp_masks = HashMap::<u16, u8>::new();
654 for id in self.table.keys().filter(|id| id.layer == u16::MAX) {
655 *mtp_masks.entry(id.ex).or_insert(0) |= 1u8 << id.proj;
656 }
657 if !mtp_masks.is_empty() {
658 let complete = mtp_masks.values().filter(|&&mask| mask == 0b111).count();
659 eprintln!(
660 "[moe-cache] frozen MTP residency: {} blocks, {complete} complete experts",
661 mtp_masks
662 .values()
663 .map(|mask| mask.count_ones() as usize)
664 .sum::<usize>()
665 );
666 }
667 }
668 }
669
670 pub(crate) fn expert_residency_shape(&self) -> (usize, usize, usize, usize, usize) {
671 let mut masks = HashMap::<(u16, u16), u8>::new();
672 for id in self.table.keys() {
673 *masks.entry((id.layer, id.ex)).or_insert(0) |= 1u8 << id.proj;
674 }
675 let complete = masks.values().filter(|&&mask| mask == 0b111).count();
676 let one_projection = masks
677 .values()
678 .filter(|&&mask| mask.count_ones() == 1)
679 .count();
680 let two_projections = masks
681 .values()
682 .filter(|&&mask| mask.count_ones() == 2)
683 .count();
684 let stranded_blocks = one_projection + 2 * two_projections;
685 (
686 masks.len(),
687 complete,
688 one_projection,
689 two_projections,
690 stranded_blocks,
691 )
692 }
693
694 #[inline]
695 pub fn max_block_bytes(&self) -> usize {
696 self.max_block_bytes
697 }
698
699 #[inline]
701 pub fn resident(&self, id: BlockId) -> Option<usize> {
702 self.table.get(&id).copied()
703 }
704
705 #[inline]
706 fn frequency_increment(&self, id: BlockId) -> f32 {
707 if id.layer == u16::MAX {
708 self.mtp_frequency_weight
709 } else {
710 1.0
711 }
712 }
713
714 pub(crate) fn note_profile_hit(&mut self, id: BlockId) {
718 if self.frozen || !self.table.contains_key(&id) {
719 return;
720 }
721 let increment = self.frequency_increment(id);
722 *self.frequencies.entry(id).or_insert(0.0) += increment;
723 }
724
725 fn on_hit(&mut self, slot: usize) {
743 let class_index = self.slot_class[slot];
744 let class = &mut self.classes[class_index];
745 if !class.free.is_empty() {
746 return;
747 }
748 class.on_hit_full(slot, &mut self.links);
749 }
750
751 fn remove_occupant(&mut self, slot: usize) {
752 if let Some(old) = self.occupant[slot].take() {
753 self.table.remove(&old);
754 self.on_block_evicted(old.layer);
755 }
756 }
757
758 fn frequency_victim_in_class(&mut self, class_index: usize, keep: &[BlockId]) -> Option<usize> {
765 let class = &self.classes[class_index];
766 let candidate = class
767 .probation
768 .iter(&self.links)
769 .chain(class.protected.iter(&self.links))
770 .enumerate()
771 .filter_map(|(position, slot)| {
772 let id = self.occupant[slot]?;
773 (!keep.contains(&id)).then_some((
774 self.frequencies.get(&id).copied().unwrap_or(0.0),
775 position,
776 slot,
777 ))
778 })
779 .min_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
780 let (_, _, slot) = candidate?;
781 self.classes[class_index].unlink_from_segment(slot, &mut self.links);
782 Some(slot)
783 }
784
785 fn evict_one(&mut self, required: usize) -> Option<usize> {
787 for class_index in 0..self.classes.len() {
788 if self.classes[class_index].capacity < required {
789 continue;
790 }
791 let slot = if self.frequency_evict {
792 self.frequency_victim_in_class(class_index, &[])
793 } else {
794 self.classes[class_index].pop_lru(&mut self.links)
795 };
796 if let Some(slot) = slot {
797 self.remove_occupant(slot);
798 return Some(slot);
799 }
800 }
801 None
802 }
803
804 fn evict_one_excluding(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
810 fn take(
811 q: &mut SlruList,
812 links: &mut [SlotLink],
813 occupant: &[Option<BlockId>],
814 keep: &[BlockId],
815 ) -> Option<usize> {
816 let slot = q
817 .iter(links)
818 .find(|&s| occupant[s].is_some_and(|id| !keep.contains(&id)))?;
819 q.unlink(slot, links);
820 Some(slot)
821 }
822 for class_index in 0..self.classes.len() {
823 if self.classes[class_index].capacity < required {
824 continue;
825 }
826 let slot = if self.frequency_evict {
827 self.frequency_victim_in_class(class_index, keep)
828 } else {
829 let class = &mut self.classes[class_index];
830 take(&mut class.probation, &mut self.links, &self.occupant, keep)
831 .or_else(|| take(&mut class.protected, &mut self.links, &self.occupant, keep))
832 };
833 if let Some(slot) = slot {
834 self.remove_occupant(slot);
835 return Some(slot);
836 }
837 }
838 None
839 }
840
841 fn on_block_evicted(&mut self, layer: u16) {
847 if let Some(c) = self.per_layer.get_mut(&layer) {
848 *c -= 1;
849 }
850 self.dev_rows.remove(&layer);
851 }
852
853 fn reserve_slot(&mut self, required: usize) -> Option<usize> {
854 for class in &mut self.classes {
855 if class.capacity >= required {
856 if let Some(slot) = class.free.pop() {
857 return Some(slot);
858 }
859 }
860 }
861 self.evict_one(required)
862 }
863
864 fn release_reserved_slot(&mut self, slot: usize) {
865 debug_assert!(self.occupant[slot].is_none());
866 self.classes[self.slot_class[slot]].free.push(slot);
867 }
868
869 fn publish(&mut self, id: BlockId, slot: usize) {
870 self.occupant[slot] = Some(id);
871 self.table.insert(id, slot);
872 self.classes[self.slot_class[slot]].probation.push_back(
873 slot,
874 SEG_PROBATION,
875 &mut self.links,
876 );
877 *self.per_layer.entry(id.layer).or_insert(0) += 1;
878 }
879
880 fn reap_copy_sources(&mut self) {
881 self.inflight_sources
882 .retain(|(ready, _)| !ready.is_complete());
883 }
884
885 fn retain_compute_source(&mut self, owner: Option<ExpertKeepalive>) {
886 if let Some(owner) = owner {
887 let key = KeepaliveKey::from_owner(&owner);
888 self.compute_sources.entry(key).or_insert(owner);
889 }
890 }
891
892 fn admit(
895 &mut self,
896 id: BlockId,
897 host_bytes: &[u8],
898 e: &Engine,
899 ) -> Result<usize, Box<dyn std::error::Error>> {
900 let slot = self.reserve_slot(host_bytes.len()).ok_or_else(|| {
901 std::io::Error::other(format!(
902 "no MoE cache slot can hold {} bytes (max class {})",
903 host_bytes.len(),
904 self.classes.last().map(|class| class.capacity).unwrap_or(0)
905 ))
906 })?;
907 if let Err(err) = e.stage_expert(host_bytes, &mut self.slots[slot], 0) {
910 return match e.stream().synchronize() {
911 Ok(()) => {
912 self.release_reserved_slot(slot);
913 Err(err)
914 }
915 Err(sync_err) => {
916 self.compute_stream_unknown = true;
919 Err(std::io::Error::other(format!(
920 "compute-stream H2D setup failed ({err}); stream drain also failed ({sync_err})"
921 )).into())
922 }
923 };
924 }
925 self.staged_bytes += host_bytes.len() as u64;
926 self.publish(id, slot);
927 Ok(slot)
928 }
929
930 fn note_pread_fallback(&mut self, reason: &dyn std::fmt::Display) {
931 self.pread_fallbacks += 1;
932 if let Some(pool) = self.pread.as_mut() {
933 pool.note_fallback();
934 }
935 if self.pread_fallbacks <= 3 {
936 eprintln!("[spill-pread] falling back to mmap: {reason}");
937 }
938 }
939
940 pub(crate) fn begin_worker_scope(&mut self) {
944 if self.worker_reads.is_empty() {
945 return;
946 }
947 let tickets: Vec<_> = self
948 .worker_reads
949 .drain()
950 .map(|(_, read)| read.ticket)
951 .collect();
952 if let Some(pool) = self.pread.as_mut().filter(|pool| pool.is_worker()) {
953 for ticket in tickets {
954 let _ = pool.cancel_worker(ticket);
955 }
956 }
957 }
958
959 pub(crate) fn begin_forward_epoch(&mut self, layer: u16, t: usize) {
969 let Some(decay) = self.frequency_decay else {
970 self.last_forward_layer = Some(layer);
971 self.last_forward_t = t;
972 return;
973 };
974 let new_sweep = self
975 .last_forward_layer
976 .is_some_and(|previous| layer <= previous);
977 if new_sweep && t == 1 {
978 if self.last_forward_t != 1 {
979 self.frequencies.clear();
980 } else {
981 self.frequencies.retain(|_, score| {
982 *score *= decay;
983 *score >= 1.0e-3
984 });
985 }
986 }
987 self.last_forward_layer = Some(layer);
988 self.last_forward_t = t;
989 }
990
991 pub(crate) fn promote_worker_reads_at_safe_boundary(
998 &mut self,
999 order: &[BlockId],
1000 keep: &[BlockId],
1001 e: &Engine,
1002 ) -> Result<usize, Box<dyn std::error::Error>> {
1003 if !crate::spill_pread::copy_h2d_enabled() {
1004 return Ok(0);
1005 }
1006 let mut promoted = 0usize;
1007 for &id in order {
1008 if self.table.contains_key(&id) || self.pending.contains_key(&id) {
1009 if let Some(read) = self.worker_reads.remove(&id) {
1010 if let Some(pool) = self.pread.as_mut() {
1011 let _ = pool.cancel_worker(read.ticket);
1012 }
1013 }
1014 continue;
1015 }
1016 let Some(read) = self.worker_reads.get(&id).copied() else {
1017 continue;
1018 };
1019 let Some(slot) = self.reserve_prefetch_slot(read.len, keep) else {
1020 continue;
1021 };
1022 self.worker_reads.remove(&id);
1023
1024 let index = match self.pread.as_mut().unwrap().wait_worker(read.ticket) {
1025 Ok(index) => index,
1026 Err(err) => {
1027 let _ = self.pread.as_mut().unwrap().cancel_worker(read.ticket);
1028 self.release_reserved_slot(slot);
1029 self.note_pread_fallback(err.as_ref());
1030 continue;
1031 }
1032 };
1033 let ready = {
1034 let bytes = match self.pread.as_ref().unwrap().bytes(index, read.len) {
1035 Ok(bytes) => bytes,
1036 Err(err) => {
1037 self.pread.as_mut().unwrap().abort_read(index);
1038 self.release_reserved_slot(slot);
1039 self.note_pread_fallback(err.as_ref());
1040 continue;
1041 }
1042 };
1043 stage_pread_prefetch_on_copy_stream(e, bytes, &mut self.slots[slot])
1044 };
1045 let ready = match ready {
1046 Ok(ready) => ready,
1047 Err((err, reusable)) => {
1048 if reusable {
1049 self.pread.as_mut().unwrap().abort_read(index);
1050 self.release_reserved_slot(slot);
1051 self.note_pread_fallback(err.as_ref());
1052 continue;
1053 }
1054 self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1055 self.copy_stream_unknown = true;
1056 return Err(err);
1057 }
1058 };
1059 self.pread.as_mut().unwrap().mark_h2d(index, ready.clone());
1060 self.occupant[slot] = Some(id);
1061 self.pending.insert(
1062 id,
1063 PendingBlock {
1064 slot,
1065 ready,
1066 keepalive: None,
1067 },
1068 );
1069 self.staged_bytes += read.len as u64;
1070 promoted += 1;
1071 }
1072 Ok(promoted)
1073 }
1074
1075 fn dispatch_disk(
1076 &mut self,
1077 id: BlockId,
1078 file: &Arc<std::fs::File>,
1079 offset: u64,
1080 len: usize,
1081 fallback: &[u8],
1082 e: &Engine,
1083 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1084 if self.pread.is_none() {
1085 if self.pread_requested {
1086 self.note_pread_fallback(&"pinned-buffer backend unavailable");
1087 }
1088 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1089 }
1090
1091 let pending = self.worker_reads.remove(&id);
1092 let pool = self.pread.as_mut().unwrap();
1093 let read = if pool.is_worker() {
1094 let ticket = match pending {
1095 Some(read) => Ok(Some(read.ticket)),
1096 None => pool.submit_worker(file.clone(), offset, len),
1097 };
1098 match ticket {
1099 Ok(Some(ticket)) => match pool.wait_worker(ticket) {
1100 Ok(index) => Ok(index),
1101 Err(err) => {
1102 let _ = pool.cancel_worker(ticket);
1105 Err(err)
1106 }
1107 },
1108 Ok(None) => Err(std::io::Error::other("worker read ring is busy").into()),
1109 Err(err) => Err(err),
1110 }
1111 } else {
1112 debug_assert!(pending.is_none());
1113 pool.read(file.as_ref(), offset, len)
1114 };
1115 let index = match read {
1116 Ok(index) => index,
1117 Err(err) => {
1118 self.note_pread_fallback(err.as_ref());
1119 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1120 }
1121 };
1122
1123 let slot = self.reserve_slot(len).ok_or_else(|| {
1126 std::io::Error::other(format!(
1127 "no MoE cache slot can hold {len} bytes (max class {})",
1128 self.classes.last().map(|class| class.capacity).unwrap_or(0)
1129 ))
1130 })?;
1131 let ready = {
1132 let bytes = match self.pread.as_ref().unwrap().bytes(index, len) {
1133 Ok(bytes) => bytes,
1134 Err(err) => {
1135 self.pread.as_mut().unwrap().abort_read(index);
1136 self.release_reserved_slot(slot);
1137 self.note_pread_fallback(err.as_ref());
1138 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1139 }
1140 };
1141 stage_pread_on_compute_stream(e, bytes, &mut self.slots[slot])
1142 };
1143 let ready = match ready {
1144 Ok(ready) => ready,
1145 Err(err) => {
1146 match e.stream().synchronize() {
1150 Ok(()) => {
1151 self.pread.as_mut().unwrap().abort_read(index);
1152 self.release_reserved_slot(slot);
1153 self.note_pread_fallback(err.as_ref());
1154 return Ok(DispatchSlot::Resident(self.admit(id, fallback, e)?));
1155 }
1156 Err(sync_err) => {
1157 self.pread.as_mut().unwrap().mark_unknown_h2d(index);
1158 return Err(std::io::Error::other(format!(
1159 "pread H2D setup failed ({err}); CUDA stream drain also failed ({sync_err})"
1160 )).into());
1161 }
1162 }
1163 }
1164 };
1165 self.pread.as_mut().unwrap().mark_h2d(index, ready);
1166 self.staged_bytes += len as u64;
1169 self.publish(id, slot);
1170 Ok(DispatchSlot::Resident(slot))
1171 }
1172
1173 pub fn dispatch(
1181 &mut self,
1182 id: BlockId,
1183 host_bytes: &[u8],
1184 e: &Engine,
1185 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1186 self.dispatch_source(
1187 id,
1188 ExpertSource::Memory {
1189 bytes: host_bytes,
1190 keepalive: None,
1191 },
1192 e,
1193 )
1194 }
1195
1196 pub(crate) fn dispatch_source(
1197 &mut self,
1198 id: BlockId,
1199 source: ExpertSource<'_>,
1200 e: &Engine,
1201 ) -> Result<DispatchSlot, Box<dyn std::error::Error>> {
1202 self.reap_copy_sources();
1203 let increment = self.frequency_increment(id);
1204 *self.frequencies.entry(id).or_insert(0.0) += increment;
1205 if let Some(s) = self.table.get(&id).copied() {
1206 self.hits += 1;
1207 self.on_hit(s);
1208 return Ok(DispatchSlot::Resident(s));
1209 }
1210 if let Some(pending) = self.pending.remove(&id) {
1211 if let Err(err) = e.compute_wait(pending.ready.as_ref()) {
1212 self.pending.insert(id, pending);
1213 return Err(err);
1214 }
1215 self.misses += 1;
1216 let slot = pending.slot;
1217 if let Some(keepalive) = pending.keepalive {
1218 self.inflight_sources.push((pending.ready, keepalive));
1219 }
1220 self.publish(id, slot);
1221 return Ok(DispatchSlot::Resident(slot));
1222 }
1223 self.misses += 1;
1224 match source {
1234 ExpertSource::Memory { bytes, keepalive } => {
1235 self.retain_compute_source(keepalive);
1238 let slot = self.admit(id, bytes, e)?;
1239 Ok(DispatchSlot::Resident(slot))
1240 }
1241 ExpertSource::Disk {
1242 file,
1243 offset,
1244 len,
1245 fallback,
1246 keepalive,
1247 } => {
1248 self.retain_compute_source(Some(keepalive));
1251 self.dispatch_disk(id, file, offset, len, fallback, e)
1252 }
1253 }
1254 }
1255
1256 pub fn prefetch(
1265 &mut self,
1266 id: BlockId,
1267 host_bytes: &[u8],
1268 keep: &[BlockId],
1269 e: &Engine,
1270 ) -> Result<bool, Box<dyn std::error::Error>> {
1271 self.prefetch_source(
1272 id,
1273 ExpertSource::Memory {
1274 bytes: host_bytes,
1275 keepalive: None,
1276 },
1277 keep,
1278 e,
1279 )
1280 }
1281
1282 fn reserve_prefetch_slot(&mut self, required: usize, keep: &[BlockId]) -> Option<usize> {
1283 for class in &mut self.classes {
1284 if class.capacity >= required {
1285 if let Some(slot) = class.free.pop() {
1286 return Some(slot);
1287 }
1288 }
1289 }
1290 self.evict_one_excluding(required, keep)
1291 }
1292
1293 fn prefetch_bytes(
1294 &mut self,
1295 id: BlockId,
1296 host_bytes: &[u8],
1297 keepalive: Option<ExpertKeepalive>,
1298 keep: &[BlockId],
1299 e: &Engine,
1300 ) -> Result<bool, Box<dyn std::error::Error>> {
1301 let Some(slot) = self.reserve_prefetch_slot(host_bytes.len(), keep) else {
1302 return Ok(false);
1303 };
1304 let ready = match stage_on_copy_stream(e, host_bytes, &mut self.slots[slot]) {
1305 Ok(ready) => ready,
1306 Err((err, reusable)) => {
1307 if reusable {
1308 self.release_reserved_slot(slot);
1309 } else {
1310 self.copy_stream_unknown = true;
1313 if let Some(keepalive) = keepalive {
1314 self.quarantined_sources.push(keepalive);
1315 }
1316 eprintln!(
1317 "[moe-cache] quarantining slot {slot} after unprovable copy completion"
1318 );
1319 }
1320 return Err(err);
1321 }
1322 };
1323 self.occupant[slot] = Some(id);
1324 self.pending.insert(
1325 id,
1326 PendingBlock {
1327 slot,
1328 ready,
1329 keepalive,
1330 },
1331 );
1332 self.staged_bytes += host_bytes.len() as u64;
1333 Ok(true)
1334 }
1335
1336 pub(crate) fn prefetch_source(
1337 &mut self,
1338 id: BlockId,
1339 source: ExpertSource<'_>,
1340 keep: &[BlockId],
1341 e: &Engine,
1342 ) -> Result<bool, Box<dyn std::error::Error>> {
1343 self.reap_copy_sources();
1344 if self.table.contains_key(&id)
1345 || self.pending.contains_key(&id)
1346 || self.worker_reads.contains_key(&id)
1347 {
1348 return Ok(false);
1349 }
1350 match source {
1351 ExpertSource::Memory { bytes, keepalive } => {
1352 self.prefetch_bytes(id, bytes, keepalive, keep, e)
1353 }
1354 ExpertSource::Disk {
1355 file,
1356 offset,
1357 len,
1358 fallback,
1359 keepalive,
1360 } => {
1361 if self.pread.as_ref().is_some_and(PreadPool::is_worker) {
1362 match self.pread.as_mut().unwrap().submit_worker_speculative(
1363 file.clone(),
1364 offset,
1365 len,
1366 ) {
1367 Ok(Some(ticket)) => {
1368 self.worker_reads.insert(id, WorkerRead { ticket, len });
1369 Ok(true)
1370 }
1371 Ok(None) => Ok(false),
1372 Err(err) => {
1373 self.note_pread_fallback(err.as_ref());
1374 Ok(false)
1375 }
1376 }
1377 } else if self.pread.is_some() {
1378 Ok(false)
1380 } else {
1381 self.prefetch_bytes(id, fallback, Some(keepalive), keep, e)
1382 }
1383 }
1384 }
1385 }
1386
1387 pub fn force_admit(
1389 &mut self,
1390 id: BlockId,
1391 host_bytes: &[u8],
1392 e: &Engine,
1393 ) -> Result<usize, Box<dyn std::error::Error>> {
1394 if let Some(s) = self.table.get(&id).copied() {
1395 return Ok(s);
1396 }
1397 self.admit(id, host_bytes, e)
1398 }
1399
1400 pub fn export_residency(&self) -> Vec<(u16, u8, u16)> {
1408 self.occupant
1409 .iter()
1410 .flatten()
1411 .map(|id| (id.layer, id.proj, id.ex))
1412 .collect()
1413 }
1414
1415 pub fn restage_block(
1420 &mut self,
1421 id: BlockId,
1422 m: &crate::hybrid::MoeWeights,
1423 e: &Engine,
1424 ) -> Result<bool, Box<dyn std::error::Error>> {
1425 if self.table.contains_key(&id) {
1426 return Ok(true);
1427 }
1428 let exps = match id.proj {
1429 PROJ_GATE => &m.gate_exps,
1430 PROJ_UP => &m.up_exps,
1431 PROJ_DOWN => &m.down_exps,
1432 _ => return Ok(false),
1433 };
1434 if id.ex as usize >= exps.n_expert {
1435 return Ok(false);
1436 }
1437 if m.active_experts
1438 .as_ref()
1439 .is_some_and(|active| !active[id.ex as usize])
1440 {
1441 return Ok(false);
1442 }
1443 if exps.expert_layout(id.ex as usize).len == 0 {
1444 return Ok(false);
1445 }
1446 match exps.expert_source(id.ex as usize) {
1447 ExpertSource::Memory { bytes, keepalive } => {
1448 self.retain_compute_source(keepalive);
1449 self.admit(id, bytes, e)?;
1450 }
1451 ExpertSource::Disk {
1452 fallback,
1453 keepalive,
1454 ..
1455 } => {
1456 self.retain_compute_source(Some(keepalive));
1457 self.admit(id, fallback, e)?;
1458 }
1459 }
1460 Ok(true)
1461 }
1462
1463 pub fn prewarm_layer(
1464 &mut self,
1465 layer: u16,
1466 m: &crate::hybrid::MoeWeights,
1467 e: &Engine,
1468 ) -> Result<(), Box<dyn std::error::Error>> {
1469 if !self.prewarm_tried.insert(layer) {
1470 return Ok(());
1471 }
1472 let n_expert = m.gate_exps.n_expert;
1473 if self.pread.is_some()
1474 && (0..n_expert).any(|ex| {
1475 matches!(m.gate_exps.expert_source(ex), ExpertSource::Disk { .. })
1476 || matches!(m.up_exps.expert_source(ex), ExpertSource::Disk { .. })
1477 || matches!(m.down_exps.expert_source(ex), ExpertSource::Disk { .. })
1478 })
1479 {
1480 return Ok(());
1483 }
1484 let resident = self.per_layer.get(&layer).copied().unwrap_or(0) as usize;
1485 let missing = 3 * n_expert - resident;
1486 if self.size_aware {
1487 return Ok(());
1488 } if self
1490 .classes
1491 .iter()
1492 .map(|class| class.free.len())
1493 .sum::<usize>()
1494 < missing
1495 {
1496 return Ok(()); }
1498 for ex in 0..n_expert {
1499 for (proj, exps) in [
1500 (PROJ_GATE, &m.gate_exps),
1501 (PROJ_UP, &m.up_exps),
1502 (PROJ_DOWN, &m.down_exps),
1503 ] {
1504 let id = BlockId::new(layer, proj, ex as u16);
1505 if self.table.contains_key(&id) {
1506 continue;
1507 }
1508 match exps.expert_source(ex) {
1509 ExpertSource::Memory { bytes, keepalive } => {
1510 self.retain_compute_source(keepalive);
1511 self.admit(id, bytes, e)?;
1512 }
1513 ExpertSource::Disk {
1514 fallback,
1515 keepalive,
1516 ..
1517 } => {
1518 self.retain_compute_source(Some(keepalive));
1519 self.admit(id, fallback, e)?;
1520 }
1521 }
1522 }
1523 }
1524 Ok(())
1525 }
1526
1527 pub fn layer_dev_row(
1533 &mut self,
1534 layer: u16,
1535 n_expert: usize,
1536 e: &Engine,
1537 ) -> Result<Option<&CudaSlice<u64>>, Box<dyn std::error::Error>> {
1538 if self.per_layer.get(&layer).copied().unwrap_or(0) as usize != 3 * n_expert {
1539 return Ok(None);
1540 }
1541 if !self.dev_rows.contains_key(&layer) {
1542 use cudarc::driver::DevicePtr;
1543 let mut host = vec![0u64; 3 * n_expert];
1544 for proj in 0..3u8 {
1545 for ex in 0..n_expert {
1546 let Some(&s) = self.table.get(&BlockId::new(layer, proj, ex as u16)) else {
1547 return Ok(None);
1549 };
1550 let __s_ev = e.stream();
1551 let (p, _ev) = self.slots[s].device_ptr(&__s_ev);
1552 host[proj as usize * n_expert + ex] = p as u64;
1553 }
1554 }
1555 let row = e.stream().clone_htod(&host)?;
1556 self.dev_rows.insert(layer, row);
1557 }
1558 Ok(self.dev_rows.get(&layer))
1559 }
1560
1561 #[inline]
1563 pub fn buf(&self, d: DispatchSlot) -> &CudaSlice<u8> {
1564 match d {
1565 DispatchSlot::Resident(s) => &self.slots[s],
1566 }
1567 }
1568
1569 #[inline]
1571 pub fn slot(&self, s: usize) -> &CudaSlice<u8> {
1572 &self.slots[s]
1573 }
1574
1575 pub fn hit_rate(&self) -> f64 {
1577 let tot = self.hits + self.misses;
1578 if tot == 0 {
1579 0.0
1580 } else {
1581 self.hits as f64 / tot as f64
1582 }
1583 }
1584
1585 pub fn reset_counters(&mut self) {
1587 self.hits = 0;
1588 self.misses = 0;
1589 self.staged_bytes = 0;
1590 }
1591
1592 pub(crate) fn pread_stats(&self) -> Option<PreadStats> {
1593 if !self.pread_requested {
1594 return None;
1595 }
1596 let mut stats = self
1597 .pread
1598 .as_ref()
1599 .map(PreadPool::stats)
1600 .unwrap_or_default();
1601 stats.fallbacks = self.pread_fallbacks;
1602 Some(stats)
1603 }
1604}
1605
1606fn cache_lfu_decay() -> Option<f32> {
1607 let raw = std::env::var("MEMRA_MOE_LFU_DECAY").ok()?;
1608 match parse_cache_lfu_decay(Some(&raw)) {
1609 Ok(value) => value,
1610 Err(reason) => {
1611 eprintln!(
1612 "[moe-cache] invalid MEMRA_MOE_LFU_DECAY={raw:?} ({reason}); disabling LFU decay"
1613 );
1614 None
1615 }
1616 }
1617}
1618
1619fn cache_lfu_mtp_weight() -> f32 {
1620 const DEFAULT: f32 = 1.0;
1621 let raw = std::env::var("MEMRA_MOE_LFU_MTP_WEIGHT").ok();
1622 match parse_cache_lfu_mtp_weight(raw.as_deref()) {
1623 Ok(value) => value,
1624 Err(reason) => {
1625 eprintln!(
1626 "[moe-cache] invalid MEMRA_MOE_LFU_MTP_WEIGHT={:?} ({reason}); using {DEFAULT}",
1627 raw.as_deref().unwrap_or("")
1628 );
1629 DEFAULT
1630 }
1631 }
1632}
1633
1634fn parse_cache_lfu_mtp_weight(raw: Option<&str>) -> Result<f32, &'static str> {
1635 let value = raw
1636 .unwrap_or("1")
1637 .parse::<f32>()
1638 .map_err(|_| "expected a number")?;
1639 if value.is_finite() && (0.25..=64.0).contains(&value) {
1640 Ok(value)
1641 } else {
1642 Err("expected a finite multiplier from 0.25 through 64")
1643 }
1644}
1645
1646fn parse_cache_lfu_decay(raw: Option<&str>) -> Result<Option<f32>, &'static str> {
1647 let Some(raw) = raw else { return Ok(None) };
1648 let value = raw.parse::<f32>().map_err(|_| "expected a number")?;
1649 if value.is_finite() && value > 0.0 && value <= 1.0 {
1650 Ok(Some(value))
1651 } else {
1652 Err("expected a finite fraction greater than 0 and at most 1")
1653 }
1654}
1655
1656fn cache_hard_vram_frac() -> f64 {
1657 const DEFAULT: f64 = 0.80;
1658 let raw = std::env::var("MEMRA_MOE_HARD_VRAM_FRAC").ok();
1659 match parse_cache_hard_vram_frac(raw.as_deref()) {
1660 Ok(value) => value,
1661 Err(reason) => {
1662 eprintln!(
1663 "[moe-cache] invalid MEMRA_MOE_HARD_VRAM_FRAC={:?} ({reason}); using {DEFAULT}",
1664 raw.as_deref().unwrap_or("")
1665 );
1666 DEFAULT
1667 }
1668 }
1669}
1670
1671fn parse_cache_hard_vram_frac(raw: Option<&str>) -> Result<f64, &'static str> {
1672 let value = raw
1673 .unwrap_or("0.80")
1674 .parse::<f64>()
1675 .map_err(|_| "expected a number")?;
1676 if value.is_finite() && (0.10..=0.95).contains(&value) {
1677 Ok(value)
1678 } else {
1679 Err("expected a finite fraction from 0.10 through 0.95")
1680 }
1681}
1682
1683#[cfg(test)]
1684mod slru_intrusive_tests {
1685 use super::{SEG_PROBATION, SEG_PROTECTED, SlotClass, SlotLink, SlruList};
1692 use std::collections::VecDeque;
1693
1694 struct OldSlru {
1696 probation: VecDeque<usize>,
1697 protected: VecDeque<usize>,
1698 protected_cap: usize,
1699 }
1700 impl OldSlru {
1701 fn on_hit_full(&mut self, slot: usize) {
1702 if let Some(pos) = self.probation.iter().position(|&x| x == slot) {
1703 self.probation.remove(pos);
1704 self.push_protected(slot);
1705 } else if let Some(pos) = self.protected.iter().position(|&x| x == slot) {
1706 self.protected.remove(pos);
1707 self.protected.push_back(slot); } else {
1709 self.push_protected(slot);
1710 }
1711 }
1712 fn push_protected(&mut self, slot: usize) {
1713 self.protected.push_back(slot);
1714 while self.protected.len() > self.protected_cap {
1715 if let Some(demoted) = self.protected.pop_front() {
1716 self.probation.push_back(demoted);
1717 } else {
1718 break;
1719 }
1720 }
1721 }
1722 fn pop_lru(&mut self) -> Option<usize> {
1723 self.probation
1724 .pop_front()
1725 .or_else(|| self.protected.pop_front())
1726 }
1727 fn take_excluding(&mut self, banned: &[usize]) -> Option<usize> {
1728 let take = |q: &mut VecDeque<usize>| {
1729 q.iter()
1730 .position(|&s| !banned.contains(&s))
1731 .and_then(|pos| q.remove(pos))
1732 };
1733 take(&mut self.probation).or_else(|| take(&mut self.protected))
1734 }
1735 }
1736
1737 fn new_pair(n: usize, protected_cap: usize) -> (SlotClass, Vec<SlotLink>, OldSlru) {
1738 let class = SlotClass {
1739 capacity: 1,
1740 probation: SlruList::new(),
1741 protected: SlruList::new(),
1742 free: Vec::new(),
1743 protected_cap,
1744 };
1745 let links = vec![SlotLink::none(); n];
1746 let old = OldSlru {
1747 probation: VecDeque::new(),
1748 protected: VecDeque::new(),
1749 protected_cap,
1750 };
1751 (class, links, old)
1752 }
1753
1754 fn orders_match(class: &SlotClass, links: &[SlotLink], old: &OldSlru) -> bool {
1755 let np: Vec<usize> = class.probation.iter(links).collect();
1756 let nt: Vec<usize> = class.protected.iter(links).collect();
1757 let op: Vec<usize> = old.probation.iter().copied().collect();
1758 let ot: Vec<usize> = old.protected.iter().copied().collect();
1759 np == op && nt == ot
1760 }
1761
1762 struct Rng(u64);
1764 impl Rng {
1765 fn next(&mut self) -> u64 {
1766 self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
1767 let mut z = self.0;
1768 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1769 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1770 z ^ (z >> 31)
1771 }
1772 fn below(&mut self, n: usize) -> usize {
1773 (self.next() % n as u64) as usize
1774 }
1775 }
1776
1777 #[test]
1778 fn same_eviction_decisions_randomized_soak() {
1779 for &(n, cap) in &[(8usize, 1usize), (16, 12), (64, 51), (128, 102)] {
1781 let (mut class, mut links, mut old) = new_pair(n, cap);
1782 let mut rng = Rng(0xC0FFEE ^ (n as u64) << 8 ^ cap as u64);
1783 let mut resident: Vec<usize> = Vec::new();
1784 let mut free: Vec<usize> = (0..n).rev().collect();
1785 for step in 0..200_000 {
1786 let op = rng.below(100);
1787 if op < 55 && !resident.is_empty() {
1788 let slot = resident[rng.below(resident.len())];
1790 class.on_hit_full(slot, &mut links);
1791 old.on_hit_full(slot);
1792 } else if op < 80 {
1793 let slot = if let Some(s) = free.pop() {
1795 s
1796 } else {
1797 let v_new = class.pop_lru(&mut links);
1798 let v_old = old.pop_lru();
1799 assert_eq!(
1800 v_new, v_old,
1801 "victim diverged at step {step} (n={n} cap={cap})"
1802 );
1803 let v = v_new.unwrap();
1804 resident.retain(|&s| s != v);
1805 v
1806 };
1807 class.probation.push_back(slot, SEG_PROBATION, &mut links);
1808 old.probation.push_back(slot);
1809 resident.push(slot);
1810 } else if op < 92 && resident.len() > 2 {
1811 let banned: Vec<usize> = (0..3.min(resident.len()))
1813 .map(|_| resident[rng.below(resident.len())])
1814 .collect();
1815 let take_new = {
1816 let q = &mut class.probation;
1817 let found = q.iter(&links).find(|s| !banned.contains(s));
1818 match found {
1819 Some(s) => {
1820 q.unlink(s, &mut links);
1821 Some(s)
1822 }
1823 None => {
1824 let q = &mut class.protected;
1825 q.iter(&links).find(|s| !banned.contains(s)).inspect(|&s| {
1826 q.unlink(s, &mut links);
1827 })
1828 }
1829 }
1830 };
1831 let take_old = old.take_excluding(&banned);
1832 assert_eq!(
1833 take_new, take_old,
1834 "excluding-victim diverged at step {step}"
1835 );
1836 if let Some(v) = take_new {
1837 resident.retain(|&s| s != v);
1838 free.push(v);
1839 }
1840 } else if !resident.is_empty() {
1841 let slot = resident[rng.below(resident.len())];
1843 match links[slot].seg {
1844 SEG_PROBATION => class.probation.unlink(slot, &mut links),
1845 SEG_PROTECTED => class.protected.unlink(slot, &mut links),
1846 _ => {}
1847 }
1848 if let Some(pos) = old.probation.iter().position(|&x| x == slot) {
1849 old.probation.remove(pos);
1850 } else if let Some(pos) = old.protected.iter().position(|&x| x == slot) {
1851 old.protected.remove(pos);
1852 }
1853 class.on_hit_full(slot, &mut links);
1854 old.on_hit_full(slot);
1855 }
1856 assert!(
1857 orders_match(&class, &links, &old),
1858 "segment order diverged at step {step} (n={n} cap={cap})"
1859 );
1860 }
1861 }
1862 }
1863
1864 #[test]
1865 fn hit_promotion_is_o1_not_on() {
1866 fn bench(n: usize, hits: usize) -> std::time::Duration {
1871 let (mut class, mut links, _) = new_pair(n, (n as f64 * 0.8) as usize);
1872 for s in 0..n {
1873 class.probation.push_back(s, SEG_PROBATION, &mut links);
1874 }
1875 let mut rng = Rng(0xBEEF);
1876 let t0 = std::time::Instant::now();
1877 for _ in 0..hits {
1878 class.on_hit_full(rng.below(n), &mut links);
1879 }
1880 t0.elapsed()
1881 }
1882 bench(1_000, 10_000);
1884 bench(46_000, 10_000);
1885 let small = bench(1_000, 850_000).as_secs_f64() / 850_000.0;
1886 let large = bench(46_000, 850_000).as_secs_f64() / 850_000.0;
1887 assert!(
1888 large < small * 8.0,
1889 "per-hit cost scaled with n_slots: {:.1}ns @1k vs {:.1}ns @46k",
1890 small * 1e9,
1891 large * 1e9
1892 );
1893 }
1894
1895 #[test]
1896 fn slru_list_basic_invariants() {
1897 let mut links = vec![SlotLink::none(); 4];
1898 let mut l = SlruList::new();
1899 assert_eq!(l.pop_front(&mut links), None);
1900 l.push_back(2, SEG_PROBATION, &mut links);
1901 l.push_back(0, SEG_PROBATION, &mut links);
1902 l.push_back(3, SEG_PROBATION, &mut links);
1903 assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 0, 3]);
1904 assert_eq!(l.len, 3);
1905 l.unlink(0, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2, 3]);
1907 l.unlink(3, &mut links); assert_eq!(l.iter(&links).collect::<Vec<_>>(), vec![2]);
1909 assert_eq!(l.pop_front(&mut links), Some(2)); assert_eq!(l.len, 0);
1911 assert_eq!(l.head, super::NIL);
1912 assert_eq!(l.tail, super::NIL);
1913 assert!(links.iter().all(|k| k.seg == super::SEG_NONE));
1914 }
1915}
1916
1917#[cfg(test)]
1918mod vram_fraction_tests {
1919 use super::{
1920 parse_cache_hard_vram_frac, parse_cache_lfu_decay, parse_cache_lfu_mtp_weight,
1921 size_class_plan,
1922 };
1923
1924 #[test]
1925 fn hard_vram_fraction_defaults_and_rejects_unsafe_values() {
1926 assert_eq!(parse_cache_hard_vram_frac(None), Ok(0.80));
1927 assert_eq!(parse_cache_hard_vram_frac(Some("0.82")), Ok(0.82));
1928 assert!(parse_cache_hard_vram_frac(Some("NaN")).is_err());
1929 assert_eq!(parse_cache_hard_vram_frac(Some("0.95")), Ok(0.95));
1930 assert!(parse_cache_hard_vram_frac(Some("0.96")).is_err());
1931 assert!(parse_cache_hard_vram_frac(Some("1.0")).is_err());
1932 assert!(parse_cache_hard_vram_frac(Some("bad")).is_err());
1933 }
1934
1935 #[test]
1936 fn lfu_decay_is_opt_in_and_bounded() {
1937 assert_eq!(parse_cache_lfu_decay(None), Ok(None));
1938 assert_eq!(parse_cache_lfu_decay(Some("0.8")), Ok(Some(0.8)));
1939 assert_eq!(parse_cache_lfu_decay(Some("1")), Ok(Some(1.0)));
1940 for value in ["0", "-0.1", "1.1", "NaN", "bad"] {
1941 assert!(
1942 parse_cache_lfu_decay(Some(value)).is_err(),
1943 "accepted {value}"
1944 );
1945 }
1946 }
1947
1948 #[test]
1949 fn lfu_mtp_weight_defaults_and_is_bounded() {
1950 assert_eq!(parse_cache_lfu_mtp_weight(None), Ok(1.0));
1951 assert_eq!(parse_cache_lfu_mtp_weight(Some("4")), Ok(4.0));
1952 for value in ["0", "0.1", "65", "NaN", "bad"] {
1953 assert!(
1954 parse_cache_lfu_mtp_weight(Some(value)).is_err(),
1955 "accepted {value}"
1956 );
1957 }
1958 }
1959
1960 #[test]
1961 fn size_class_plan_preserves_classes_and_never_exceeds_budget() {
1962 let blocks = [100usize, 100, 100, 200, 200, 400];
1963 let budget = (108 * 2) + 208 + 408;
1964 let plan = size_class_plan(&blocks, budget);
1965 assert!(plan.iter().all(|(_, count)| *count > 0));
1966 assert!(
1967 plan.iter()
1968 .map(|(bytes, count)| (bytes + 8) * count)
1969 .sum::<usize>()
1970 <= budget
1971 );
1972 assert!(plan.iter().all(|(bytes, count)| {
1973 *count <= blocks.iter().filter(|block| **block == *bytes).count()
1974 }));
1975 }
1976
1977 #[test]
1978 fn size_class_plan_returns_full_inventory_when_it_fits() {
1979 let blocks = [100usize, 100, 200, 400];
1980 let budget: usize = blocks.iter().map(|bytes| bytes + 8).sum();
1981 assert_eq!(
1982 size_class_plan(&blocks, budget),
1983 vec![(100, 2), (200, 1), (400, 1)]
1984 );
1985 }
1986
1987 #[test]
1988 fn size_class_plan_does_not_overflow_on_pathological_sizes() {
1989 let plan = size_class_plan(&[usize::MAX, usize::MAX], usize::MAX);
1990 assert!(plan.is_empty());
1991 }
1992}
1993
1994impl Drop for MoeSlotCache {
1995 fn drop(&mut self) {
1996 let mut safe_to_drop_slots = true;
1999 if self.compute_stream_unknown || !self.compute_sources.is_empty() {
2000 if let Err(err) = self.compute_stream.synchronize() {
2001 safe_to_drop_slots = false;
2002 eprintln!(
2003 "[moe-cache] unknown compute-stream drain failed ({err}); leaking GPU slots for safety"
2004 );
2005 for (_, keepalive) in self.compute_sources.drain() {
2006 std::mem::forget(keepalive);
2007 }
2008 } else {
2009 self.compute_stream_unknown = false;
2010 self.compute_sources.clear();
2011 }
2012 }
2013 let need_copy_drain = self.copy_stream_unknown
2014 || !self.pending.is_empty()
2015 || !self.inflight_sources.is_empty()
2016 || !self.quarantined_sources.is_empty();
2017 if need_copy_drain {
2018 if let Err(err) = self.copy_stream.synchronize() {
2019 safe_to_drop_slots = false;
2020 eprintln!(
2021 "[moe-cache] unknown copy-stream drain failed ({err}); leaking GPU slots for safety"
2022 );
2023 for (_, keepalive) in self.inflight_sources.drain(..) {
2024 std::mem::forget(keepalive);
2025 }
2026 for keepalive in self.quarantined_sources.drain(..) {
2027 std::mem::forget(keepalive);
2028 }
2029 for (_, pending) in self.pending.drain() {
2030 if let Some(keepalive) = pending.keepalive {
2031 std::mem::forget(keepalive);
2032 }
2033 }
2034 } else {
2035 self.copy_stream_unknown = false;
2036 self.inflight_sources.clear();
2037 self.quarantined_sources.clear();
2038 self.pending.clear();
2039 }
2040 }
2041 if let Some(pool) = self.pread.as_mut() {
2042 safe_to_drop_slots &= pool.drain();
2043 } else if self.pread_requested && self.pread_fallbacks != 0 {
2044 eprintln!(
2045 "[spill-pread] backend unavailable; mmap_fallbacks={}",
2046 self.pread_fallbacks
2047 );
2048 }
2049 if !safe_to_drop_slots {
2050 for slot in self.slots.drain(..) {
2051 std::mem::forget(slot);
2052 }
2053 }
2054 }
2055}