spark_model/lora/types_slots.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `LoraWeights` SLOT LIFECYCLE: acquiring, releasing, ref-counting and the
4//! LRU bookkeeping the victim search reads.
5//!
6//! Split from `types.rs` for the 500-LoC cap, which the file crossed when the
7//! GDN `out_proj` module joined `LoraModule`. Exact piecewise copy — no method
8//! changed in the move. The pool LAYOUT stays in `types.rs`; this is the half
9//! that mutates which adapter is resident.
10
11use std::sync::atomic::Ordering;
12
13use anyhow::Result;
14use spark_runtime::gpu::{DevicePtr, GpuBackend};
15
16use super::types::{LoraLayerWeights, LoraModule, LoraWeights, SlotView};
17
18impl LoraWeights {
19 /// Task #25: resolve `slot` (`>= 0` → that slot, `-1` → active) to a concrete
20 /// pool index and `+1` its ref_count, returning the RESOLVED index so the
21 /// caller can release EXACTLY that index later (immune to an intervening
22 /// rotate changing `active`). Returns `-1` — "nothing acquired" — when the
23 /// resolved index is out of range (bad request slot); the active slot is
24 /// always in range so `-1 -> active` never no-ops here for a loaded pool.
25 pub fn acquire_slot(&self, slot: i32) -> i32 {
26 let resolved = if slot >= 0 {
27 slot as usize
28 } else {
29 self.active
30 };
31 match self.ref_counts.get(resolved) {
32 Some(rc) => {
33 rc.fetch_add(1, Ordering::AcqRel);
34 // Task #27: stamp the RESOLVED slot as most-recently-used so the
35 // LRU victim policy ages the slot a request actually touched
36 // (including `-1 -> active`). Ticks are strictly increasing.
37 if let Some(lu) = self.last_used.get(resolved) {
38 let t = self.lru_tick.fetch_add(1, Ordering::Relaxed) + 1;
39 lu.store(t, Ordering::Relaxed);
40 }
41 resolved as i32
42 }
43 None => -1,
44 }
45 }
46
47 /// Task #27: stamp `slot` as most-recently-used WITHOUT taking a ref. Called
48 /// right after a promote so a freshly-staged (ref_count==0) slot is NOT the
49 /// immediate LRU victim of a back-to-back promote before its own request has
50 /// acquired — otherwise two distinct cold adapters promoted in quick
51 /// succession would collide on the same slot (the second evicting the first).
52 pub fn touch_slot(&self, slot: usize) {
53 if let Some(lu) = self.last_used.get(slot) {
54 let t = self.lru_tick.fetch_add(1, Ordering::Relaxed) + 1;
55 lu.store(t, Ordering::Relaxed);
56 }
57 }
58
59 /// Task #27: current LRU tick of pool `slot` (larger = more recently
60 /// acquired). Out-of-range → 0 (never used).
61 pub fn slot_last_used(&self, slot: usize) -> u64 {
62 self.last_used
63 .get(slot)
64 .map(|lu| lu.load(Ordering::Relaxed))
65 .unwrap_or(0)
66 }
67
68 /// Task #26: refresh `slot`'s cell in the `[max_loras]` a/b pointer tables +
69 /// the per-slot scale table from `layers` (the just-staged adapter's actual
70 /// per-module coverage). A re-staged adapter whose module coverage DIFFERS
71 /// from the evicted one would otherwise keep a STALE table entry: the
72 /// bgmv-routed path would SKIP a module the new adapter adds (`a_table[slot]`
73 /// stale-NULL → missed delta), keep applying an evicted module (stale non-NULL
74 /// → wrong delta), or use the wrong per-slot scale. Shared by BOTH the disk
75 /// swap (`pack_store_into_slot`) and the RDMA swap (`swap_lora_slot_from_peer`).
76 /// Only the `[slot]` cell of each fixed-address device array is rewritten.
77 pub fn refresh_slot_tables(
78 &self,
79 slot: usize,
80 layers: &[Option<LoraLayerWeights>],
81 scale: f32,
82 gpu: &dyn GpuBackend,
83 ) -> Result<()> {
84 for ((layer, module), (a_dev, b_dev)) in &self.tables {
85 let pair = layers
86 .get(*layer)
87 .and_then(|o| o.as_ref())
88 .and_then(|lw| match module {
89 LoraModule::QProj => lw.q_proj.as_ref(),
90 LoraModule::KProj => lw.k_proj.as_ref(),
91 LoraModule::VProj => lw.v_proj.as_ref(),
92 LoraModule::OProj => lw.o_proj.as_ref(),
93 LoraModule::GateProj => lw.gate_proj.as_ref(),
94 LoraModule::UpProj => lw.up_proj.as_ref(),
95 LoraModule::DownProj => lw.down_proj.as_ref(),
96 LoraModule::OutProj => lw.out_proj.as_ref(),
97 });
98 let (a_ptr, b_ptr) = pair.map(|p| (p.a.weight.0, p.b.weight.0)).unwrap_or((0, 0));
99 gpu.copy_h2d(&a_ptr.to_le_bytes(), DevicePtr(a_dev.0 + (slot * 8) as u64))?;
100 gpu.copy_h2d(&b_ptr.to_le_bytes(), DevicePtr(b_dev.0 + (slot * 8) as u64))?;
101 }
102 if self.scale_table.0 != 0 {
103 gpu.copy_h2d(
104 &scale.to_le_bytes(),
105 DevicePtr(self.scale_table.0 + (slot * 4) as u64),
106 )?;
107 }
108 Ok(())
109 }
110
111 /// Task #27: snapshot the CACHE region `[pinned, max_loras)` as
112 /// `(slot_index, SlotView)` for `select_victim_slot`. `filled` = the slot
113 /// holds a non-placeholder adapter (non-empty name). Read on the model
114 /// thread at a quiescent point.
115 pub fn cache_slot_views(&self) -> Vec<(usize, SlotView)> {
116 (self.pinned..self.max_loras)
117 .map(|k| {
118 let filled = self.slots.get(k).is_some_and(|s| !s.name.is_empty());
119 (
120 k,
121 SlotView {
122 filled,
123 ref_count: self.slot_ref_count(k),
124 last_used: self.slot_last_used(k),
125 },
126 )
127 })
128 .collect()
129 }
130
131 /// Task #25: release a ref previously taken by [`Self::acquire_slot`], by the
132 /// RESOLVED index it returned. `-1` (nothing acquired) is a no-op. Saturating
133 /// so a stray double-release can never wrap the counter below 0.
134 pub fn release_slot(&self, resolved: i32) {
135 if resolved < 0 {
136 return;
137 }
138 if let Some(rc) = self.ref_counts.get(resolved as usize) {
139 let _ = rc.fetch_update(Ordering::Release, Ordering::Acquire, |v| {
140 Some(v.saturating_sub(1))
141 });
142 }
143 }
144
145 /// Task #25: current in-flight ref_count of pool `slot` (the exact read the
146 /// swap busy-slot gate branches on). Out-of-range → 0.
147 pub fn slot_ref_count(&self, slot: usize) -> usize {
148 self.ref_counts
149 .get(slot)
150 .map(|rc| rc.load(Ordering::Acquire))
151 .unwrap_or(0)
152 }
153}