spark_storage/
high_speed_swap.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// `HighSpeedSwap` orchestrator: combines Predictor + ScratchPool +
4// IoUringBackend + TiledAttention + EvictionPolicy behind a two-method API
5// (`offload_block`, `attend_layer`). Designed to be the primitive a future
6// scheduler integration in `spark-model` plugs into.
7
8use anyhow::{Context, Result};
9
10// The tier's backend is io_uring on Linux and the portable positional-I/O
11// backend everywhere else. Selected by alias so the orchestrator below has ONE
12// body: only the submission mechanism differs, not the layout, the eviction
13// policy or the scratch pool.
14#[cfg(target_os = "linux")]
15use crate::backend::IoUringBackend as TierBackend;
16#[cfg(not(target_os = "linux"))]
17use crate::backend::PosixBackend as TierBackend;
18use crate::config::HighSpeedSwapConfig;
19use crate::cuda_min::{CudaCtx, DeviceBuffer};
20use crate::eviction::EvictionPolicy;
21use crate::group::GroupLayout;
22use crate::layout::Layout;
23use crate::predictor::{Predictor, PredictorDims};
24use crate::scratch_pool::{ScratchDims, ScratchPool};
25use crate::tiled_attention::{TiledAttention, TiledAttentionDims};
26
27// `ModelDims` lives in `crate::model_dims` so it stays available on
28// non-cuda builds where the swap orchestrator below isn't compiled.
29pub use crate::model_dims::ModelDims;
30
31pub struct HighSpeedSwap {
32    cfg: HighSpeedSwapConfig,
33    model: ModelDims,
34    predictor: Predictor,
35    pool: ScratchPool,
36    backend: TierBackend,
37    attn: TiledAttention,
38    eviction: EvictionPolicy,
39    // Reusable scratch buffers.
40    q_proj: DeviceBuffer,
41    block_scores_dev: DeviceBuffer, // [max_blocks] f32
42    block_table_dev: DeviceBuffer,  // [tile_capacity] i32
43    counts_dev: DeviceBuffer,       // [1] i32 (single seq)
44    score_host_buf: Vec<f32>,
45    // Disk-block-ID allocator (Phase 6.1.a, refactored). One global
46    // allocator: a `disk_block_id` indexes the SAME logical position
47    // across every layer's file, so allocation, refcount, and free list
48    // are layer-agnostic. Each layer's file independently stores its
49    // K/V at `offset(layer, disk_block_id)`.
50    disk_state: DiskState,
51}
52
53#[derive(Debug)]
54struct DiskState {
55    next_id: u32,
56    free_list: Vec<u32>,
57    refcount: Vec<u32>,
58}
59
60impl DiskState {
61    fn new() -> Self {
62        Self {
63            next_id: 0,
64            free_list: Vec::new(),
65            refcount: Vec::new(),
66        }
67    }
68}
69
70impl HighSpeedSwap {
71    pub fn new(ctx: &CudaCtx, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<Self> {
72        Self::new_on_stream(ctx.stream, cfg, model)
73    }
74
75    /// Stream-only constructor for production callers that already own a
76    /// CUDA context (spark-model). The provided `stream` is used only for
77    /// init-time copies (uploading the projection matrix P); subsequent
78    /// per-step calls take their own stream argument.
79    pub fn new_on_stream(stream: u64, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<Self> {
80        cfg.validate_and_prepare()?;
81        let group_layout = GroupLayout::new(
82            model.num_layers,
83            model.max_blocks_per_layer,
84            model.num_kv_heads,
85            model.block_size as u32,
86            model.head_dim as u32,
87            2, // BF16
88            4096,
89        );
90        let layout = Layout::create(&cfg.dir, group_layout).context("create layout")?;
91        // qd is the io_uring submission-queue depth; the portable backend
92        // serialises on a single pinned bounce buffer and has no queue.
93        #[cfg(target_os = "linux")]
94        let backend = TierBackend::new(layout, cfg.qd as usize)?;
95        #[cfg(not(target_os = "linux"))]
96        let backend = TierBackend::new(layout)?;
97        let pool = ScratchPool::new(ScratchDims {
98            num_slots: cfg.resident_blocks,
99            num_kv_heads: model.num_kv_heads,
100            group_stride: group_layout.group_stride,
101        })?;
102        let predictor = Predictor::new_on_stream(
103            stream,
104            PredictorDims {
105                num_layers: model.num_layers as usize,
106                num_q_heads: model.num_q_heads as usize,
107                num_kv_heads: model.num_kv_heads as usize,
108                head_dim: model.head_dim as usize,
109                r: cfg.rank as usize,
110                block_size: model.block_size as usize,
111                max_blocks: model.max_blocks_per_layer as usize,
112            },
113            cfg.projection_seed,
114        )?;
115        let attn = TiledAttention::new(TiledAttentionDims {
116            max_seqs: 1, // single-seq for the orchestrator's first iteration
117            num_q_heads: model.num_q_heads as usize,
118            num_kv_heads: model.num_kv_heads as usize,
119            head_dim: model.head_dim as usize,
120            block_size: model.block_size as usize,
121            tile_capacity: cfg.resident_blocks as usize,
122        })?;
123        let eviction = EvictionPolicy::new(cfg.resident_blocks);
124        let q_proj = DeviceBuffer::new(model.num_q_heads as usize * cfg.rank as usize * 2)?;
125        let block_scores_dev = DeviceBuffer::new(model.max_blocks_per_layer as usize * 4)?;
126        let block_table_dev = DeviceBuffer::new(cfg.resident_blocks as usize * 4)?;
127        let counts_dev = DeviceBuffer::new(4)?;
128        let score_host_buf = vec![0.0_f32; model.max_blocks_per_layer as usize];
129        let disk_state = DiskState::new();
130        Ok(Self {
131            cfg,
132            model,
133            predictor,
134            pool,
135            backend,
136            attn,
137            eviction,
138            q_proj,
139            block_scores_dev,
140            block_table_dev,
141            counts_dev,
142            score_host_buf,
143            disk_state,
144        })
145    }
146
147    // ── Disk-block-ID allocator (Phase 6.1.a) ─────────────────────────
148    // Each layer has an independent ID space. Capacity == max_blocks_per_layer.
149    // alloc / free list / refcount semantics:
150    //   - alloc_disk_block_id(layer) -> Some(id) if room, else None
151    //   - inc_disk_ref(layer, id) increments (panics if id is unallocated)
152    //   - dec_disk_ref(layer, id) -> new refcount; on 0 returns id to free list
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; // capacity exhausted
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    /// Aggregated diagnostic summary across all layers (Phase 6.1.j).
199    /// Use to log periodic state during long-running decode loops; the
200    /// scheduler can call this once per N steps to verify HBM-shrink
201    /// behavior is on track.
202    pub fn diagnostic_summary(&self) -> HighSpeedSwapDiagnostic {
203        let st = &self.disk_state;
204        let active = st.next_id.saturating_sub(st.free_list.len() as u32);
205        HighSpeedSwapDiagnostic {
206            num_layers: self.model.num_layers,
207            active_disk_blocks: active,
208            disk_block_capacity: self.model.max_blocks_per_layer,
209            scratch_pool_resident: self.pool.dims().num_slots,
210            scratch_pool_free: self.pool.free_count(),
211        }
212    }
213}
214
215#[derive(Debug, Clone, Copy)]
216pub struct HighSpeedSwapDiagnostic {
217    pub num_layers: u32,
218    pub active_disk_blocks: u32,
219    pub disk_block_capacity: u32,
220    pub scratch_pool_resident: u32,
221    pub scratch_pool_free: u32,
222}
223
224#[cfg(test)]
225mod disk_id_tests;
226
227mod impl_more;
228
229// ── Thread-local installation for production callers (spark-model) ──
230//
231// The scheduler thread, after `bind_gpu_to_thread`, calls `install_local`
232// to register the orchestrator. Per-layer attention code in spark-model
233// then accesses it via `with_local`. The orchestrator's HBM allocations
234// live as long as the thread; cleanup happens on thread exit (or
235// explicit drop via `take_local`).
236
237use std::cell::RefCell;
238// THREAD-LOCAL, DELIBERATELY. The orchestrator owns HBM allocations bound to
239// ONE thread's CUDA stream, so it is not shareable across threads and cannot
240// be an `Arc` on a shared context. `install_local` is called on the scheduler
241// thread and `with_local` reads it from the per-layer attention code far down
242// the same call stack; the alternative is threading it through every layer
243// signature to reach one thread's own state. Cleanup is thread exit or an
244// explicit `take_local`, which a model teardown calls.
245thread_local! {
246    static LOCAL: RefCell<Option<HighSpeedSwap>> = const { RefCell::new(None) };
247}
248
249/// Install the orchestrator on the current thread. Idempotent (overwrites
250/// any prior installation, dropping it).
251pub fn install_local(stream: u64, cfg: HighSpeedSwapConfig, model: ModelDims) -> Result<()> {
252    let hss = HighSpeedSwap::new_on_stream(stream, cfg, model)?;
253    LOCAL.with(|cell| {
254        *cell.borrow_mut() = Some(hss);
255    });
256    Ok(())
257}
258
259/// True iff `install_local` has populated this thread's slot.
260pub fn local_installed() -> bool {
261    LOCAL.with(|cell| cell.borrow().is_some())
262}
263
264/// Run `f` with a `&mut HighSpeedSwap` if installed; returns `None` if not.
265pub fn with_local<R>(f: impl FnOnce(&mut HighSpeedSwap) -> Result<R>) -> Option<Result<R>> {
266    LOCAL.with(|cell| cell.borrow_mut().as_mut().map(f))
267}
spark_storage/high_speed_swap.rs

spark_storage/
high_speed_swap.rs
