metrale_storage/
scratch_pool.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Device scratch pool of the high-speed-swap tier: slots that hold one
4//! block's K and V for every KV head, filled from the storage backend.
5//!
6//! Owner: storage, high-speed swap.
7//! Invariants:
8//! - `lookup[k] == s` exactly when `residents[s] == Some(k)`.
9//! - A slot is in `free_list` only while `residents` holds `None` for it.
10//!
11//! Slot `S` is `slot_bytes = 2 * num_kv_heads * group_stride` bytes. The K stripe of
12//! head `h` starts at `S * slot_bytes + h * group_stride` and its V stripe at
13//! `S * slot_bytes + (num_kv_heads + h) * group_stride`. That is the layout
14//! `TiledAttention::scratch_pool_strides` describes when `group_stride` is exactly
15//! `block_size * head_dim * 2` bytes, with no filesystem padding.
16
17use anyhow::{Result, bail};
18use std::collections::{HashMap, VecDeque};
19
20use crate::cuda_min::DeviceBuffer;
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
23pub struct ResidentKey {
24    pub layer: u32,
25    pub block: u32,
26}
27
28#[derive(Clone, Copy, Debug)]
29pub struct ScratchDims {
30    pub num_slots: u32,
31    pub num_kv_heads: u16,
32    // 2026-09-25: Bytes per (block, kv_head) stripe.
33    pub group_stride: u64,
34}
35
36impl ScratchDims {
37    pub fn slot_bytes(&self) -> usize {
38        (2 * self.num_kv_heads as u64 * self.group_stride) as usize
39    }
40    pub fn pool_bytes(&self) -> usize {
41        self.num_slots as usize * self.slot_bytes()
42    }
43}
44
45pub struct ScratchPool {
46    dims: ScratchDims,
47    pool: DeviceBuffer,
48    residents: Vec<Option<ResidentKey>>,
49    lookup: HashMap<ResidentKey, u32>,
50    free_list: VecDeque<u32>,
51}
52
53impl ScratchPool {
54    pub fn new(dims: ScratchDims) -> Result<Self> {
55        if dims.num_slots == 0 {
56            bail!("ScratchPool requires at least one slot");
57        }
58        let pool = DeviceBuffer::new(dims.pool_bytes())?;
59        let residents = vec![None; dims.num_slots as usize];
60        let free_list = (0..dims.num_slots).collect();
61        Ok(Self {
62            dims,
63            pool,
64            residents,
65            lookup: HashMap::new(),
66            free_list,
67        })
68    }
69
70    pub fn dims(&self) -> ScratchDims {
71        self.dims
72    }
73    pub fn pool_dev_ptr(&self) -> u64 {
74        self.pool.ptr
75    }
76    pub fn slot_dev_ptr(&self, slot: u32) -> u64 {
77        self.pool.ptr + (slot as u64) * (self.dims.slot_bytes() as u64)
78    }
79    pub fn slot_k_ptr(&self, slot: u32, kv_head: u16) -> u64 {
80        self.slot_dev_ptr(slot) + (kv_head as u64) * self.dims.group_stride
81    }
82    pub fn slot_v_ptr(&self, slot: u32, kv_head: u16) -> u64 {
83        self.slot_dev_ptr(slot)
84            + (self.dims.num_kv_heads as u64 + kv_head as u64) * self.dims.group_stride
85    }
86
87    pub fn lookup(&self, key: ResidentKey) -> Option<u32> {
88        self.lookup.get(&key).copied()
89    }
90
91    /// 2026-09-25: Free `key`'s slot, if any, so the next `assign` of `key` needs a
92    /// fresh read. The offload path calls it after rewriting a block on disk, so
93    /// attention does not read the old resident copy.
94    pub fn invalidate(&mut self, key: ResidentKey) {
95        if let Some(slot) = self.lookup.remove(&key) {
96            self.residents[slot as usize] = None;
97            self.free_list.push_back(slot);
98        }
99    }
100
101    pub fn capacity(&self) -> u32 {
102        self.dims.num_slots
103    }
104    pub fn free_count(&self) -> u32 {
105        self.free_list.len() as u32
106    }
107
108    /// 2026-09-25: Slot for `key`: its current slot if resident, else the first free
109    /// slot, else the first resident slot in `evict_candidates`, whose key is evicted.
110    /// With no free slot and no resident candidate, an error. The caller fills the
111    /// slot.
112    pub fn assign(&mut self, key: ResidentKey, evict_candidates: &[u32]) -> Result<u32> {
113        if let Some(&slot) = self.lookup.get(&key) {
114            return Ok(slot);
115        }
116        let slot = match self.free_list.pop_front() {
117            Some(s) => s,
118            None => {
119                // 2026-09-25: Pinned slots are the caller's to leave out of
120                // `evict_candidates`.
121                let mut chosen = None;
122                for &c in evict_candidates {
123                    if self
124                        .residents
125                        .get(c as usize)
126                        .and_then(|r| r.as_ref())
127                        .is_some()
128                    {
129                        chosen = Some(c);
130                        break;
131                    }
132                }
133                let s = chosen.ok_or_else(|| {
134                    anyhow::anyhow!("no slot available and no eviction candidate is resident")
135                })?;
136                if let Some(prev) = self.residents[s as usize].take() {
137                    self.lookup.remove(&prev);
138                }
139                s
140            }
141        };
142        self.residents[slot as usize] = Some(key);
143        self.lookup.insert(key, slot);
144        Ok(slot)
145    }
146
147    /// 2026-09-25: Every resident slot with its key, in slot order.
148    pub fn residents(&self) -> Vec<(u32, ResidentKey)> {
149        self.residents
150            .iter()
151            .enumerate()
152            .filter_map(|(i, r)| r.map(|k| (i as u32, k)))
153            .collect()
154    }
155
156    pub fn clear(&mut self) {
157        self.lookup.clear();
158        for r in self.residents.iter_mut() {
159            *r = None;
160        }
161        self.free_list.clear();
162        for s in 0..self.dims.num_slots {
163            self.free_list.push_back(s);
164        }
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    fn ctx() -> crate::cuda_min::CudaCtx {
173        crate::cuda_min::CudaCtx::new(0).expect("cuda init")
174    }
175
176    #[test]
177    #[ignore = "requires GPU"]
178    fn assign_and_lookup() {
179        let _ctx = ctx();
180        let mut pool = ScratchPool::new(ScratchDims {
181            num_slots: 4,
182            num_kv_heads: 2,
183            group_stride: 4096,
184        })
185        .unwrap();
186        let k0 = ResidentKey { layer: 0, block: 7 };
187        let s0 = pool.assign(k0, &[]).unwrap();
188        assert_eq!(pool.lookup(k0), Some(s0));
189        let s0_again = pool.assign(k0, &[]).unwrap();
190        assert_eq!(s0, s0_again);
191        for b in 8..11 {
192            pool.assign(ResidentKey { layer: 0, block: b }, &[])
193                .unwrap();
194        }
195        assert_eq!(pool.free_count(), 0);
196        let evicted = pool
197            .assign(
198                ResidentKey {
199                    layer: 0,
200                    block: 99,
201                },
202                &[s0],
203            )
204            .unwrap();
205        assert_eq!(evicted, s0);
206        assert_eq!(pool.lookup(k0), None);
207    }
208
209    #[test]
210    #[ignore = "requires GPU"]
211    fn slot_pointer_layout() {
212        let _ctx = ctx();
213        let pool = ScratchPool::new(ScratchDims {
214            num_slots: 2,
215            num_kv_heads: 4,
216            group_stride: 4096,
217        })
218        .unwrap();
219        let base = pool.pool_dev_ptr();
220        assert_eq!(pool.slot_dev_ptr(0), base);
221        assert_eq!(pool.slot_dev_ptr(1), base + 8 * 4096);
222        assert_eq!(pool.slot_k_ptr(0, 2), base + 2 * 4096);
223        assert_eq!(pool.slot_v_ptr(0, 2), base + (4 + 2) * 4096);
224    }
225}
metrale_storage/scratch_pool.rs

metrale_storage/
scratch_pool.rs
