metrale_storage/
high_speed_swap.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `HighSpeedSwap`, the high-speed-swap orchestrator. It writes KV
4//! blocks to per-layer files (`offload_block*`) and streams a sequence's blocks
5//! back through a scratch pool into tiled attention (`attend_layer*`), using the
6//! predictor, the file backend and the eviction policy. It also keeps the disk
7//! block ids and the per-thread slot (`install_local`, `with_local`) through
8//! which model code reaches it.
9//!
10//! Owner: metrale-storage high-speed swap.
11//! Invariants:
12//! - `alloc_disk_block_id` reuses ids from `free_list` before it issues a new
13//!   id, and never issues one at or above `max_blocks_per_layer`.
14
15use anyhow::{Context, Result};
16
17// 2026-09-25: io_uring on Linux, the portable backend elsewhere; the alias
18// keeps one orchestrator body for both.
19#[cfg(target_os = "linux")]
20use crate::backend::IoUringBackend as TierBackend;
21#[cfg(not(target_os = "linux"))]
22use crate::backend::PosixBackend as TierBackend;
23use crate::config::HighSpeedSwapConfig;
24use crate::cuda_min::{CudaCtx, DeviceBuffer};
25use crate::eviction::EvictionPolicy;
26use crate::group::GroupLayout;
27use crate::layout::Layout;
28use crate::predictor::{Predictor, PredictorDims};
29use crate::scratch_pool::{ScratchDims, ScratchPool};
30use crate::tiled_attention::{TiledAttention, TiledAttentionDims};
31
32// 2026-09-25: `ModelDims` is in `crate::model_dims`, which builds without the
33// `cuda` feature; this module does not.
34pub use crate::model_dims::ModelDims;
35
36pub struct HighSpeedSwap {
37    cfg: HighSpeedSwapConfig,
38    model: ModelDims,
39    predictor: Predictor,
40    pool: ScratchPool,
41    backend: TierBackend,
42    attn: TiledAttention,
43    eviction: EvictionPolicy,
44    q_proj: DeviceBuffer,
45    // 2026-09-25: f32 scores `[max_blocks_per_layer]`, the tile's i32 block
46    // table `[resident_blocks]`, and the i32 block count of the one sequence.
47    block_scores_dev: DeviceBuffer,
48    block_table_dev: DeviceBuffer,
49    counts_dev: DeviceBuffer,
50    score_host_buf: Vec<f32>,
51    // 2026-09-25: One disk-block-id allocator for all layers: an id names the
52    // same block position in every layer's file.
53    disk_state: DiskState,
54}
55
56#[derive(Debug)]
57struct DiskState {
58    next_id: u32,
59    free_list: Vec<u32>,
60    refcount: Vec<u32>,
61}
62
63impl DiskState {
64    fn new() -> Self {
65        Self {
66            next_id: 0,
67            free_list: Vec::new(),
68            refcount: Vec::new(),
69        }
70    }
71}
72
73impl HighSpeedSwap {
74    pub fn new(ctx: &CudaCtx, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<Self> {
75        Self::new_on_stream(ctx.stream, cfg, model)
76    }
77
78    /// 2026-09-25: Build on `stream`, which must belong to the current thread's
79    /// CUDA context. The stream carries only the construction-time upload of the
80    /// predictor's projection; per-step methods take their own stream.
81    pub fn new_on_stream(stream: u64, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<Self> {
82        cfg.validate_and_prepare()?;
83        let group_layout = GroupLayout::new(
84            model.num_layers,
85            model.max_blocks_per_layer,
86            model.num_kv_heads,
87            model.block_size as u32,
88            model.head_dim as u32,
89            2,
90            4096,
91        );
92        let layout = Layout::create(&cfg.dir, group_layout).context("create layout")?;
93        // 2026-09-25: The portable backend has one bounce buffer and no queue
94        // depth.
95        #[cfg(target_os = "linux")]
96        let backend = TierBackend::new(layout, cfg.qd as usize)?;
97        #[cfg(not(target_os = "linux"))]
98        let backend = TierBackend::new(layout)?;
99        let pool = ScratchPool::new(ScratchDims {
100            num_slots: cfg.resident_blocks,
101            num_kv_heads: model.num_kv_heads,
102            group_stride: group_layout.group_stride,
103        })?;
104        let predictor = Predictor::new_on_stream(
105            stream,
106            PredictorDims {
107                num_layers: model.num_layers as usize,
108                num_q_heads: model.num_q_heads as usize,
109                num_kv_heads: model.num_kv_heads as usize,
110                head_dim: model.head_dim as usize,
111                r: cfg.rank as usize,
112                block_size: model.block_size as usize,
113                max_blocks: model.max_blocks_per_layer as usize,
114            },
115            cfg.projection_seed,
116        )?;
117        let attn = TiledAttention::new(TiledAttentionDims {
118            max_seqs: 1,
119            num_q_heads: model.num_q_heads as usize,
120            num_kv_heads: model.num_kv_heads as usize,
121            head_dim: model.head_dim as usize,
122            block_size: model.block_size as usize,
123            tile_capacity: cfg.resident_blocks as usize,
124        })?;
125        let eviction = EvictionPolicy::new(cfg.resident_blocks);
126        let q_proj = DeviceBuffer::new(model.num_q_heads as usize * cfg.rank as usize * 2)?;
127        let block_scores_dev = DeviceBuffer::new(model.max_blocks_per_layer as usize * 4)?;
128        let block_table_dev = DeviceBuffer::new(cfg.resident_blocks as usize * 4)?;
129        let counts_dev = DeviceBuffer::new(4)?;
130        let score_host_buf = vec![0.0_f32; model.max_blocks_per_layer as usize];
131        let disk_state = DiskState::new();
132        Ok(Self {
133            cfg,
134            model,
135            predictor,
136            pool,
137            backend,
138            attn,
139            eviction,
140            q_proj,
141            block_scores_dev,
142            block_table_dev,
143            counts_dev,
144            score_host_buf,
145            disk_state,
146        })
147    }
148
149    // 2026-09-25: Disk block ids, shared by all layers, at most
150    // `max_blocks_per_layer` of them. `alloc_disk_block_id` returns `None` when
151    // none is free; `inc_disk_ref` panics on a freed id; `dec_disk_ref` returns
152    // the new count and frees the id at 0.
153
154    pub fn alloc_disk_block_id(&mut self) -> Option<u32> {
155        let st = &mut self.disk_state;
156        if let Some(id) = st.free_list.pop() {
157            st.refcount[id as usize] = 1;
158            return Some(id);
159        }
160        if st.next_id >= self.model.max_blocks_per_layer {
161            return None;
162        }
163        let id = st.next_id;
164        st.next_id += 1;
165        st.refcount.push(1);
166        Some(id)
167    }
168
169    pub fn inc_disk_ref(&mut self, id: u32) {
170        let rc = &mut self.disk_state.refcount[id as usize];
171        if *rc == 0 {
172            panic!("inc_disk_ref on freed disk_block_id {id}; caller must hold a live ref");
173        }
174        *rc += 1;
175    }
176
177    pub fn dec_disk_ref(&mut self, id: u32) -> u32 {
178        let st = &mut self.disk_state;
179        let rc = &mut st.refcount[id as usize];
180        debug_assert!(*rc > 0, "dec_disk_ref on already-freed id {id}");
181        *rc = rc.saturating_sub(1);
182        let new_rc = *rc;
183        if new_rc == 0 {
184            st.free_list.push(id);
185        }
186        new_rc
187    }
188
189    pub fn disk_refcount(&self, id: u32) -> u32 {
190        self.disk_state.refcount[id as usize]
191    }
192
193    pub fn disk_free_count(&self) -> usize {
194        let st = &self.disk_state;
195        st.free_list.len() + (self.model.max_blocks_per_layer - st.next_id) as usize
196    }
197
198    /// 2026-09-25: Disk-id and scratch-pool occupancy.
199    pub fn diagnostic_summary(&self) -> HighSpeedSwapDiagnostic {
200        let st = &self.disk_state;
201        let active = st.next_id.saturating_sub(st.free_list.len() as u32);
202        HighSpeedSwapDiagnostic {
203            num_layers: self.model.num_layers,
204            active_disk_blocks: active,
205            disk_block_capacity: self.model.max_blocks_per_layer,
206            scratch_pool_resident: self.pool.dims().num_slots,
207            scratch_pool_free: self.pool.free_count(),
208        }
209    }
210}
211
212#[derive(Debug, Clone, Copy)]
213pub struct HighSpeedSwapDiagnostic {
214    pub num_layers: u32,
215    pub active_disk_blocks: u32,
216    pub disk_block_capacity: u32,
217    pub scratch_pool_resident: u32,
218    pub scratch_pool_free: u32,
219}
220
221#[cfg(test)]
222mod disk_id_tests;
223
224mod impl_more;
225
226use std::cell::RefCell;
227// 2026-09-25: Thread-local because the orchestrator's device memory and
228// streams belong to one thread's CUDA context. The server installs it
229// (`install_local`), and model-layer and model-engine code reach it with
230// `with_local` instead of passing it through every layer signature. It is
231// dropped when the thread exits or another install replaces it.
232thread_local! {
233    static LOCAL: RefCell<Option<HighSpeedSwap>> = const { RefCell::new(None) };
234}
235
236/// 2026-09-25: Build an orchestrator and install it on the current thread,
237/// dropping any earlier one. On a build error the earlier one stays.
238pub fn install_local(stream: u64, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<()> {
239    let hss = HighSpeedSwap::new_on_stream(stream, cfg, model)?;
240    LOCAL.with(|cell| {
241        *cell.borrow_mut() = Some(hss);
242    });
243    Ok(())
244}
245
246/// 2026-09-25: Whether this thread has an installed orchestrator.
247pub fn local_installed() -> bool {
248    LOCAL.with(|cell| cell.borrow().is_some())
249}
250
251/// 2026-09-25: Run `f` on this thread's orchestrator; `None` when none is
252/// installed.
253pub fn with_local<R>(f: impl FnOnce(&mut HighSpeedSwap) -> Result<R>) -> Option<Result<R>> {
254    LOCAL.with(|cell| cell.borrow_mut().as_mut().map(f))
255}
metrale_storage/high_speed_swap.rs

metrale_storage/
high_speed_swap.rs
