metrale_storage/backend/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The `StorageBackend` trait, which moves KV groups and blocks
4//! between a storage tier and device memory, and the two file backends.
5//!
6//! Owner: metrale-storage high-speed swap.
7//! Invariants:
8//! - A backend that keeps the default `read_blocks` issues, for each block,
9//!   the per-head reads of `expand_blocks_to_groups`, in that order (pinned
10//!   by `mod_tests.rs`).
11
12use anyhow::Result;
13
14use crate::group::{GroupKey, GroupLayout, KvKind};
15
16// 2026-09-25: io_uring is Linux-only; `posix` builds on every target.
17#[cfg(target_os = "linux")]
18pub mod io_uring;
19pub mod posix;
20
21#[cfg(target_os = "linux")]
22pub use self::io_uring::IoUringBackend;
23pub use posix::PosixBackend;
24
25/// 2026-09-25: Read `group` from the tier into device memory at `dst_dev_ptr`.
26#[derive(Clone, Copy, Debug)]
27pub struct ReadRequest {
28    pub group: GroupKey,
29    pub dst_dev_ptr: u64,
30}
31
32/// 2026-09-25: Read one whole block (`block_bytes()`: every kv_head's K, then
33/// every V) into device memory starting at `dst_dev_ptr`. Only `base_key`'s
34/// `layer` and `block` are read.
35#[derive(Clone, Copy, Debug)]
36pub struct BlockReadRequest {
37    pub base_key: GroupKey,
38    pub dst_dev_ptr: u64,
39}
40
41/// 2026-09-25: Expand each block request into its `2·nkv` per-head requests,
42/// ordered `K(kh), V(kh)` for `kh` in `0..nkv`, with device destinations
43/// `dst + kh·gs` (K) and `dst + (nkv + kh)·gs` (V). The default `read_blocks`
44/// and `read_blocks_async` fan out through it.
45pub fn expand_blocks_to_groups(spec: &GroupLayout, reqs: &[BlockReadRequest]) -> Vec<ReadRequest> {
46    let nkv = spec.num_kv_heads;
47    let gs = spec.group_stride;
48    let mut out = Vec::with_capacity(reqs.len() * 2 * nkv as usize);
49    for r in reqs {
50        let layer = r.base_key.layer;
51        let block = r.base_key.block;
52        for kh in 0..nkv {
53            out.push(ReadRequest {
54                group: GroupKey::new(layer, block, kh, KvKind::K),
55                dst_dev_ptr: r.dst_dev_ptr + (kh as u64) * gs,
56            });
57            out.push(ReadRequest {
58                group: GroupKey::new(layer, block, kh, KvKind::V),
59                dst_dev_ptr: r.dst_dev_ptr + (nkv as u64 + kh as u64) * gs,
60            });
61        }
62    }
63    out
64}
65
66pub trait StorageBackend: Send + Sync {
67    /// 2026-09-25: Fill every request's destination. Returns once the data is
68    /// in place, so work the caller enqueues on `stream` afterwards reads it.
69    fn read(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()>;
70
71    /// 2026-09-25: Like `read`, but an implementation may return before the
72    /// copies on `stream` finish. The default calls `read`.
73    fn read_async(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()> {
74        self.read(requests, stream)
75    }
76
77    /// 2026-09-25: Write one group from host memory; `src` is `group_bytes()`
78    /// long.
79    fn write_from_host(&mut self, key: GroupKey, src: &[u8]) -> Result<()>;
80
81    /// 2026-09-25: The geometry the default block methods fan out with.
82    fn group_layout(&self) -> GroupLayout;
83
84    /// 2026-09-25: Read whole blocks, with the stream contract of `read`. The
85    /// default fans out to `read` through `expand_blocks_to_groups`.
86    fn read_blocks(&mut self, requests: &[BlockReadRequest], stream: u64) -> Result<()> {
87        let groups = expand_blocks_to_groups(&self.group_layout(), requests);
88        self.read(&groups, stream)
89    }
90
91    /// 2026-09-25: The block form of `read_async`. The default fans out to
92    /// `read_async`.
93    fn read_blocks_async(&mut self, requests: &[BlockReadRequest], stream: u64) -> Result<()> {
94        let groups = expand_blocks_to_groups(&self.group_layout(), requests);
95        self.read_async(&groups, stream)
96    }
97
98    /// 2026-09-25: Write one whole block. `src` is `block_bytes()` laid out
99    /// `[K0, …, K(nkv-1), V0, …, V(nkv-1)]` at `group_stride` pitch; only
100    /// `base_key`'s `layer` and `block` are read. The default rejects any other
101    /// length, then calls `write_from_host` for K(kh) and V(kh) of each head.
102    fn write_block_from_host(&mut self, base_key: GroupKey, src: &[u8]) -> Result<()> {
103        let spec = self.group_layout();
104        let nkv = spec.num_kv_heads as usize;
105        let gs = spec.group_stride as usize;
106        let expect = 2 * nkv * gs;
107        if src.len() != expect {
108            anyhow::bail!(
109                "write_block_from_host: src len {} != block bytes {expect}",
110                src.len()
111            );
112        }
113        let layer = base_key.layer;
114        let block = base_key.block;
115        for kh in 0..nkv {
116            let k_off = kh * gs;
117            let v_off = (nkv + kh) * gs;
118            self.write_from_host(
119                GroupKey::new(layer, block, kh as u16, KvKind::K),
120                &src[k_off..k_off + gs],
121            )?;
122            self.write_from_host(
123                GroupKey::new(layer, block, kh as u16, KvKind::V),
124                &src[v_off..v_off + gs],
125            )?;
126        }
127        Ok(())
128    }
129
130    /// 2026-09-25: Write `run_len` consecutive blocks of one layer, starting at
131    /// `base_key`'s block; `src` is `run_len · block_bytes()`. The default
132    /// rejects any other length, then calls `write_block_from_host` once per
133    /// block.
134    fn write_blocks_run(&mut self, base_key: GroupKey, run_len: usize, src: &[u8]) -> Result<()> {
135        let spec = self.group_layout();
136        let block_bytes = spec.block_bytes() as usize;
137        let expect = run_len * block_bytes;
138        if src.len() != expect {
139            anyhow::bail!(
140                "write_blocks_run: src len {} != run bytes {expect} ({run_len} × {block_bytes})",
141                src.len()
142            );
143        }
144        for i in 0..run_len {
145            let off = i * block_bytes;
146            self.write_block_from_host(
147                GroupKey::new(base_key.layer, base_key.block + i as u32, 0, KvKind::K),
148                &src[off..off + block_bytes],
149            )?;
150        }
151        Ok(())
152    }
153
154    /// 2026-09-25: Whether this backend serves `write_blocks_run` as one wide
155    /// op. The default is `false`.
156    fn supports_write_run_coalescing(&self) -> bool {
157        false
158    }
159
160    /// 2026-09-25: Register `[base, base + len)` as the region reads land in.
161    /// The RDMA backends register it as one memory region per rail; the default
162    /// does nothing.
163    fn register_landing_region(&mut self, base: u64, len: usize) -> Result<()> {
164        let _ = (base, len);
165        Ok(())
166    }
167}
168
169#[cfg(test)]
170#[path = "mod_tests.rs"]
171mod coalesce_tests;
metrale_storage/backend/mod.rs

metrale_storage/backend/
mod.rs
