metrale_gpu_runtime/buffers/
decode_meta.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Fixed-stride batched-decode metadata layout, derived from the serve `max_batch_size`.
4//!
5//! `sizes.rs` sizes `scratch` from it, and metrale-model-engine's
6//! `upload_batch_metadata_fixed` / `_at` take their offsets from it. Byte
7//! offsets, with `R = rows`:
8//!   positions  u32  [0,       4R)
9//!   seq_slot   i32  [4R,      8R)   (per-request LoRA slot)
10//!   slots      i64  [8R,     16R)
11//!   seq_lens   i32  [16R,    20R)
12//!   (unused         [20R,    24R))
13//!   block_tbl  i32  [24R,    24R + R·max_blocks·4)
14//!
15//! Owner: gpu-runtime.
16//! Invariants:
17//! - `rows() >= DECODE_META_MIN_ROWS`.
18//! - The regions do not overlap, and `slots_off()` is a multiple of 8.
19
20/// 2026-09-25: The row floor: `rows = max(DECODE_META_MIN_ROWS, max_batch_size)`.
21pub const DECODE_META_MIN_ROWS: usize = 32;
22
23/// 2026-09-25: The largest serve `max_batch_size` accepted; `serve_load.rs`
24/// refuses a larger one at startup. The layout itself has no upper bound.
25pub const DECODE_META_MAX_ROWS: usize = 128;
26
27/// 2026-09-25: The decode-metadata layout for `rows` rows (see the module docs).
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
29pub struct DecodeMetaLayout {
30    rows: usize,
31}
32
33impl DecodeMetaLayout {
34    /// 2026-09-25: The layout for the serve `max_batch_size`. It applies only the
35    /// floor; the ceiling `DECODE_META_MAX_ROWS` is checked by the serve path.
36    pub fn for_max_batch_size(max_batch_size: usize) -> Self {
37        Self {
38            rows: max_batch_size.max(DECODE_META_MIN_ROWS),
39        }
40    }
41
42    /// 2026-09-25: Row capacity, the widest `padded_n` the upload accepts.
43    pub fn rows(&self) -> usize {
44        self.rows
45    }
46
47    pub fn positions_off(&self) -> usize {
48        0
49    }
50
51    pub fn seq_slot_off(&self) -> usize {
52        4 * self.rows
53    }
54
55    pub fn slots_off(&self) -> usize {
56        8 * self.rows
57    }
58
59    pub fn seq_lens_off(&self) -> usize {
60        16 * self.rows
61    }
62
63    /// 2026-09-25: Offset of the flattened block table, row stride `max_blocks · 4` bytes.
64    pub fn block_table_off(&self) -> usize {
65        24 * self.rows
66    }
67
68    /// 2026-09-25: Bytes of the whole block for `max_blocks` blocks per row.
69    pub fn meta_bytes(&self, max_blocks: usize) -> usize {
70        self.block_table_off() + self.rows * max_blocks * 4
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    /// 2026-09-25: At or below the floor the offsets are 0/128/256/512/768.
79    #[test]
80    fn legacy_layout_at_or_below_32() {
81        for bs in [1usize, 31, 32] {
82            let l = DecodeMetaLayout::for_max_batch_size(bs);
83            assert_eq!(l.rows(), 32, "bs={bs}");
84            assert_eq!(l.positions_off(), 0);
85            assert_eq!(l.seq_slot_off(), 128);
86            assert_eq!(l.slots_off(), 256);
87            assert_eq!(l.seq_lens_off(), 512);
88            assert_eq!(l.block_table_off(), 768);
89            assert_eq!(l.meta_bytes(257), 768 + 32 * 257 * 4);
90        }
91    }
92
93    /// 2026-09-25: Above the floor, rows follow `max_batch_size`, the regions
94    /// are back to back except the 4R gap before the block table, and the i64
95    /// slots are 8-byte aligned.
96    #[test]
97    fn widened_layout_arithmetic() {
98        for bs in [33usize, 64, 128] {
99            let l = DecodeMetaLayout::for_max_batch_size(bs);
100            let r = l.rows();
101            assert_eq!(r, bs, "bs={bs}: rows derive from bs above the floor");
102            assert_eq!(l.seq_slot_off(), l.positions_off() + 4 * r);
103            assert_eq!(l.slots_off(), l.seq_slot_off() + 4 * r);
104            assert_eq!(l.slots_off() % 8, 0);
105            assert_eq!(l.seq_lens_off(), l.slots_off() + 8 * r);
106            assert_eq!(l.block_table_off(), l.seq_lens_off() + 8 * r);
107            assert_eq!(l.meta_bytes(257), 24 * r + r * 257 * 4);
108        }
109    }
110
111    #[test]
112    fn policy_ceiling_and_floor_are_exact() {
113        assert_eq!(DECODE_META_MIN_ROWS, 32);
114        assert_eq!(DECODE_META_MAX_ROWS, 128);
115    }
116}
metrale_gpu_runtime/buffers/decode_meta.rs

metrale_gpu_runtime/buffers/
decode_meta.rs
