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}