metrale_circuit/memory/
copies.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Load-time derived weight copies (`kernels/circuits/COPIES.toml`): every
4//! device allocation the loader makes for a weight beyond the checkpoint tensors it stores
5//! (requantized serving copies, transposed prefill twins, repacks, widened scales, and the
6//! leaked intermediates nothing reads). A rule matches nodes by op, by the checkpoint's declared
7//! weight format, by the served format and by serve settings, and states the byte law of its copy.
8//!
9//! Owner: metrale-circuit (memory).
10//! Invariants:
11//! - A copy's bytes come from the node's own shape (`out x k`, per expert for routed experts) and
12//!   the rule's law; nothing is sized from a measured total.
13//! - Every rule cites the loader site that allocates it; a rule that matches no node of a circuit
14//!   is not an error (it belongs to another checkpoint), but an unknown op, format, law or
15//!   setting spelling is refused at parse.
16
17use std::collections::BTreeMap;
18
19use serde::Deserialize;
20
21use crate::format::Format;
22use crate::ir::{OpKind, Section};
23
24/// 2026-10-02: How a copy's bytes follow from the weight's `out x k` shape.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum CopyLaw {
27    /// 2026-10-02: A whole weight in this format ([`Format::weight_bytes`]).
28    Weight(Format),
29    /// 2026-10-02: One F32 scale per output row (`out x 4`).
30    RowScaleF32,
31    /// 2026-10-02: One F32 scale per `r x c` block (`ceil(out/r) x ceil(k/c) x 4`).
32    BlockGridF32(u32, u32),
33}
34
35impl CopyLaw {
36    /// 2026-10-02: `bf16`, `nvfp4/g16`, `fp8/...` (a [`Format`]), `row_scale_f32` or
37    /// `block_grid_f32/<r>x<c>`.
38    pub fn parse(s: &str) -> Option<Self> {
39        if s == "row_scale_f32" {
40            return Some(Self::RowScaleF32);
41        }
42        if let Some(rc) = s.strip_prefix("block_grid_f32/") {
43            let (r, c) = rc.split_once('x')?;
44            let pos = |v: &str| v.parse::<u32>().ok().filter(|&v| v > 0);
45            return Some(Self::BlockGridF32(pos(r)?, pos(c)?));
46        }
47        Format::parse(s).ok().map(Self::Weight)
48    }
49
50    /// 2026-10-02: Bytes of one copy of an `out x k` weight; `None` when the shape does not fit
51    /// the format's groups, or on overflow.
52    pub fn bytes(self, out: u64, k: u64) -> Option<u64> {
53        match self {
54            Self::Weight(f) => f.weight_bytes(out, k),
55            Self::RowScaleF32 => out.checked_mul(4),
56            Self::BlockGridF32(r, c) => out
57                .div_ceil(u64::from(r))
58                .checked_mul(k.div_ceil(u64::from(c)))?
59                .checked_mul(4),
60        }
61    }
62}
63
64/// 2026-10-02: One copy rule.
65#[derive(Debug, Clone, PartialEq)]
66pub struct CopyRule {
67    /// 2026-10-02: Rule id (unique).
68    pub id: String,
69    /// 2026-10-02: Arches it applies to.
70    pub arch: Vec<String>,
71    /// 2026-10-02: Op spellings it matches ([`OpKind::name`]: `linear:q`, `lm_head`, ...).
72    pub ops: Vec<String>,
73    /// 2026-10-02: The section it matches (`main` | `draft`).
74    pub section: Section,
75    /// 2026-10-02: Declared (stored) weight formats it matches; empty matches any.
76    pub stored: Vec<Format>,
77    /// 2026-10-02: Served weight formats it matches; empty matches any.
78    pub served: Vec<Format>,
79    /// 2026-10-02: Serve settings that must hold (`speculative = "on"`).
80    pub when: BTreeMap<String, String>,
81    /// 2026-10-02: The copy's byte law.
82    pub law: CopyLaw,
83    /// 2026-10-02: Copies per matched weight.
84    pub count: u64,
85    /// 2026-10-02: Nothing reads it and nothing frees it (a loader defect, resident all the same).
86    pub leaked: bool,
87    /// 2026-10-02: The loader site that allocates it.
88    pub site: String,
89}
90
91#[derive(Deserialize)]
92#[serde(deny_unknown_fields)]
93struct File {
94    schema: u32,
95    copy: Vec<RuleFile>,
96}
97
98#[derive(Deserialize)]
99#[serde(deny_unknown_fields)]
100struct RuleFile {
101    id: String,
102    arch: Vec<String>,
103    ops: Vec<String>,
104    section: String,
105    #[serde(default)]
106    stored: Vec<String>,
107    #[serde(default)]
108    served: Vec<String>,
109    #[serde(default)]
110    when: BTreeMap<String, String>,
111    law: String,
112    count: u64,
113    #[serde(default)]
114    leaked: bool,
115    site: String,
116}
117
118/// 2026-10-02: Why the copy rules did not load.
119#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
120#[error("COPIES.toml: {0}")]
121pub struct CopyError(pub String);
122
123/// 2026-10-02: The settings a rule's `when` may read: the serve facts the loader branches on.
124/// `latent_moe` is `on` when the checkpoint's MoE runs through a latent projection
125/// (`moe_latent_size > 0`, Nemotron-3-Super), `off` otherwise. `moe_expert_tables` is the serve's
126/// plan-based decision on the MoE's transposed prefill tables (`build` or `skip`).
127pub const COPY_SETTINGS: [&str; 6] = [
128    "speculative",
129    "weight_quantization",
130    "lm_head_dtype",
131    "expert_quantization",
132    "latent_moe",
133    "moe_expert_tables",
134];
135
136/// 2026-10-02: Parse the rules text.
137pub fn parse_copies(text: &str) -> Result<Vec<CopyRule>, CopyError> {
138    let file: File = toml::from_str(text).map_err(|e| CopyError(e.to_string()))?;
139    if file.schema != 1 {
140        return Err(CopyError(format!(
141            "schema {} (this build reads 1)",
142            file.schema
143        )));
144    }
145    let mut out: Vec<CopyRule> = Vec::with_capacity(file.copy.len());
146    for r in file.copy {
147        let bad = |what: String| CopyError(format!("copy `{}`: {what}", r.id));
148        if out.iter().any(|o| o.id == r.id) {
149            return Err(bad("listed twice".into()));
150        }
151        for op in &r.ops {
152            let (name, q) = op
153                .split_once(':')
154                .map_or((op.as_str(), None), |(a, b)| (a, Some(b)));
155            let parsed = match name {
156                "linear" => OpKind::parse(name, q, None),
157                _ if q.is_some() => return Err(bad(format!("op `{op}` takes no qualifier"))),
158                _ => OpKind::parse(name, None, None),
159            };
160            if !parsed.is_ok_and(|k| k.reads_linear_weight()) {
161                return Err(bad(format!("op `{op}` is no weight-reading op")));
162            }
163        }
164        let formats = |v: &[String]| -> Result<Vec<Format>, CopyError> {
165            v.iter()
166                .map(|f| Format::parse(f).map_err(|e| bad(e.to_string())))
167                .collect()
168        };
169        for k in r.when.keys() {
170            if !COPY_SETTINGS.contains(&k.as_str()) {
171                return Err(bad(format!(
172                    "setting `{k}` is not one of {}",
173                    COPY_SETTINGS.join(", ")
174                )));
175            }
176        }
177        let section = match r.section.as_str() {
178            "main" => Section::Main,
179            "draft" => Section::Draft,
180            s => return Err(bad(format!("section `{s}` is not main or draft"))),
181        };
182        if r.count == 0 || r.arch.is_empty() || r.ops.is_empty() || r.site.is_empty() {
183            return Err(bad("needs arch, ops, a site and a positive count".into()));
184        }
185        out.push(CopyRule {
186            law: CopyLaw::parse(&r.law).ok_or_else(|| bad(format!("law `{}`", r.law)))?,
187            stored: formats(&r.stored)?,
188            served: formats(&r.served)?,
189            id: r.id,
190            arch: r.arch,
191            ops: r.ops,
192            section,
193            when: r.when,
194            count: r.count,
195            leaked: r.leaked,
196            site: r.site,
197        });
198    }
199    Ok(out)
200}
201
202impl CopyRule {
203    /// 2026-10-02: The rule applies to a node of `op` in `section` of `arch`, stored as `stored`
204    /// and served as `served`, under `settings`.
205    pub fn matches(
206        &self,
207        arch: &str,
208        op: &OpKind,
209        section: Section,
210        stored: Format,
211        served: Format,
212        settings: &BTreeMap<String, String>,
213    ) -> bool {
214        self.arch.iter().any(|a| a == arch)
215            && self.section == section
216            && self.ops.iter().any(|o| *o == op.name())
217            && (self.stored.is_empty() || self.stored.contains(&stored))
218            && (self.served.is_empty() || self.served.contains(&served))
219            && self.when.iter().all(|(k, v)| settings.get(k) == Some(v))
220    }
221}
metrale_circuit/memory/copies.rs

metrale_circuit/memory/
copies.rs
