metrale_circuit/hardware/
estimate.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: Roofline estimates on one device: the decode step at a batch, prefill of a prompt,
4//! and whether weights plus KV and recurrent state fit in the device's usable memory.
5//!
6//! Owner: metrale-circuit (hardware).
7//! Invariants:
8//! - Every node is costed by `venn::roofline::node_cost` (one cost model for Venn ranking and
9//!   hardware plans); the decode step is summed in [`super::gaps`], where a plan group that
10//!   loops per row is costed per row, so a missing multi-row kernel shows in the step time.
11//! - Prefill of `T` tokens is the main section at `T` rows with a causal half-context (`T/2`).
12//! - Datasheet peaks are ceilings, not predictions: the report names the basis of every number.
13//! - Not counted: activations and workspaces, conv windows, the CUDA context, and launch
14//!   overhead. The fit check says so beside its table.
15
16use std::collections::BTreeMap;
17
18use super::device::Device;
19use super::exec::fp4_fallback;
20use crate::format::Format;
21use crate::fuser::section_of;
22use crate::ir::{Circuit, NodeIdx, OpKind, Section};
23use crate::rules::Mode;
24use crate::venn::families::Roofline;
25use crate::venn::roofline::{CostError, node_cost};
26
27/// 2026-09-30: The roofline constants used on a device, and where they come from.
28#[derive(Debug, Clone, PartialEq)]
29pub struct DeviceRoofline {
30    /// 2026-09-30: Constants. A node's own are [`super::plan::Resolved::roofline_of`].
31    pub roofline: Roofline,
32    /// 2026-09-30: `measured (...)` or `datasheet (...)`.
33    pub basis: String,
34    /// 2026-10-01: The constants are a measurement on the device's own class; otherwise every
35    /// estimate from them is a roofline projection, unmeasured on the device.
36    pub measured: bool,
37}
38
39/// 2026-09-30: The device's constants: the measured ones when the registry points at them and
40/// they are `measured` (the class's own manifest), else the datasheet with `context_tokens`
41/// taken from `assumptions`. A device without the FP4 MMA states its NVFP4 slot at the peak of
42/// [`fp4_fallback`].
43pub fn device_roofline(
44    device: &Device,
45    measured: Option<&Roofline>,
46    assumptions: &Roofline,
47) -> DeviceRoofline {
48    if let (Some(src), Some(m)) = (&device.measured, measured) {
49        return DeviceRoofline {
50            roofline: *m,
51            basis: format!("measured achievable ({src})"),
52            measured: true,
53        };
54    }
55    let p = device.peaks;
56    let mut roofline = Roofline {
57        dram_gbps: device.bandwidth_gbps,
58        bf16_tflops: p.bf16,
59        fp8_tflops: p.fp8.unwrap_or(p.bf16),
60        nvfp4_tflops: 0.0,
61        context_tokens: assumptions.context_tokens,
62    };
63    roofline.nvfp4_tflops = p
64        .fp4_block_scale
65        .unwrap_or_else(|| fp4_fallback(device).peak(&roofline).0);
66    DeviceRoofline {
67        roofline,
68        basis: format!(
69            "datasheet ceiling (sources {}{})",
70            device.sources.join(", "),
71            if device.derived.is_empty() {
72                String::new()
73            } else {
74                format!("; derived, not published: {}", device.derived.join(", "))
75            }
76        ),
77        measured: false,
78    }
79}
80
81/// 2026-09-30: The estimated prefill time of `tokens` tokens, microseconds, each node costed
82/// with `roofline_of` its index.
83pub fn prefill_us(
84    c: &Circuit,
85    settings: &BTreeMap<String, String>,
86    roofline_of: &dyn Fn(NodeIdx) -> Roofline,
87    tokens: u64,
88) -> Result<f64, CostError> {
89    let main = section_of(Mode::Decode);
90    let mut total = 0.0;
91    for b in c.blocks.iter().filter(|b| b.section == main) {
92        for i in b.first..b.end {
93            let causal = Roofline {
94                context_tokens: (tokens / 2).max(1),
95                ..roofline_of(i)
96            };
97            total += node_cost(c, &c.nodes[i], Mode::Decode, tokens, settings, &causal)?.time_us;
98        }
99    }
100    Ok(total)
101}
102
103/// 2026-09-30: What one sequence and the whole model hold in device memory.
104#[derive(Debug, Clone, Copy, PartialEq)]
105pub struct Footprint {
106    /// 2026-09-30: Every weight of every section (embedding at BF16), bytes.
107    pub weights: f64,
108    /// 2026-09-30: KV bytes per token of one sequence, over every attention node.
109    pub kv_per_token: f64,
110    /// 2026-09-30: Recurrent state bytes of one sequence (GatedDeltaNet and Mamba2).
111    pub state_per_seq: f64,
112}
113
114impl Footprint {
115    /// 2026-09-30: Bytes for `seqs` sequences of `context` tokens.
116    pub fn bytes(&self, seqs: u64, context: u64) -> f64 {
117        self.weights + seqs as f64 * (self.kv_per_token * context as f64 + self.state_per_seq)
118    }
119}
120
121/// 2026-09-30: The footprint of `c` under `settings` (`kv_cache_dtype`, `ssm_h_dtype`).
122pub fn footprint(c: &Circuit, settings: &BTreeMap<String, String>) -> Result<Footprint, String> {
123    let dim = |d: &str| {
124        c.dims
125            .get(d)
126            .map(|v| *v as f64)
127            .ok_or_else(|| format!("the footprint needs dim `{d}`"))
128    };
129    let kv = match settings.get("kv_cache_dtype").map(String::as_str) {
130        Some("bf16") => 2.0,
131        Some("fp8") => 1.0,
132        other => return Err(format!("kv_cache_dtype {other:?} has no element size")),
133    };
134    let h = match settings.get("ssm_h_dtype").map(String::as_str) {
135        Some("f32") => 4.0,
136        Some("f16" | "f16-pool" | "bf16") => 2.0,
137        other => return Err(format!("ssm_h_dtype {other:?} has no element size")),
138    };
139    let mut weights = 0.0;
140    let (mut kv_per_token, mut state) = (0.0, 0.0);
141    let main: Vec<bool> = {
142        let mut v = vec![false; c.nodes.len()];
143        for b in c.blocks.iter().filter(|b| b.section == Section::Main) {
144            v[b.first..b.end].iter_mut().for_each(|x| *x = true);
145        }
146        v
147    };
148    for (i, n) in c.nodes.iter().enumerate() {
149        // 2026-09-30: The draft head reads the main embedding table; count it once.
150        if n.op == OpKind::Embed && main[i] {
151            weights += dim("vocab")? * dim("hidden")? * 2.0;
152        }
153        if let Some(w) = n.weight {
154            let k = n.inputs.first().map_or(0, |&e| c.edges[e].dim_value);
155            let out: u64 = n.outputs.iter().map(|&e| c.edges[e].dim_value).sum();
156            let one = w
157                .weight_bytes(out, k)
158                .ok_or_else(|| format!("node `{}`: weight {} has no size", n.id, w.name()))?
159                as f64;
160            let copies = match n.op {
161                OpKind::ExpertGateUp | OpKind::ExpertDown => dim("experts")?,
162                _ => 1.0,
163            };
164            weights += one * copies;
165        }
166        match n.op {
167            OpKind::PagedAttention => {
168                kv_per_token += 2.0 * dim("kv_heads")? * dim("head_dim")? * kv
169            }
170            OpKind::GdnRecurrence => {
171                state += dim("lin_v_heads")? * dim("lin_k_dim")? * dim("lin_v_dim")? * h;
172            }
173            OpKind::SsmUpdate => {
174                state += dim("mamba_heads")? * dim("mamba_head_dim")? * dim("ssm_state")? * 4.0;
175            }
176            _ => {}
177        }
178    }
179    Ok(Footprint {
180        weights,
181        kv_per_token,
182        state_per_seq: state,
183    })
184}
185
186/// 2026-09-30: Weight bytes a decode step reads at one row: every main-section weight once
187/// (routed experts at `top_k` of them). Divided by bandwidth, the C=1 floor.
188pub fn weight_floor_bytes(c: &Circuit) -> Result<f64, String> {
189    let mut total = 0.0;
190    for b in c.blocks.iter().filter(|b| b.section == Section::Main) {
191        for n in &c.nodes[b.first..b.end] {
192            let Some(w) = n.weight else { continue };
193            let k = n.inputs.first().map_or(0, |&e| c.edges[e].dim_value);
194            let out: u64 = n.outputs.iter().map(|&e| c.edges[e].dim_value).sum();
195            let copies = match n.op {
196                OpKind::ExpertGateUp | OpKind::ExpertDown => {
197                    *c.dims.get("top_k").ok_or_else(|| {
198                        format!("node `{}`: routed experts need dim `top_k`", n.id)
199                    })? as f64
200                }
201                _ => 1.0,
202            };
203            let one = w
204                .weight_bytes(out, k)
205                .ok_or_else(|| format!("node `{}`: weight {} has no size", n.id, w.name()))?;
206            total += one as f64 * copies;
207        }
208    }
209    Ok(total)
210}
211
212/// 2026-09-30: `Format` of a node's first input.
213pub fn activation_of(c: &Circuit, n: &crate::ir::Node) -> Option<Format> {
214    n.inputs.first().map(|&e| c.edges[e].format)
215}
metrale_circuit/hardware/estimate.rs

metrale_circuit/hardware/
estimate.rs
