metrale_gpu_runtime/buffers/sizes/
total.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: [`BufferSizes::total_bytes`], moved from `sizes.rs` unchanged to keep that
4//! file under the 500-line cap.
5//!
6//! Owner: gpu-runtime.
7//! Invariants: every size field of [`BufferSizes`] except `o_latent` and `norm_unit_w` is in
8//! the sum.
9
10use super::BufferSizes;
11
12impl BufferSizes {
13    /// 2026-09-25: The sum of the sizes, which preflight reserves for the arena.
14    /// `o_latent` and `norm_unit_w` are not in it.
15    pub fn total_bytes(&self) -> usize {
16        self.hidden_states
17            + self.residual
18            + self.norm_output
19            + self.qkv_output
20            + self.attn_output
21            + self.gate_logits
22            + self.gate_logits_f32
23            + self.moe_router_in_f32
24            + self.moe_output
25            + self.logits
26            + self.ssm_qkvz
27            + self.ssm_ba
28            + self.ssm_deinterleaved
29            + self.ssm_gates
30            + self.ssm_conv_out_f32
31            + self.scratch
32            + self.expert_gate_out
33            + self.expert_up_out
34            + self.hc_lowrank_scratch
35            + self.qsa_select_scratch
36            + self.expert_down_out
37            + self.splitk_workspace
38            + self.gdn_fla_scratch
39            + self.ssd_scratch
40            + self.hc_streams
41            + self.hc_post
42            + self.hc_comb
43            + self.token_ids
44            + self.ffn_act_q8
45            + self.ffn_act_a
46            + self.ffn_gate_up_fused
47            + self.ffn_act_scale
48            + self.ffn_act_scale_kmajor
49            + self.fp8_act
50            + self.moe_fp8_scratch
51            + self.fp8_act_scale
52            + self.fp8_act_scale_kmajor
53            + self.lora_xa
54            + self.lora_delta
55            + self.lora_hact
56            + self.lora_seq_slot
57            + self.q2_dequant_scratch
58            + self.q2_act_q8
59            + self.ssm_rowwise_w_bf16
60    }
61}