metrale_storage/
group.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Group addressing for the high-speed-swap layer. A group is the K
4//! or V stripe of one (layer, block, kv_head): `block_size × head_dim ×
5//! elem_bytes` bytes, padded to a multiple of `fs_block_size`.
6//!
7//! Owner: metrale-storage high-speed swap.
8//! Invariants:
9//! - Within a layer, block `b` occupies `block_bytes()` contiguous bytes at
10//!   `b · block_bytes()`: the K stripes of every kv_head, then the V stripes.
11//! - For keys inside the layout, `group_id` maps each (layer, block, kind,
12//!   kv_head) to a distinct id in `0 .. num_layers · 2 · num_blocks ·
13//!   num_kv_heads`, in the same order as `file_offset`.
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
16pub struct GroupId(pub u64);
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19pub enum KvKind {
20    K = 0,
21    V = 1,
22}
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
25pub struct GroupKey {
26    pub layer: u32,
27    pub block: u32,
28    pub kv_head: u16,
29    // 2026-09-25: `KvKind as u8`; `kind()` reads 0 as K and any other value as V.
30    pub kv_kind: u8,
31}
32
33impl GroupKey {
34    pub fn new(layer: u32, block: u32, kv_head: u16, kv_kind: KvKind) -> Self {
35        Self {
36            layer,
37            block,
38            kv_head,
39            kv_kind: kv_kind as u8,
40        }
41    }
42    pub fn kind(self) -> KvKind {
43        match self.kv_kind {
44            0 => KvKind::K,
45            _ => KvKind::V,
46        }
47    }
48}
49
50#[derive(Clone, Copy, Debug)]
51pub struct GroupLayout {
52    pub num_layers: u32,
53    pub num_blocks: u32,
54    pub num_kv_heads: u16,
55    /// 2026-09-25: `block_size × head_dim × elem_bytes`, rounded up to a
56    /// multiple of `fs_block_size`.
57    pub group_stride: u64,
58    /// 2026-09-25: The alignment unit `group_stride` is padded to.
59    pub fs_block_size: u64,
60}
61
62impl GroupLayout {
63    pub fn new(
64        num_layers: u32,
65        num_blocks: u32,
66        num_kv_heads: u16,
67        block_size: u32,
68        head_dim: u32,
69        elem_bytes: u32,
70        fs_block_size: u64,
71    ) -> Self {
72        let raw = (block_size as u64) * (head_dim as u64) * (elem_bytes as u64);
73        let group_stride = raw.div_ceil(fs_block_size) * fs_block_size;
74        Self {
75            num_layers,
76            num_blocks,
77            num_kv_heads,
78            group_stride,
79            fs_block_size,
80        }
81    }
82
83    /// 2026-09-25: Bytes of one layer's file: K and V of every block and kv_head.
84    pub fn bytes_per_layer(&self) -> u64 {
85        2 * (self.num_blocks as u64) * (self.num_kv_heads as u64) * self.group_stride
86    }
87
88    /// 2026-09-25: Offset of `key` within its layer's file. Debug builds assert
89    /// that `block` and `kv_head` are inside the layout.
90    pub fn file_offset(&self, key: GroupKey) -> u64 {
91        debug_assert!(key.block < self.num_blocks);
92        debug_assert!(key.kv_head < self.num_kv_heads);
93        let kv_stride = (self.num_kv_heads as u64) * self.group_stride;
94        (key.block as u64) * (2 * kv_stride)
95            + (key.kv_kind as u64) * kv_stride
96            + (key.kv_head as u64) * self.group_stride
97    }
98
99    /// 2026-09-25: Dense `GroupId` for `key` across all layers.
100    pub fn group_id(&self, key: GroupKey) -> GroupId {
101        let per_layer = 2 * (self.num_blocks as u64) * (self.num_kv_heads as u64);
102        let per_block = 2 * (self.num_kv_heads as u64);
103        GroupId(
104            (key.layer as u64) * per_layer
105                + (key.block as u64) * per_block
106                + (key.kv_kind as u64) * (self.num_kv_heads as u64)
107                + (key.kv_head as u64),
108        )
109    }
110
111    /// 2026-09-25: Bytes one group occupies on disk: `group_stride`.
112    pub fn group_bytes(&self) -> u64 {
113        self.group_stride
114    }
115
116    /// 2026-09-25: Bytes in one block: K and V of every kv_head, each at
117    /// `group_stride` pitch. The unit of the `StorageBackend` block methods.
118    pub fn block_bytes(&self) -> u64 {
119        2 * (self.num_kv_heads as u64) * self.group_stride
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn offset_and_id_are_consistent() {
129        let l = GroupLayout::new(80, 4096, 8, 16, 128, 2, 4096);
130        // 2026-09-25: 16 × 128 × 2 = 4096 bytes needs no padding.
131        assert_eq!(l.group_stride, 4096);
132        let k0 = GroupKey::new(0, 0, 0, KvKind::K);
133        assert_eq!(l.file_offset(k0), 0);
134        let k1 = GroupKey::new(0, 0, 0, KvKind::V);
135        assert_eq!(l.file_offset(k1), (8u64) * 4096);
136        let k2 = GroupKey::new(0, 1, 0, KvKind::K);
137        assert_eq!(l.file_offset(k2), 2 * 8 * 4096);
138        let k3 = GroupKey::new(0, 0, 7, KvKind::V);
139        assert_eq!(l.file_offset(k3), 8 * 4096 + 7 * 4096);
140    }
141
142    #[test]
143    fn rounds_up_to_fs_block() {
144        // 2026-09-25: 16 × 96 × 2 = 3072 bytes pads to 4096.
145        let l = GroupLayout::new(1, 1, 1, 16, 96, 2, 4096);
146        assert_eq!(l.group_stride, 4096);
147    }
148
149    #[test]
150    fn bytes_per_layer_correct() {
151        let l = GroupLayout::new(1, 4, 2, 16, 128, 2, 4096);
152        // 2026-09-25: 4 blocks × 2 kv_heads × 2 (K, V) × 4096 bytes.
153        assert_eq!(l.bytes_per_layer(), 65536);
154    }
155}
metrale_storage/group.rs

metrale_storage/
group.rs
