metrale_circuit/venn/
roofline.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: The roofline estimate of one node's time in one step, from its edge shapes and
4//! formats: `max(bytes / DRAM bandwidth, FLOPs / tensor peak)`. It ranks a Venn report's rows
5//! until measured profiles replace it.
6//!
7//! Owner: metrale-circuit (venn).
8//! Invariants:
9//! - Bytes are every edge the node reads or writes at `rows` rows, plus its weights, plus the
10//!   state it reads and writes: the paged KV cache over `context_tokens` per sequence, the
11//!   GatedDeltaNet and Mamba2 recurrent states (FP32, read and written once per sequence).
12//! - Routed experts read `E * (1 - (1 - k/E)^T)` distinct experts' weights for `T` tokens
13//!   (uniform routing).
14//! - The peak is the MMA class of the node's input activation: FP8 or NVFP4 activations run the
15//!   FP8 / NVFP4 MMA; anything else (W4A16, W8A16, BF16) runs the BF16 MMA.
16//! - Not counted: launch overhead, conv windows, KV writes beyond their input edges. A dim an op
17//!   needs that the circuit lacks is an error, never a guess.
18
19use std::collections::BTreeMap;
20
21use super::families::Roofline;
22use crate::format::Format;
23use crate::ir::{Circuit, Node, OpKind};
24use crate::rules::Mode;
25
26/// 2026-09-29: One node's estimate.
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct Cost {
29    /// 2026-09-29: DRAM bytes.
30    pub bytes: f64,
31    /// 2026-09-29: FLOPs.
32    pub flops: f64,
33    /// 2026-09-29: `max(bytes / bandwidth, flops / peak)`, microseconds.
34    pub time_us: f64,
35}
36
37/// 2026-09-29: Why a node could not be estimated.
38#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
39pub enum CostError {
40    /// 2026-09-29: A dim the op's cost needs.
41    #[error("node `{node}`: the estimate needs dim `{dim}`, which the circuit does not define")]
42    MissingDim {
43        /// 2026-09-29: Node id.
44        node: String,
45        /// 2026-09-29: The dim.
46        dim: String,
47    },
48    /// 2026-09-29: An edge whose rows do not evaluate or whose bytes do not divide.
49    #[error("node `{node}`: edge `{edge}` has no byte size at {rows} rows")]
50    Edge {
51        /// 2026-09-29: Node id.
52        node: String,
53        /// 2026-09-29: Edge id.
54        edge: String,
55        /// 2026-09-29: Rows.
56        rows: u64,
57    },
58    /// 2026-09-29: A KV-cache dtype the estimate has no element size for.
59    #[error("kv_cache_dtype `{0}` has no element size (bf16 | fp8)")]
60    KvDtype(String),
61}
62
63/// 2026-09-29: Estimate `node` at `rows` rows of `mode`.
64pub fn node_cost(
65    c: &Circuit,
66    n: &Node,
67    mode: Mode,
68    rows: u64,
69    settings: &BTreeMap<String, String>,
70    r: &Roofline,
71) -> Result<Cost, CostError> {
72    let dim = |name: &str| {
73        c.dims
74            .get(name)
75            .map(|v| *v as f64)
76            .ok_or_else(|| CostError::MissingDim {
77                node: n.id.clone(),
78                dim: name.to_string(),
79            })
80    };
81    let seqs = if mode == Mode::MultiSeq { rows } else { 1 } as f64;
82    let t = rows as f64;
83    let mut bytes = 0.0;
84    let contiguous_split =
85        n.op == OpKind::Split && n.params.get("layout").map(String::as_str) == Some("contiguous");
86    if !contiguous_split {
87        for &e in n.inputs.iter().chain(&n.outputs) {
88            bytes += edge_bytes(c, n, e, rows)?;
89        }
90    }
91    let k = n
92        .inputs
93        .first()
94        .map_or(0.0, |&e| c.edges[e].dim_value as f64);
95    let out: f64 = n.outputs.iter().map(|&e| c.edges[e].dim_value as f64).sum();
96    let mut flops = 0.0;
97    match n.op {
98        OpKind::Linear(_) | OpKind::LmHead | OpKind::Router => {
99            bytes += weight_bytes(n, out, k)?;
100            flops = 2.0 * t * k * out;
101        }
102        OpKind::ExpertGateUp | OpKind::ExpertDown => {
103            let (e, top) = (dim("experts")?, dim("top_k")?);
104            let distinct = e * (1.0 - (1.0 - top / e).powf(t));
105            bytes += distinct * weight_bytes(n, out, k)?;
106            flops = 2.0 * t * top * k * out;
107        }
108        OpKind::PagedAttention => {
109            let kv = match settings.get("kv_cache_dtype").map(String::as_str) {
110                Some("bf16") => 2.0,
111                Some("fp8") => 1.0,
112                other => return Err(CostError::KvDtype(other.unwrap_or("<unset>").to_string())),
113            };
114            let ctx = r.context_tokens as f64;
115            let (kvh, hd, qh) = (dim("kv_heads")?, dim("head_dim")?, dim("q_heads")?);
116            bytes += seqs * ctx * kvh * hd * 2.0 * kv;
117            flops = 4.0 * t * ctx * qh * hd;
118        }
119        OpKind::GdnRecurrence | OpKind::SsmUpdate => {
120            let elems = if n.op == OpKind::GdnRecurrence {
121                dim("lin_v_heads")? * dim("lin_k_dim")? * dim("lin_v_dim")?
122            } else {
123                dim("mamba_heads")? * dim("mamba_head_dim")? * dim("ssm_state")?
124            };
125            // 2026-09-29: FP32 state, read and written once per sequence; about four FLOPs per
126            // state element per row.
127            bytes += 2.0 * seqs * elems * 4.0;
128            flops = 4.0 * t * elems;
129        }
130        _ => {}
131    }
132    let peak = match n.inputs.first().map(|&e| c.edges[e].format) {
133        Some(Format::Fp8E4m3 { .. }) if n.op.reads_linear_weight() => r.fp8_tflops,
134        _ if nvfp4_mma(c, n) => r.nvfp4_tflops,
135        _ => r.bf16_tflops,
136    };
137    let time_us = (bytes / (r.dram_gbps * 1e3)).max(flops / (peak * 1e6));
138    Ok(Cost {
139        bytes,
140        flops,
141        time_us,
142    })
143}
144
145/// 2026-10-01: `n` multiplies an NVFP4 activation by a linear weight: the node the NVFP4 peak
146/// costs.
147pub fn nvfp4_mma(c: &Circuit, n: &Node) -> bool {
148    n.op.reads_linear_weight()
149        && matches!(
150            n.inputs.first().map(|&e| c.edges[e].format),
151            Some(Format::Nvfp4 { .. })
152        )
153}
154
155fn edge_bytes(c: &Circuit, n: &Node, e: usize, rows: u64) -> Result<f64, CostError> {
156    let edge = &c.edges[e];
157    let mut dims = c.dims.clone();
158    dims.insert("n".into(), rows);
159    let bad = || CostError::Edge {
160        node: n.id.clone(),
161        edge: edge.id.clone(),
162        rows,
163    };
164    let r = edge.rows.eval(&dims).map_err(|_| bad())?;
165    edge.format
166        .bytes(r, edge.dim_value)
167        .map(|b| b as f64)
168        .ok_or_else(bad)
169}
170
171fn weight_bytes(n: &Node, out: f64, k: f64) -> Result<f64, CostError> {
172    let Some(w) = n.weight else {
173        return Ok(0.0);
174    };
175    w.weight_bytes(out as u64, k as u64)
176        .map(|b| b as f64)
177        .ok_or_else(|| CostError::Edge {
178            node: n.id.clone(),
179            edge: "<weight>".into(),
180            rows: out as u64,
181        })
182}
metrale_circuit/venn/roofline.rs

metrale_circuit/venn/
roofline.rs
