metrale_model_layers/layers/ops/
ssm_gdn_carry_sizes.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The byte sizes of the carried-state GDN verify's device allocations, split from
4//! `ssm_gdn_carry.rs` unchanged: the binder (`TransformerModel::gdn_carry_bind`) allocates exactly
5//! these and the preflight reserve plans their total.
6//!
7//! Owner: model-layers ops (GDN).
8//! Invariants: none beyond the types.
9
10/// 2026-10-01: The device allocations `TransformerModel::gdn_carry_bind` makes for `layers`
11/// GDN layers and `slots` carry slots (the MTP slots plus the dummy), in bytes, one field per
12/// allocation. The binder allocates exactly these and the preflight reserve plans `total()`.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct GdnCarrySizes {
15    /// 2026-10-01: `[layers][slots][seq_floats]` f32.
16    pub stash: usize,
17    /// 2026-10-01: `[layers][slots][conv_seq_elems]` bf16.
18    pub conv_stash: usize,
19    /// 2026-10-01: `[layers][table_seqs]` u64 conv pointers, for the verify and for the flush.
20    pub conv_tab: usize,
21    pub flush_conv_tab: usize,
22    /// 2026-10-01: `[layers][slots]` u32 pending counts.
23    pub pend: usize,
24    /// 2026-10-01: `[table_seqs]` u32 slot table.
25    pub slot_tab: usize,
26    /// 2026-10-01: `[layers][table_seqs]` u32 flags.
27    pub flags: usize,
28    /// 2026-10-01: `[layers][table_seqs]` u64 flush pointers and `[table_seqs]` u32 flush slots.
29    pub flush_tab: usize,
30    pub flush_slots: usize,
31}
32
33impl GdnCarrySizes {
34    /// 2026-10-01: The sizes for `layers` layers, `slots` slots, the per-slot widths and the
35    /// verify table's `table_seqs` rows (`VERIFY_WY_TABLE_SEQS`).
36    pub const fn new(
37        layers: usize,
38        slots: usize,
39        seq_floats: usize,
40        conv_seq_elems: usize,
41        table_seqs: usize,
42    ) -> Self {
43        Self {
44            stash: layers * slots * seq_floats * 4,
45            conv_stash: layers * slots * conv_seq_elems * 2,
46            conv_tab: layers * table_seqs * 8,
47            flush_conv_tab: layers * table_seqs * 8,
48            pend: layers * slots * 4,
49            slot_tab: table_seqs * 4,
50            flags: layers * table_seqs * 4,
51            flush_tab: layers * table_seqs * 8,
52            flush_slots: table_seqs * 4,
53        }
54    }
55
56    /// 2026-10-01: Every allocation's bytes, summed.
57    pub const fn total(&self) -> usize {
58        self.stash
59            + self.conv_stash
60            + self.conv_tab
61            + self.flush_conv_tab
62            + self.pend
63            + self.slot_tab
64            + self.flags
65            + self.flush_tab
66            + self.flush_slots
67    }
68}