spark_storage/
scratch_pool.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// HBM scratch pool for the high-speed-swap path. Holds N "slots", each large
4// enough for one block worth of K + V across all kv_heads:
5//
6//   slot_bytes = 2 * num_kv_heads * group_stride
7//
8// The pool is laid out so that, for a given slot S, the K bytes for kv_head
9// `h` start at `pool_base + S*slot_bytes + h*group_stride`, and V bytes start
10// at `pool_base + S*slot_bytes + (num_kv_heads + h)*group_stride`. This
11// matches the BHND layout the tiled-attention kernel expects when treating
12// the scratch pool itself as the K/V "block pool" and using slot indices as
13// block IDs.
14//
15// The pool maintains:
16//   - A `Vec<Option<(u32 layer, u32 block)>>` of slot residents.
17//   - A free-list of available slot indices.
18//   - A `HashMap<(layer, block), slot_idx>` for lookup.
19//
20// Phase 2 keeps the API intentionally simple: `assign(layer, block)` returns
21// a slot index for a fresh block, evicting the head of the free-list (or, if
22// empty, the resident with the lowest predictor score — see eviction.rs).
23// No epoch counters yet; threading and fence safety arrive in Phase 3 when
24// the I/O thread lives on a separate stream.
25
26use anyhow::{Result, bail};
27use std::collections::{HashMap, VecDeque};
28
29use crate::cuda_min::DeviceBuffer;
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
32pub struct ResidentKey {
33    pub layer: u32,
34    pub block: u32,
35}
36
37#[derive(Clone, Copy, Debug)]
38pub struct ScratchDims {
39    pub num_slots: u32,
40    pub num_kv_heads: u16,
41    pub group_stride: u64, // bytes per (block, kv_head) stripe
42}
43
44impl ScratchDims {
45    pub fn slot_bytes(&self) -> usize {
46        (2 * self.num_kv_heads as u64 * self.group_stride) as usize
47    }
48    pub fn pool_bytes(&self) -> usize {
49        self.num_slots as usize * self.slot_bytes()
50    }
51}
52
53pub struct ScratchPool {
54    dims: ScratchDims,
55    pool: DeviceBuffer,
56    residents: Vec<Option<ResidentKey>>, // indexed by slot idx
57    lookup: HashMap<ResidentKey, u32>,
58    free_list: VecDeque<u32>,
59}
60
61impl ScratchPool {
62    pub fn new(dims: ScratchDims) -> Result<Self> {
63        if dims.num_slots == 0 {
64            bail!("ScratchPool requires at least one slot");
65        }
66        let pool = DeviceBuffer::new(dims.pool_bytes())?;
67        let residents = vec![None; dims.num_slots as usize];
68        let free_list = (0..dims.num_slots).collect();
69        Ok(Self {
70            dims,
71            pool,
72            residents,
73            lookup: HashMap::new(),
74            free_list,
75        })
76    }
77
78    pub fn dims(&self) -> ScratchDims {
79        self.dims
80    }
81    pub fn pool_dev_ptr(&self) -> u64 {
82        self.pool.ptr
83    }
84    pub fn slot_dev_ptr(&self, slot: u32) -> u64 {
85        self.pool.ptr + (slot as u64) * (self.dims.slot_bytes() as u64)
86    }
87    /// K stripe device pointer for (slot, kv_head).
88    pub fn slot_k_ptr(&self, slot: u32, kv_head: u16) -> u64 {
89        self.slot_dev_ptr(slot) + (kv_head as u64) * self.dims.group_stride
90    }
91    /// V stripe device pointer for (slot, kv_head).
92    pub fn slot_v_ptr(&self, slot: u32, kv_head: u16) -> u64 {
93        self.slot_dev_ptr(slot)
94            + (self.dims.num_kv_heads as u64 + kv_head as u64) * self.dims.group_stride
95    }
96
97    pub fn lookup(&self, key: ResidentKey) -> Option<u32> {
98        self.lookup.get(&key).copied()
99    }
100
101    /// Drop the resident slot for `key` (if any). Returns the slot to the
102    /// free list so the next `assign(key, _)` triggers a fresh disk read.
103    ///
104    /// Used by the offload path to discard the cached copy of a block after
105    /// its on-disk image has been overwritten — without this, streaming
106    /// attention would keep serving the stale resident copy and never see
107    /// the freshly-offloaded K/V (e.g., decode steps re-writing the active
108    /// block every step).
109    pub fn invalidate(&mut self, key: ResidentKey) {
110        if let Some(slot) = self.lookup.remove(&key) {
111            self.residents[slot as usize] = None;
112            self.free_list.push_back(slot);
113        }
114    }
115
116    pub fn capacity(&self) -> u32 {
117        self.dims.num_slots
118    }
119    pub fn free_count(&self) -> u32 {
120        self.free_list.len() as u32
121    }
122
123    /// Reserve a slot for `key`. If the pool is full, picks an evictable slot
124    /// from `evict_candidates` (callers pass them in score-ascending order;
125    /// the lowest-scoring one is kicked first). Returns the slot index. The
126    /// caller is responsible for issuing the disk read into `slot_dev_ptr`.
127    pub fn assign(&mut self, key: ResidentKey, evict_candidates: &[u32]) -> Result<u32> {
128        if let Some(&slot) = self.lookup.get(&key) {
129            return Ok(slot); // already resident
130        }
131        let slot = match self.free_list.pop_front() {
132            Some(s) => s,
133            None => {
134                // Find the first candidate that is currently resident (still
135                // backed by a known key) and is not pinned.
136                let mut chosen = None;
137                for &c in evict_candidates {
138                    if self
139                        .residents
140                        .get(c as usize)
141                        .and_then(|r| r.as_ref())
142                        .is_some()
143                    {
144                        chosen = Some(c);
145                        break;
146                    }
147                }
148                let s = chosen.ok_or_else(|| {
149                    anyhow::anyhow!("no slot available and no eviction candidate is resident")
150                })?;
151                if let Some(prev) = self.residents[s as usize].take() {
152                    self.lookup.remove(&prev);
153                }
154                s
155            }
156        };
157        self.residents[slot as usize] = Some(key);
158        self.lookup.insert(key, slot);
159        Ok(slot)
160    }
161
162    /// Return all currently-resident slot indices in arbitrary order (for the
163    /// eviction policy to score against the predictor).
164    pub fn residents(&self) -> Vec<(u32, ResidentKey)> {
165        self.residents
166            .iter()
167            .enumerate()
168            .filter_map(|(i, r)| r.map(|k| (i as u32, k)))
169            .collect()
170    }
171
172    /// Free all slots; use between decode steps that don't share residency.
173    pub fn clear(&mut self) {
174        self.lookup.clear();
175        for r in self.residents.iter_mut() {
176            *r = None;
177        }
178        self.free_list.clear();
179        for s in 0..self.dims.num_slots {
180            self.free_list.push_back(s);
181        }
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    fn ctx() -> crate::cuda_min::CudaCtx {
190        crate::cuda_min::CudaCtx::new(0).expect("cuda init")
191    }
192
193    #[test]
194    #[ignore = "requires GPU"]
195    fn assign_and_lookup() {
196        let _ctx = ctx();
197        let mut pool = ScratchPool::new(ScratchDims {
198            num_slots: 4,
199            num_kv_heads: 2,
200            group_stride: 4096,
201        })
202        .unwrap();
203        let k0 = ResidentKey { layer: 0, block: 7 };
204        let s0 = pool.assign(k0, &[]).unwrap();
205        assert_eq!(pool.lookup(k0), Some(s0));
206        // Repeated assign returns the same slot.
207        let s0_again = pool.assign(k0, &[]).unwrap();
208        assert_eq!(s0, s0_again);
209        // Fill the pool.
210        for b in 8..11 {
211            pool.assign(ResidentKey { layer: 0, block: b }, &[])
212                .unwrap();
213        }
214        assert_eq!(pool.free_count(), 0);
215        // Evict the lowest-scoring slot (caller passes it in).
216        let evicted = pool
217            .assign(
218                ResidentKey {
219                    layer: 0,
220                    block: 99,
221                },
222                &[s0],
223            )
224            .unwrap();
225        assert_eq!(evicted, s0); // s0 was the eviction candidate
226        assert_eq!(pool.lookup(k0), None); // k0 displaced
227    }
228
229    #[test]
230    #[ignore = "requires GPU"]
231    fn slot_pointer_layout() {
232        let _ctx = ctx();
233        let pool = ScratchPool::new(ScratchDims {
234            num_slots: 2,
235            num_kv_heads: 4,
236            group_stride: 4096,
237        })
238        .unwrap();
239        let base = pool.pool_dev_ptr();
240        assert_eq!(pool.slot_dev_ptr(0), base);
241        assert_eq!(pool.slot_dev_ptr(1), base + 8 * 4096);
242        assert_eq!(pool.slot_k_ptr(0, 2), base + 2 * 4096);
243        assert_eq!(pool.slot_v_ptr(0, 2), base + (4 + 2) * 4096);
244    }
245}
spark_storage/scratch_pool.rs

spark_storage/
scratch_pool.rs
