spark_runtime/buffers/
decode_meta.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Fixed-stride batched-decode metadata layout, derived from the serve
4//! `max_batch_size` (SSOT — consumed by `sizes.rs` for the scratch envelope
5//! and by `spark-model`'s `upload_batch_metadata_fixed`/`_at` for the
6//! upload offsets, replacing the former hardcoded 0/128/256/512/768 gaps
7//! that fit exactly 32 rows).
8//!
9//! Region shapes (byte offsets, `R = rows`):
10//!   positions  u32  [0,       4R)
11//!   seq_slot   i32  [4R,      8R)   (per-request LoRA routing)
12//!   slots      i64  [8R,     16R)
13//!   seq_lens   i32  [16R,    20R)
14//!   (gap            [20R,    24R))  — legacy [640,768) pad, scaled
15//!   block_tbl  i32  [24R,    24R + R·max_blocks·4)
16//!
17//! At `R = 32` this reproduces the legacy layout BYTE-FOR-BYTE
18//! (0/128/256/512/768), so every boot with `max_batch_size <= 32` is
19//! byte-identical in addresses, strides and upload sizes.
20
21/// Layout floor: the legacy fixed layout was sized for exactly 32 rows;
22/// deriving `rows = max(32, bs)` keeps every `bs <= 32` boot byte-identical.
23pub const DECODE_META_MIN_ROWS: usize = 32;
24
25/// Layout ceiling, checked at serve time (`serve.rs`). The metadata gaps
26/// themselves derive cleanly to any width; the binding constraints on this
27/// tip are downstream row consumers sized at 96+ rows:
28/// * the logits arena (`sizes.rs`) — derived `max(96, rows+1)` rows, where
29///   `rows+1` covers the run_standard mixed path parking prefill logits at
30///   row `padded_n`;
31/// * the scratch block-table envelope (`sizes.rs`) — derived
32///   `max(verify 96-row overlay, decode `rows`-row layout)`.
33///
34/// Batched-decode kernels are row-count parametric (grid.y = n, smem per
35/// CTA constant, split-K workspace derived from the pinned max batch), so
36/// 128 is a policy cap for the widths validated by the enterprise-
37/// concurrency campaign, not an smem wall.
38pub const DECODE_META_MAX_ROWS: usize = 128;
39
40/// Derived fixed-stride decode-metadata layout.
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub struct DecodeMetaLayout {
43    rows: usize,
44}
45
46impl DecodeMetaLayout {
47    /// Derive the layout from the serve `max_batch_size`. Callers gate
48    /// `max_batch_size <= DECODE_META_MAX_ROWS` at serve time; this
49    /// constructor only applies the byte-identity floor.
50    pub fn for_max_batch_size(max_batch_size: usize) -> Self {
51        Self {
52            rows: max_batch_size.max(DECODE_META_MIN_ROWS),
53        }
54    }
55
56    /// Row capacity of the metadata block (== the widest `padded_n` the
57    /// upload accepts).
58    pub fn rows(&self) -> usize {
59        self.rows
60    }
61
62    /// `positions` u32 stream offset.
63    pub fn positions_off(&self) -> usize {
64        0
65    }
66
67    /// Per-request LoRA adapter-slot i32 stream offset.
68    pub fn seq_slot_off(&self) -> usize {
69        4 * self.rows
70    }
71
72    /// KV `slots` i64 stream offset (8-byte aligned: `8R`).
73    pub fn slots_off(&self) -> usize {
74        8 * self.rows
75    }
76
77    /// `seq_lens` i32 stream offset.
78    pub fn seq_lens_off(&self) -> usize {
79        16 * self.rows
80    }
81
82    /// Flattened block-table offset (row stride `max_blocks · 4` bytes).
83    pub fn block_table_off(&self) -> usize {
84        24 * self.rows
85    }
86
87    /// Total bytes of the metadata block for `max_blocks` blocks per row.
88    pub fn meta_bytes(&self, max_blocks: usize) -> usize {
89        self.block_table_off() + self.rows * max_blocks * 4
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    /// bs <= 32 must reproduce the legacy hardcoded layout byte-for-byte.
98    #[test]
99    fn legacy_layout_at_or_below_32() {
100        for bs in [1usize, 31, 32] {
101            let l = DecodeMetaLayout::for_max_batch_size(bs);
102            assert_eq!(l.rows(), 32, "bs={bs}");
103            assert_eq!(l.positions_off(), 0);
104            assert_eq!(l.seq_slot_off(), 128);
105            assert_eq!(l.slots_off(), 256);
106            assert_eq!(l.seq_lens_off(), 512);
107            assert_eq!(l.block_table_off(), 768);
108            // Legacy total: 768 + 32·mb·4 (decode bt region in sizes.rs).
109            assert_eq!(l.meta_bytes(257), 768 + 32 * 257 * 4);
110        }
111    }
112
113    /// Widened layouts: regions must be contiguous-or-gapped exactly like
114    /// the legacy shape scaled by R/32, non-overlapping, and 8-byte aligned
115    /// where i64 lands.
116    #[test]
117    fn widened_layout_arithmetic() {
118        for bs in [33usize, 64, 128] {
119            let l = DecodeMetaLayout::for_max_batch_size(bs);
120            let r = l.rows();
121            assert_eq!(r, bs, "bs={bs}: rows derive from bs above the floor");
122            // positions [0,4R) then seq_slot [4R,8R): no overlap.
123            assert_eq!(l.seq_slot_off(), l.positions_off() + 4 * r);
124            // slots i64 begins exactly after seq_slot and is 8-byte aligned.
125            assert_eq!(l.slots_off(), l.seq_slot_off() + 4 * r);
126            assert_eq!(l.slots_off() % 8, 0);
127            // seq_lens begins exactly after the 8R-byte slots region.
128            assert_eq!(l.seq_lens_off(), l.slots_off() + 8 * r);
129            // block table begins after seq_lens (4R) + the scaled legacy pad (4R).
130            assert_eq!(l.block_table_off(), l.seq_lens_off() + 8 * r);
131            assert_eq!(l.meta_bytes(257), 24 * r + r * 257 * 4);
132        }
133    }
134
135    #[test]
136    fn policy_ceiling_and_floor_are_exact() {
137        assert_eq!(DECODE_META_MIN_ROWS, 32);
138        assert_eq!(DECODE_META_MAX_ROWS, 128);
139    }
140}
spark_runtime/buffers/decode_meta.rs

spark_runtime/buffers/
decode_meta.rs
