Skip to main content

cortiq_core/
format.rs

1//! CMF v2 binary container — envelope, tensor directory, mmap access.
2//!
3//! See `docs/CMF_V2_SPEC.md`. Layout summary:
4//!
5//! ```text
6//! [0x00]  magic "CMF\x01" | version u32 = 2 | flags u32 | required_features u32
7//! [0x10]  header_off/len | dir_off/len | data_off/len   (u64 LE each)
8//! [0x40]  masks_off/len  | vocab_off/len | index_off/len
9//! [0x70]  16 reserved bytes (zero)
10//! [0x80]  header JSON → tensor directory → weight blob (4096-aligned,
11//!         tensors 64-aligned) → masks → vocab → sparse index
12//! ```
13//!
14//! The tensor directory is the ONLY source of truth for the weight blob
15//! layout — there is no computable layout, by design (v1 bug class #1).
16//! Every validation failure is a hard error: no silent fallbacks.
17
18use crate::hash::hash64;
19use crate::mask::{MaskCatalog, TaskMask, decode_masks_section, encode_masks_section};
20use crate::quant::expected_nbytes;
21use crate::types::{ModelArch, QuantType, TensorDtype};
22use serde::{Deserialize, Serialize};
23use std::collections::HashMap;
24use std::fs::File;
25use std::io::{self, BufWriter, Seek, SeekFrom, Write};
26use std::path::{Path, PathBuf};
27
28pub const CMF_MAGIC: [u8; 4] = *b"CMF\x01";
29pub const CMF_VERSION: u32 = 2;
30pub const ENVELOPE_LEN: usize = 128;
31/// Weight blob is page-aligned for mmap.
32pub const DATA_ALIGNMENT: u64 = 4096;
33/// Every tensor inside the blob is 64-byte aligned (SIMD / cache line).
34pub const TENSOR_ALIGNMENT: u64 = 64;
35/// Tensors at least this large are additionally page-aligned to
36/// [`LARGE_TENSOR_ALIGN`], so a cold skill / MoE-expert / mask weight sits on
37/// its own page(s): the "unused weights cost 0 RSS" guarantee then holds at
38/// page granularity (a lazily-paged tensor pulls exactly its own bytes, not a
39/// neighbour's), and per-layer `madvise(WILLNEED)` covers clean ranges. Small
40/// tensors (norms, biases, 1-D f16) keep the 64-byte SIMD alignment so the
41/// padding stays negligible.
42pub const LARGE_TENSOR_MIN: u64 = 16 * 1024;
43/// Page alignment applied to large tensors. 4096 is the common page size
44/// (x86, most ARM/Android); it is a multiple of [`TENSOR_ALIGNMENT`], so
45/// existing readers — which only require `off % 64 == 0` — accept these files
46/// unchanged. Purely a writer-side, backward-compatible layout choice.
47pub const LARGE_TENSOR_ALIGN: u64 = 4096;
48/// One directory record is 56 bytes (see `.vmfc` v2).
49pub const DIR_RECORD_LEN: usize = 56;
50pub const DIR_MAX_NDIM: usize = 6;
51
52/// `required_features` bits. A reader MUST refuse a file with any bit
53/// it does not support.
54pub mod features {
55    pub const TENSOR_DIR: u32 = 1 << 0;
56    pub const BINARY_MASKS: u32 = 1 << 1;
57    pub const QUANT_2F: u32 = 1 << 2;
58    pub const DELTA_MASKS: u32 = 1 << 3;
59    pub const HOT_PACKS: u32 = 1 << 4;
60
61    /// Features this reader implements today.
62    pub const SUPPORTED: u32 = TENSOR_DIR | BINARY_MASKS | QUANT_2F;
63}
64
65/// JSON header — architecture and provenance (human-readable part;
66/// machine-critical data lives in binary sections).
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct CmfHeader {
69    #[serde(default = "default_format")]
70    pub format: String,
71    pub version: u32,
72    pub arch: ModelArch,
73    /// Informational default; per-tensor truth is in the directory.
74    pub quant_type: QuantType,
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub provenance: Option<serde_json::Value>,
77    /// Chat/eos bundle (spec §6.1): the file — not the binary — defines
78    /// chat behavior. Additive: absent in older files.
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub tokenizer_config: Option<TokenizerBundle>,
81    /// Section-level integrity (spec §8.1): hex hash64 of the raw bytes
82    /// of the optional sections. header/dir hashes live in the envelope
83    /// reserved bytes — JSON cannot protect the JSON that carries it.
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub section_hashes: Option<SectionHashes>,
86    /// Per-skill records (spec §9): replacement tensors live in the
87    /// directory as `skill.{id}.{name}`; this registry carries the
88    /// selection descriptor and the honest quality contract.
89    #[serde(default, skip_serializing_if = "Vec::is_empty")]
90    pub skills: Vec<SkillRecord>,
91    /// Sharding (spec §10): this file is shard `no` of `count`; every
92    /// shard is a standalone valid .cmf carrying a tensor subset.
93    /// Naming convention: `…-{no:05}-of-{count:05}.cmf`.
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub shard: Option<ShardInfo>,
96    /// Measured confidence calibration (B1): a temperature fit on held-out
97    /// so the displayed Born-mass confidence is a true property of the
98    /// model (softmax(logits/T)), not a raw estimate. Additive; absent =
99    /// use raw (T=1). Written by `set_calibration.py` after `cortiq
100    /// calibrate` measures the reliability/ECE.
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub calibration: Option<Calibration>,
103}
104
105/// Confidence-calibration record (spec §6.2). `temperature` scales the
106/// logits before softmax when reporting confidence; `ece_before`/`after`
107/// are the measured Expected Calibration Error (honest provenance).
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct Calibration {
110    pub temperature: f32,
111    #[serde(default, skip_serializing_if = "Option::is_none")]
112    pub ece_before: Option<f32>,
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub ece_after: Option<f32>,
115}
116
117/// Shard coordinates (1-based, gguf-split style).
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct ShardInfo {
120    pub no: usize,
121    pub count: usize,
122}
123
124/// Recon-argmin routing parameters (spec §9; P1 signal-consistency):
125/// E = ‖(φ−mean) − B·Bᵀ(φ−mean)‖² / ‖φ−mean‖²; pick argmin over skills.
126#[derive(Debug, Clone, Serialize, Deserialize)]
127pub struct SelectionDescriptor {
128    /// "mse" (normalized reconstruction error) — the only metric today.
129    pub metric: String,
130    /// Backbone layer whose mean-pooled hidden is φ(x).
131    pub phi_layer: usize,
132    /// Subspace mean, f16 LE base64, len = hidden.
133    pub mean: String,
134    /// Orthonormal basis rows, f16 LE base64, len = rank·hidden.
135    pub basis: String,
136    pub rank: usize,
137}
138
139/// One skill of the swarm (spec §9; Patent 15 per-skill record).
140#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct SkillRecord {
142    pub id: String,
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub name: Option<String>,
145    /// Layers this skill specializes (a proper subset).
146    #[serde(default)]
147    pub layers: Vec<usize>,
148    /// Selection descriptor for recon-argmin routing (208c, P1):
149    /// per-skill affine subspace over φ(x) = mean-pooled hidden state.
150    #[serde(default, skip_serializing_if = "Option::is_none")]
151    pub selection: Option<SelectionDescriptor>,
152    /// Optional input-mask task name (208b), applied with the skill.
153    #[serde(default, skip_serializing_if = "Option::is_none")]
154    pub input_mask_task: Option<String>,
155    /// Measured quality (claim 16): overlaid vs backbone, held-out.
156    #[serde(default, skip_serializing_if = "Option::is_none")]
157    pub quality: Option<serde_json::Value>,
158}
159
160/// Hex-encoded hash64 per optional section (u64 as JSON number would
161/// lose precision past 2^53).
162#[derive(Debug, Clone, Default, Serialize, Deserialize)]
163pub struct SectionHashes {
164    #[serde(default, skip_serializing_if = "Option::is_none")]
165    pub masks: Option<String>,
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub vocab: Option<String>,
168    #[serde(default, skip_serializing_if = "Option::is_none")]
169    pub index: Option<String>,
170}
171
172/// Chat template + generation stop tokens carried by the container.
173#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
174pub struct TokenizerBundle {
175    /// Jinja chat template (chat_template.jinja / tokenizer_config.json)
176    #[serde(default, skip_serializing_if = "Option::is_none")]
177    pub chat_template: Option<String>,
178    /// All ids that terminate generation (generation_config + im_end)
179    #[serde(default)]
180    pub eos_token_ids: Vec<u32>,
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub bos_token_id: Option<u32>,
183    #[serde(default, skip_serializing_if = "Option::is_none")]
184    pub pad_token_id: Option<u32>,
185}
186
187fn default_format() -> String {
188    "cmf".to_string()
189}
190
191/// One tensor directory entry.
192#[derive(Debug, Clone, PartialEq, Eq)]
193pub struct TensorEntry {
194    pub name: String,
195    pub dtype: TensorDtype,
196    pub shape: Vec<usize>,
197    /// Offset relative to the OWNING shard's `data_off`, multiple of 64.
198    pub off: u64,
199    pub nbytes: u64,
200    /// Runtime-only: which shard's mmap holds the bytes (0 for the
201    /// single-file case; not part of the 56-byte record).
202    pub shard: usize,
203    /// `hash64` of the tensor bytes.
204    pub hash: u64,
205}
206
207impl TensorEntry {
208    pub fn n_elems(&self) -> usize {
209        self.shape.iter().product()
210    }
211}
212
213/// Input for the Rust writer: one tensor with its encoded bytes.
214#[derive(Debug, Clone)]
215pub struct TensorSpec {
216    pub name: String,
217    pub dtype: TensorDtype,
218    pub shape: Vec<usize>,
219    pub data: Vec<u8>,
220}
221
222/// `TensorSpec` with a borrowed payload — see [`CmfModel::write_ref`].
223pub struct TensorSpecRef<'a> {
224    pub name: String,
225    pub dtype: TensorDtype,
226    pub shape: Vec<usize>,
227    pub data: &'a [u8],
228}
229
230/// Sparse index entry — precomputed per-task per-layer active group IDs.
231#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
232pub struct SparseIndexEntry {
233    pub task_id: u32,
234    pub layer_idx: usize,
235    /// Active quant-group indices for FFN (sorted, group = 32 neurons).
236    pub active_ffn_groups: Vec<u16>,
237    /// Active head indices for attention (sorted).
238    pub active_heads: Vec<u8>,
239}
240
241/// Section ranges parsed from the fixed envelope.
242#[derive(Debug, Clone, Copy, Default)]
243struct Envelope {
244    required_features: u32,
245    header: (u64, u64),
246    dir: (u64, u64),
247    data: (u64, u64),
248    masks: (u64, u64),
249    vocab: (u64, u64),
250    index: (u64, u64),
251    /// hash64 of the header JSON bytes (reserved [0x70]); 0 = absent.
252    header_hash: u64,
253    /// hash64 of the tensor-directory bytes (reserved [0x78]); 0 = absent.
254    dir_hash: u64,
255}
256
257enum Backing {
258    Mmap(memmap2::Mmap),
259    Owned(Vec<u8>),
260}
261
262impl Backing {
263    fn bytes(&self) -> &[u8] {
264        match self {
265            Backing::Mmap(m) => m,
266            Backing::Owned(v) => v,
267        }
268    }
269}
270
271/// A loaded CMF model: metadata owned, weights zero-copy via mmap.
272pub struct CmfModel {
273    pub path: PathBuf,
274    /// Identifies THIS open of the file, monotonically. GPU backends cache
275    /// device weights by (model, tensor); keying that on the mapping's
276    /// address made a reloaded model inherit the previous one's buffers
277    /// whenever the new mmap landed where the old had been — silent wrong
278    /// weights in any process that unloads and reloads, which a server does.
279    uid: u64,
280    pub header: CmfHeader,
281    pub required_features: u32,
282    pub tensors: Vec<TensorEntry>,
283    /// name-hash → tensor index. Keying on the hash (not the name) avoids
284    /// cloning every tensor name into the map at `open()` — that halves the
285    /// open-time allocations and the map's footprint, which matters for large
286    /// MoE / skills files with tens of thousands of tensors. A genuine 64-bit
287    /// hash collision between two *distinct* names — astronomically unlikely —
288    /// lands in `name_overflow`, so lookups stay exact.
289    by_name: HashMap<u64, u32>,
290    name_overflow: Vec<u32>,
291    pub masks: MaskCatalog,
292    pub sparse_index: Vec<SparseIndexEntry>,
293    /// Embedded tokenizer.json bytes, if present.
294    pub vocab: Option<Vec<u8>>,
295    backing: Backing,
296    data_off: u64,
297    envelope: Envelope,
298    /// Shards 2..N (spec §10): (backing, data_off) per extra file;
299    /// `TensorEntry.shard` 0 = this file, i>0 = extra_shards[i-1].
300    extra_shards: Vec<(Backing, u64)>,
301}
302
303impl std::fmt::Debug for CmfModel {
304    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305        f.debug_struct("CmfModel")
306            .field("path", &self.path)
307            .field("arch", &self.header.arch.arch_name)
308            .field("tensors", &self.tensors.len())
309            .field("masks", &self.masks.masks.len())
310            .finish()
311    }
312}
313
314/// Hands out a fresh id per open. Wraps only after 2^64 opens.
315static MODEL_UID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
316
317impl CmfModel {
318    /// A key that is unique to this open of the file and never recycled.
319    /// Backends that cache anything derived from the weights must key on
320    /// this, not on the mapping's address, which the allocator reuses.
321    pub fn uid(&self) -> u64 {
322        self.uid
323    }
324
325    /// Open and strictly validate a CMF v2 file. Any inconsistency is an
326    /// error — this function never substitutes defaults.
327    pub fn open(path: impl AsRef<Path>) -> Result<Self, CmfError> {
328        let path = path.as_ref().to_path_buf();
329        if !path.exists() {
330            return Err(CmfError::FileNotFound(path.display().to_string()));
331        }
332        let file = File::open(&path)?;
333        let file_len = file.metadata()?.len();
334
335        let backing = match unsafe { memmap2::MmapOptions::new().map(&file) } {
336            Ok(m) => {
337                // Decode touches every weight page each token, so tell
338                // the kernel up front: WillNeed front-loads readahead
339                // (first-token page-fault storm becomes streaming I/O —
340                // this is TTFT on phones, where the file is a large
341                // share of RAM). Advisory only: a memory-pressured
342                // kernel is free to ignore it. CMF_MMAP_ADVISE=0 turns
343                // it off; CMF_MLOCK=1 additionally tries to pin the
344                // mapping (needs RLIMIT_MEMLOCK headroom — refusal is
345                // logged, not fatal).
346                #[cfg(unix)]
347                {
348                    // …unless these pages are headed for a device that will
349                    // then drop them. Reading the file ahead only to throw it
350                    // out behind the uploader has the kernel fetching the same
351                    // bytes twice. The two advices contradict each other, so
352                    // the one that asked for eviction wins.
353                    //
354                    // The conditions are the uploader's, approximated from what
355                    // is knowable at open time (the GPU is not up yet): Linux,
356                    // because `evict_ranges` is a no-op elsewhere; a backend
357                    // actually requested, because a CPU run wants the readahead;
358                    // and eviction not turned off. On UMA the mapping IS the
359                    // working copy and nothing is evicted — that is why this
360                    // cannot key on `CMF_GPU` alone.
361                    let evicting = cfg!(target_os = "linux")
362                        && std::env::var("CMF_GPU").is_ok_and(|v| v != "0" && v != "off")
363                        && std::env::var("CMF_UPLOAD_EVICT")
364                            .map(|v| v != "0")
365                            .unwrap_or(true);
366                    if !evicting
367                        && std::env::var("CMF_MMAP_ADVISE")
368                            .map(|v| v != "0")
369                            .unwrap_or(true)
370                    {
371                        let _ = m.advise(memmap2::Advice::WillNeed);
372                    }
373                    if std::env::var("CMF_MLOCK")
374                        .map(|v| v == "1")
375                        .unwrap_or(false)
376                    {
377                        if let Err(e) = m.lock() {
378                            tracing::warn!(
379                                "CMF_MLOCK=1: mlock refused ({e}) — continuing unpinned"
380                            );
381                        }
382                    }
383                }
384                Backing::Mmap(m)
385            }
386            Err(e) => {
387                tracing::warn!("mmap failed ({e}), reading file into memory");
388                Backing::Owned(std::fs::read(&path)?)
389            }
390        };
391
392        let env = Self::parse_envelope(backing.bytes(), file_len)?;
393
394        let bytes = backing.bytes();
395        let section = |off: u64, len: u64| -> &[u8] { &bytes[off as usize..(off + len) as usize] };
396
397        // Header JSON
398        let header: CmfHeader = serde_json::from_slice(section(env.header.0, env.header.1))
399            .map_err(|e| CmfError::Parse(format!("header JSON: {e}")))?;
400
401        // Tensor directory
402        let tensors = Self::decode_directory(section(env.dir.0, env.dir.1))?;
403        for t in &tensors {
404            if t.off % TENSOR_ALIGNMENT != 0 {
405                return Err(CmfError::Bounds(format!(
406                    "tensor '{}': offset {} not 64-aligned",
407                    t.name, t.off
408                )));
409            }
410            let tensor_end = t.off.checked_add(t.nbytes).ok_or_else(|| {
411                CmfError::Bounds(format!("tensor '{}': offset + length overflows", t.name))
412            })?;
413            if tensor_end > env.data.1 {
414                return Err(CmfError::Bounds(format!(
415                    "tensor '{}': [{}, {}) exceeds data section ({} bytes)",
416                    t.name, t.off, tensor_end, env.data.1
417                )));
418            }
419            t.shape
420                .iter()
421                .try_fold(1usize, |n, &dim| n.checked_mul(dim))
422                .ok_or_else(|| {
423                    CmfError::Bounds(format!(
424                        "tensor '{}': shape product overflows usize",
425                        t.name
426                    ))
427                })?;
428            if let Some(expect) = expected_nbytes(t.dtype, &t.shape) {
429                if expect as u64 != t.nbytes {
430                    return Err(CmfError::Bounds(format!(
431                        "tensor '{}': nbytes {} != expected {} for {:?}{:?}",
432                        t.name, t.nbytes, expect, t.dtype, t.shape
433                    )));
434                }
435            }
436            // Payload-dependent lengths (vbit): exact check against the
437            // width header, bounds-before-slice (roadmap §4.9).
438            if matches!(t.dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
439                let payload = section(env.data.0 + t.off, t.nbytes);
440                crate::quant::validate_payload(t.dtype, &t.shape, payload)
441                    .map_err(|e| CmfError::Bounds(format!("tensor '{}': {e}", t.name)))?;
442            }
443        }
444        // Duplicate names would silently shadow each other in the
445        // HashMap (directory scan and by_name would disagree) — refuse
446        // the file instead (roadmap §4.9).
447        let mut by_name: HashMap<u64, u32> = HashMap::with_capacity(tensors.len());
448        let mut name_overflow: Vec<u32> = Vec::new();
449        for i in 0..tensors.len() {
450            let h = hash64(tensors[i].name.as_bytes());
451            match by_name.get(&h) {
452                Some(&j) if tensors[j as usize].name == tensors[i].name => {
453                    return Err(CmfError::Parse(format!(
454                        "duplicate tensor name '{}' in directory",
455                        tensors[i].name
456                    )));
457                }
458                Some(_) => name_overflow.push(i as u32), // hash collision of distinct names
459                None => {
460                    by_name.insert(h, i as u32);
461                }
462            }
463        }
464
465        // Masks
466        let masks = if env.masks.1 > 0 {
467            decode_masks_section(section(env.masks.0, env.masks.1), &header.arch)
468                .map_err(CmfError::Parse)?
469        } else {
470            MaskCatalog::empty()
471        };
472
473        // Vocab (tokenizer.json)
474        let vocab = if env.vocab.1 > 0 {
475            Some(section(env.vocab.0, env.vocab.1).to_vec())
476        } else {
477            None
478        };
479
480        // Sparse index
481        let sparse_index = if env.index.1 > 0 {
482            decode_sparse_index(section(env.index.0, env.index.1))?
483        } else {
484            vec![]
485        };
486
487        tracing::info!(
488            "Opened CMF v2: {} | {} tensors | {} masks | vocab {} | {:.1} MB",
489            header.arch.arch_name,
490            tensors.len(),
491            masks.masks.len(),
492            if vocab.is_some() { "embedded" } else { "none" },
493            file_len as f64 / 1e6
494        );
495
496        Ok(Self {
497            uid: MODEL_UID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
498            path,
499            header,
500            required_features: env.required_features,
501            tensors,
502            by_name,
503            name_overflow,
504            masks,
505            sparse_index,
506            vocab,
507            backing,
508            data_off: env.data.0,
509            envelope: env,
510            extra_shards: Vec::new(),
511        })
512    }
513
514    /// Open a sharded model (spec §10): pass shard 1; siblings found by
515    /// the `-{no:05}-of-{count:05}.cmf` convention. Directories merge;
516    /// masks/vocab/index/skills come from shard 1.
517    pub fn open_sharded(path: impl AsRef<Path>) -> Result<Self, CmfError> {
518        let path = path.as_ref();
519        let mut first = Self::open(path)?;
520        let Some(info) = first.header.shard.clone() else {
521            return Ok(first); // not sharded — plain open
522        };
523        if info.no != 1 {
524            return Err(CmfError::Parse(format!(
525                "open shard 1, not {} (of {})",
526                info.no, info.count
527            )));
528        }
529        let name = path
530            .file_name()
531            .and_then(|n| n.to_str())
532            .ok_or_else(|| CmfError::Parse("bad shard path".into()))?;
533        let tag1 = format!("-{:05}-of-{:05}.cmf", 1, info.count);
534        if !name.ends_with(&tag1) {
535            return Err(CmfError::Parse(format!(
536                "shard file must end with '{tag1}' (got '{name}')"
537            )));
538        }
539        let stem = &name[..name.len() - tag1.len()];
540        for no in 2..=info.count {
541            let sib = path.with_file_name(format!("{stem}-{:05}-of-{:05}.cmf", no, info.count));
542            let sh = Self::open(&sib)?;
543            match &sh.header.shard {
544                Some(si) if si.no == no && si.count == info.count => {}
545                other => {
546                    return Err(CmfError::Parse(format!(
547                        "{}: wrong shard coords {other:?}",
548                        sib.display()
549                    )));
550                }
551            }
552            let shard_idx = first.extra_shards.len() + 1;
553            first.extra_shards.push((sh.backing, sh.envelope.data.0));
554            for mut t in sh.tensors {
555                t.shard = shard_idx;
556                let idx = first.tensors.len() as u32;
557                let h = hash64(t.name.as_bytes());
558                match first.by_name.get(&h) {
559                    Some(&j) if first.tensors[j as usize].name == t.name => {
560                        return Err(CmfError::Parse(format!(
561                            "duplicate tensor name '{}' across shards",
562                            t.name
563                        )));
564                    }
565                    Some(_) => first.name_overflow.push(idx),
566                    None => {
567                        first.by_name.insert(h, idx);
568                    }
569                }
570                first.tensors.push(t);
571            }
572        }
573        tracing::info!(
574            "sharded model: {} files, {} tensors total",
575            info.count,
576            first.tensors.len()
577        );
578        Ok(first)
579    }
580
581    fn parse_envelope(bytes: &[u8], file_len: u64) -> Result<Envelope, CmfError> {
582        if bytes.len() < ENVELOPE_LEN {
583            return Err(CmfError::Bounds(format!(
584                "file too small for CMF envelope: {} bytes",
585                bytes.len()
586            )));
587        }
588        if bytes[0..4] != CMF_MAGIC {
589            return Err(CmfError::InvalidMagic);
590        }
591        let u32le = |o: usize| u32::from_le_bytes(bytes[o..o + 4].try_into().unwrap());
592        let u64le = |o: usize| u64::from_le_bytes(bytes[o..o + 8].try_into().unwrap());
593
594        let version = u32le(4);
595        if version != CMF_VERSION {
596            return Err(CmfError::UnsupportedVersion(version));
597        }
598        let _flags = u32le(8); // reserved
599        let required_features = u32le(12);
600        let unknown = required_features & !features::SUPPORTED;
601        if unknown != 0 {
602            return Err(CmfError::UnsupportedFeature(unknown));
603        }
604
605        let env = Envelope {
606            required_features,
607            header: (u64le(0x10), u64le(0x18)),
608            dir: (u64le(0x20), u64le(0x28)),
609            data: (u64le(0x30), u64le(0x38)),
610            masks: (u64le(0x40), u64le(0x48)),
611            vocab: (u64le(0x50), u64le(0x58)),
612            index: (u64le(0x60), u64le(0x68)),
613            header_hash: u64le(0x70),
614            dir_hash: u64le(0x78),
615        };
616
617        for (name, (off, len), required) in [
618            ("header", env.header, true),
619            ("dir", env.dir, true),
620            ("data", env.data, false),
621            ("masks", env.masks, false),
622            ("vocab", env.vocab, false),
623            ("index", env.index, false),
624        ] {
625            if required && len == 0 {
626                return Err(CmfError::Bounds(format!("section '{name}' is required")));
627            }
628            if len > 0
629                && off
630                    .checked_add(len)
631                    .map(|end| end > file_len)
632                    .unwrap_or(true)
633            {
634                return Err(CmfError::Bounds(format!(
635                    "section '{name}' [{off}, {}) exceeds file ({file_len} bytes)",
636                    off.saturating_add(len)
637                )));
638            }
639            if len > 0
640                && (usize::try_from(off).is_err()
641                    || usize::try_from(len).is_err()
642                    || usize::try_from(off + len).is_err())
643            {
644                return Err(CmfError::Bounds(format!(
645                    "section '{name}' cannot be addressed on this platform"
646                )));
647            }
648        }
649        if env.data.1 > 0 && env.data.0 % DATA_ALIGNMENT != 0 {
650            return Err(CmfError::Bounds(format!(
651                "data section offset {} not {}-aligned",
652                env.data.0, DATA_ALIGNMENT
653            )));
654        }
655        Ok(env)
656    }
657
658    fn decode_directory(bytes: &[u8]) -> Result<Vec<TensorEntry>, CmfError> {
659        if bytes.len() < 16 {
660            return Err(CmfError::Parse("tensor directory too short".into()));
661        }
662        let count = u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
663        let pool_off = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
664        let records_len = count
665            .checked_mul(DIR_RECORD_LEN)
666            .ok_or_else(|| CmfError::Parse("tensor directory record count overflows".into()))?;
667        let records_end = 16usize
668            .checked_add(records_len)
669            .ok_or_else(|| CmfError::Parse("tensor directory size overflows".into()))?;
670        if records_end > bytes.len() || pool_off > bytes.len() || pool_off < records_end {
671            return Err(CmfError::Parse(format!(
672                "tensor directory malformed: count={count}, pool_off={pool_off}, len={}",
673                bytes.len()
674            )));
675        }
676        let pool = &bytes[pool_off..];
677
678        let mut out = Vec::with_capacity(count);
679        for i in 0..count {
680            let r = &bytes[16 + i * DIR_RECORD_LEN..16 + (i + 1) * DIR_RECORD_LEN];
681            let name_off = u32::from_le_bytes(r[0..4].try_into().unwrap()) as usize;
682            let name_len = u16::from_le_bytes(r[4..6].try_into().unwrap()) as usize;
683            let dtype_id = r[6];
684            let ndim = r[7] as usize;
685            if ndim > DIR_MAX_NDIM {
686                return Err(CmfError::Parse(format!("tensor #{i}: ndim {ndim} > 6")));
687            }
688            let mut shape = Vec::with_capacity(ndim);
689            for d in 0..ndim {
690                shape.push(
691                    u32::from_le_bytes(r[8 + d * 4..12 + d * 4].try_into().unwrap()) as usize,
692                );
693            }
694            let off = u64::from_le_bytes(r[32..40].try_into().unwrap());
695            let nbytes = u64::from_le_bytes(r[40..48].try_into().unwrap());
696            let hash = u64::from_le_bytes(r[48..56].try_into().unwrap());
697
698            let name_end = name_off
699                .checked_add(name_len)
700                .ok_or_else(|| CmfError::Parse(format!("tensor #{i}: name range overflows")))?;
701            if name_end > pool.len() {
702                return Err(CmfError::Parse(format!("tensor #{i}: name out of pool")));
703            }
704            let name = std::str::from_utf8(&pool[name_off..name_end])
705                .map_err(|_| CmfError::Parse(format!("tensor #{i}: name is not UTF-8")))?
706                .to_string();
707            let dtype = TensorDtype::from_id(dtype_id).ok_or(CmfError::UnknownDtype(dtype_id))?;
708
709            out.push(TensorEntry {
710                name,
711                dtype,
712                shape,
713                off,
714                nbytes,
715                shard: 0,
716                hash,
717            });
718        }
719        Ok(out)
720    }
721
722    // ───────────────────────── access ─────────────────────────
723
724    pub fn arch(&self) -> &ModelArch {
725        &self.header.arch
726    }
727
728    pub fn tensor(&self, name: &str) -> Option<&TensorEntry> {
729        self.tensor_index(name).map(|i| &self.tensors[i])
730    }
731
732    /// Directory index of a tensor by name (same resolution as
733    /// [`Self::tensor`] — engines must not re-scan the directory). O(1) via the
734    /// name-hash index; the name is verified against the entry so a hash
735    /// collision can never return the wrong tensor, and the rare distinct-name
736    /// collision falls back to the tiny overflow list.
737    pub fn tensor_index(&self, name: &str) -> Option<usize> {
738        let h = hash64(name.as_bytes());
739        if let Some(&i) = self.by_name.get(&h) {
740            if self.tensors[i as usize].name == name {
741                return Some(i as usize);
742            }
743        }
744        self.name_overflow
745            .iter()
746            .copied()
747            .find(|&i| self.tensors[i as usize].name == name)
748            .map(|i| i as usize)
749    }
750
751    /// Tensor-source indirection (spec §9, Patent 15 fig3/302): the
752    /// skill's replacement is read IN PLACE OF the backbone tensor —
753    /// either/or, never combined. None skill → backbone directly.
754    pub fn resolve_tensor(&self, name: &str, skill: Option<&str>) -> Option<&TensorEntry> {
755        if let Some(sid) = skill {
756            if let Some(t) = self.tensor(&format!("skill.{sid}.{name}")) {
757                return Some(t);
758            }
759        }
760        self.tensor(name)
761    }
762
763    /// The per-skill delta index view (claim 2): directory entries of
764    /// one skill — exactly the byte ranges lazy loading pages in.
765    pub fn skill_tensors(&self, skill_id: &str) -> impl Iterator<Item = &TensorEntry> {
766        let prefix = format!("skill.{skill_id}.");
767        self.tensors
768            .iter()
769            .filter(move |t| t.name.starts_with(&prefix))
770    }
771
772    /// Zero-copy bytes of a tensor from the mmap'd data section.
773    pub fn tensor_bytes(&self, name: &str) -> Result<&[u8], CmfError> {
774        let entry = self
775            .tensor(name)
776            .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
777        Ok(self.entry_bytes(entry))
778    }
779
780    /// All bytes of the primary mapping (GPU path: no-copy Metal buffer
781    /// over the same mmap — unified memory, zero copying).
782    pub fn primary_bytes(&self) -> &[u8] {
783        self.backing.bytes()
784    }
785
786    /// Absolute offset of the tensor within the primary mapping
787    /// (None for tensors from sibling shards).
788    pub fn entry_abs_offset(&self, entry: &TensorEntry) -> Option<usize> {
789        (entry.shard == 0).then(|| (self.data_off + entry.off) as usize)
790    }
791
792    pub fn entry_bytes(&self, entry: &TensorEntry) -> &[u8] {
793        let (bytes, data_off) = if entry.shard == 0 {
794            (self.backing.bytes(), self.data_off)
795        } else {
796            let (b, o) = &self.extra_shards[entry.shard - 1];
797            (b.bytes(), *o)
798        };
799        let start = (data_off + entry.off) as usize;
800        &bytes[start..start + entry.nbytes as usize]
801    }
802
803    /// The CPU is done with these byte ranges of the primary mapping
804    /// (absolute offsets, as `entry_abs_offset` hands them out): drop them
805    /// from the resident set and let the page cache release the file pages.
806    /// Both calls are advisory and the mapping stays valid — a range that
807    /// gets touched again re-faults from disk, so a caller can only cost
808    /// time here, never correctness. Ranges are aligned OUTWARD to page
809    /// boundaries; the neighbours those edges claw in re-fault the same way.
810    /// Linux + mmap backing only; everywhere else a no-op.
811    pub fn evict_ranges(&self, ranges: &[(usize, usize)]) {
812        #[cfg(target_os = "linux")]
813        {
814            let Backing::Mmap(m) = &self.backing else {
815                return;
816            };
817            let page = 4096usize;
818            // One fd for the whole batch: fadvise targets the inode's page
819            // cache, any fd on the same file will do.
820            use std::os::unix::io::AsRawFd;
821            let file = File::open(&self.path).ok();
822            for &(off, len) in ranges {
823                if len == 0 || off.saturating_add(len) > m.len() {
824                    continue;
825                }
826                let start = off & !(page - 1);
827                let end = off + len;
828                let alen = end.next_multiple_of(page).min(m.len()) - start;
829                // SAFETY: read-only MAP_SHARED file mapping — DONTNEED here
830                // only drops clean pages; the next access re-faults them.
831                let _ = unsafe {
832                    m.unchecked_advise_range(memmap2::UncheckedAdvice::DontNeed, start, alen)
833                };
834                if let Some(f) = &file {
835                    // SAFETY: plain fd + numeric range; advisory by contract.
836                    unsafe {
837                        libc::posix_fadvise(
838                            f.as_raw_fd(),
839                            start as libc::off_t,
840                            alen as libc::off_t,
841                            libc::POSIX_FADV_DONTNEED,
842                        );
843                    }
844                }
845            }
846        }
847        #[cfg(not(target_os = "linux"))]
848        let _ = ranges;
849    }
850
851    /// Tensors belonging to layer `i` (prefix `model.layers.{i}.`).
852    pub fn layer_tensors(&self, layer_idx: usize) -> Vec<&TensorEntry> {
853        let prefix = format!("model.layers.{layer_idx}.");
854        self.tensors
855            .iter()
856            .filter(|t| t.name.starts_with(&prefix))
857            .collect()
858    }
859
860    /// Total parameter count estimated from matrix tensors (ndim ≥ 2).
861    pub fn total_param_count(&self) -> u64 {
862        self.tensors
863            .iter()
864            .filter(|t| t.shape.len() >= 2)
865            .map(|t| t.n_elems() as u64)
866            .sum()
867    }
868
869    /// Recode selected tensors IN PLACE: each new payload must fit its old
870    /// slot, the entry keeps its offset and the file keeps its length — the
871    /// bytes between the new end and the old simply go dark (every reader
872    /// walks the directory, nothing addresses the gap). This is what lets a
873    /// published 100+ GB file change a tensor's layout on a disk too small
874    /// to hold two copies of it. Patches are `(directory index, new dtype,
875    /// new payload)`; entry hashes and the directory hash are recomputed so
876    /// `verify` stays clean. Not atomic: a crash between the payload writes
877    /// and the directory write leaves the old dtype over new bytes — verify
878    /// (or re-fetch the source) after an interrupted run.
879    pub fn recode_entries_in_place(
880        path: &str,
881        patches: &[(usize, TensorDtype, Vec<u8>)],
882    ) -> Result<(), CmfError> {
883        use std::io::{Read, Seek, SeekFrom, Write};
884        let mut f = std::fs::OpenOptions::new().read(true).write(true).open(path)?;
885        let file_len = f.metadata()?.len();
886        let mut head = vec![0u8; ENVELOPE_LEN];
887        f.read_exact(&mut head)?;
888        let env = Self::parse_envelope(&head, file_len)?;
889
890        let mut dir = vec![0u8; env.dir.1 as usize];
891        f.seek(SeekFrom::Start(env.dir.0))?;
892        f.read_exact(&mut dir)?;
893        let count = u64::from_le_bytes(dir[0..8].try_into().unwrap()) as usize;
894
895        for (i, dtype, data) in patches {
896            if *i >= count {
897                return Err(CmfError::Bounds(format!(
898                    "recode: tensor #{i} out of directory ({count} entries)"
899                )));
900            }
901            let rb = 16 + i * DIR_RECORD_LEN;
902            let rec = &mut dir[rb..rb + DIR_RECORD_LEN];
903            let off = u64::from_le_bytes(rec[32..40].try_into().unwrap());
904            let old_n = u64::from_le_bytes(rec[40..48].try_into().unwrap());
905            if data.len() as u64 > old_n {
906                return Err(CmfError::Bounds(format!(
907                    "recode: tensor #{i} payload {} > slot {old_n}",
908                    data.len()
909                )));
910            }
911            f.seek(SeekFrom::Start(env.data.0 + off))?;
912            f.write_all(data)?;
913            rec[6] = dtype.id();
914            rec[40..48].copy_from_slice(&(data.len() as u64).to_le_bytes());
915            rec[48..56].copy_from_slice(&hash64(data).to_le_bytes());
916        }
917
918        f.seek(SeekFrom::Start(env.dir.0))?;
919        f.write_all(&dir)?;
920        f.seek(SeekFrom::Start(0x78))?;
921        f.write_all(&hash64(&dir).to_le_bytes())?;
922        f.sync_all()?;
923        Ok(())
924    }
925
926    /// Recompute all tensor hashes; returns human-readable problems
927    /// (empty = file intact).
928    pub fn verify(&self) -> Vec<String> {
929        let mut problems = Vec::new();
930
931        // Section-level integrity (spec §8.1). Zero/absent = legacy file.
932        let bytes = self.backing.bytes();
933        let env = &self.envelope;
934        let sect = |(off, len): (u64, u64)| &bytes[off as usize..(off + len) as usize];
935        let check = |name: &str, stored: u64, span: (u64, u64)| -> Option<String> {
936            if stored != 0 && span.1 > 0 {
937                let actual = hash64(sect(span));
938                if actual != stored {
939                    return Some(format!(
940                        "section '{name}': hash mismatch (stored {stored:016x}, \
941                         actual {actual:016x})"
942                    ));
943                }
944            }
945            None
946        };
947        problems.extend(check("header", env.header_hash, env.header));
948        problems.extend(check("dir", env.dir_hash, env.dir));
949        if let Some(sh) = &self.header.section_hashes {
950            for (name, hex, span) in [
951                ("masks", &sh.masks, env.masks),
952                ("vocab", &sh.vocab, env.vocab),
953                ("index", &sh.index, env.index),
954            ] {
955                if let Some(hex) = hex {
956                    match u64::from_str_radix(hex, 16) {
957                        Ok(stored) => problems.extend(check(name, stored, span)),
958                        Err(_) => {
959                            problems.push(format!("section '{name}': malformed hash '{hex}'"))
960                        }
961                    }
962                }
963            }
964        }
965
966        for t in &self.tensors {
967            let actual = hash64(self.entry_bytes(t));
968            if actual != t.hash {
969                problems.push(format!(
970                    "tensor '{}': hash mismatch (stored {:016x}, actual {:016x})",
971                    t.name, t.hash, actual
972                ));
973            }
974        }
975        problems
976    }
977
978    /// Approximate active weight bytes under a mask, from real tensor
979    /// sizes in the directory (not from a formula).
980    pub fn compute_active_size(&self, mask: &TaskMask) -> u64 {
981        let arch = &self.header.arch;
982        let mut total = 0u64;
983        for li in 0..arch.num_layers {
984            if !mask.layer_alive(li) {
985                continue;
986            }
987            let ffn_frac = mask.ffn_active_count(li) as f64 / arch.intermediate_size.max(1) as f64;
988            let head_frac =
989                mask.active_head_count(li) as f64 / arch.num_attention_heads.max(1) as f64;
990            for t in self.layer_tensors(li) {
991                let frac = if t.name.contains(".mlp.") {
992                    ffn_frac
993                } else if t.name.contains(".self_attn.") {
994                    head_frac
995                } else {
996                    1.0
997                };
998                total += (t.nbytes as f64 * frac) as u64;
999            }
1000        }
1001        total
1002    }
1003
1004    // ───────────────────────── writer ─────────────────────────
1005
1006    /// Write a CMF v2 file. Offsets, alignment, hashes and the sparse
1007    /// index are computed here — the caller supplies content only.
1008    pub fn write(
1009        path: impl AsRef<Path>,
1010        header: &CmfHeader,
1011        tensors: &[TensorSpec],
1012        masks: Option<&MaskCatalog>,
1013        vocab: Option<&[u8]>,
1014    ) -> Result<(), CmfError> {
1015        let refs: Vec<TensorSpecRef> = tensors
1016            .iter()
1017            .map(|t| TensorSpecRef {
1018                name: t.name.clone(),
1019                dtype: t.dtype,
1020                shape: t.shape.clone(),
1021                data: &t.data,
1022            })
1023            .collect();
1024        Self::write_ref(path, header, &refs, masks, vocab)
1025    }
1026
1027    /// `write` with BORROWED tensor payloads — repack tools slice the
1028    /// source file's mmap directly, so a 19 GB container rewrites without
1029    /// materializing its tensors in RAM (the OS streams pages through).
1030    pub fn write_ref(
1031        path: impl AsRef<Path>,
1032        header: &CmfHeader,
1033        tensors: &[TensorSpecRef],
1034        masks: Option<&MaskCatalog>,
1035        vocab: Option<&[u8]>,
1036    ) -> Result<(), CmfError> {
1037        let path = path.as_ref();
1038
1039        // Directory + data layout.
1040        let mut entries = Vec::with_capacity(tensors.len());
1041        let mut data_cursor = 0u64;
1042        for t in tensors {
1043            if t.shape.len() > DIR_MAX_NDIM {
1044                return Err(CmfError::Parse(format!(
1045                    "tensor '{}': ndim {} > 6",
1046                    t.name,
1047                    t.shape.len()
1048                )));
1049            }
1050            if let Some(expect) = expected_nbytes(t.dtype, &t.shape) {
1051                if expect != t.data.len() {
1052                    return Err(CmfError::Bounds(format!(
1053                        "tensor '{}': data {} bytes != expected {} for {:?}{:?}",
1054                        t.name,
1055                        t.data.len(),
1056                        expect,
1057                        t.dtype,
1058                        t.shape
1059                    )));
1060                }
1061            }
1062            let align = if t.data.len() as u64 >= LARGE_TENSOR_MIN {
1063                LARGE_TENSOR_ALIGN
1064            } else {
1065                TENSOR_ALIGNMENT
1066            };
1067            data_cursor = align_to(data_cursor, align);
1068            entries.push(TensorEntry {
1069                name: t.name.clone(),
1070                dtype: t.dtype,
1071                shape: t.shape.clone(),
1072                off: data_cursor,
1073                nbytes: t.data.len() as u64,
1074                shard: 0,
1075                hash: hash64(&t.data),
1076            });
1077            data_cursor += t.data.len() as u64;
1078        }
1079        let data_len = data_cursor;
1080
1081        let dir_bytes = Self::encode_directory(&entries);
1082
1083        let masks_bytes = match masks {
1084            Some(catalog) if !catalog.masks.is_empty() => {
1085                Some(encode_masks_section(catalog, &header.arch).map_err(CmfError::Parse)?)
1086            }
1087            _ => None,
1088        };
1089        let index_bytes = match masks {
1090            Some(catalog) if !catalog.masks.is_empty() => {
1091                let idx = build_sparse_index(catalog, &header.arch);
1092                Some(encode_sparse_index(&idx))
1093            }
1094            _ => None,
1095        };
1096
1097        // Section hashes go INTO the header (so the envelope's header
1098        // hash transitively covers them), then the header is serialized.
1099        let hex = |b: Option<&[u8]>| b.map(|b| format!("{:016x}", hash64(b)));
1100        let mut header = header.clone();
1101        if masks_bytes.is_some() || vocab.is_some() || index_bytes.is_some() {
1102            header.section_hashes = Some(SectionHashes {
1103                masks: hex(masks_bytes.as_deref()),
1104                vocab: hex(vocab),
1105                index: hex(index_bytes.as_deref()),
1106            });
1107        }
1108        let header_json =
1109            serde_json::to_vec(&header).map_err(|e| CmfError::Parse(format!("header: {e}")))?;
1110
1111        let mut required_features = features::TENSOR_DIR;
1112        if masks_bytes.is_some() {
1113            required_features |= features::BINARY_MASKS;
1114        }
1115        if entries
1116            .iter()
1117            .any(|t| matches!(t.dtype, TensorDtype::Q8_2f | TensorDtype::Vbit))
1118        {
1119            required_features |= features::QUANT_2F;
1120        }
1121
1122        // Section offsets.
1123        let header_off = ENVELOPE_LEN as u64;
1124        let dir_off = header_off + header_json.len() as u64;
1125        let data_off = align_to(dir_off + dir_bytes.len() as u64, DATA_ALIGNMENT);
1126        let masks_off = data_off + data_len;
1127        let masks_len = masks_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1128        let vocab_off = masks_off + masks_len;
1129        let vocab_len = vocab.map(|b| b.len() as u64).unwrap_or(0);
1130        let index_off = vocab_off + vocab_len;
1131        let index_len = index_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1132
1133        // Envelope.
1134        let mut env = Vec::with_capacity(ENVELOPE_LEN);
1135        env.extend_from_slice(&CMF_MAGIC);
1136        env.extend_from_slice(&CMF_VERSION.to_le_bytes());
1137        env.extend_from_slice(&0u32.to_le_bytes()); // flags
1138        env.extend_from_slice(&required_features.to_le_bytes());
1139        for (off, len) in [
1140            (header_off, header_json.len() as u64),
1141            (dir_off, dir_bytes.len() as u64),
1142            (data_off, data_len),
1143            (if masks_len > 0 { masks_off } else { 0 }, masks_len),
1144            (if vocab_len > 0 { vocab_off } else { 0 }, vocab_len),
1145            (if index_len > 0 { index_off } else { 0 }, index_len),
1146        ] {
1147            env.extend_from_slice(&off.to_le_bytes());
1148            env.extend_from_slice(&len.to_le_bytes());
1149        }
1150        // Reserved bytes carry header/dir integrity (spec §8.1).
1151        env.extend_from_slice(&hash64(&header_json).to_le_bytes());
1152        env.extend_from_slice(&hash64(&dir_bytes).to_le_bytes());
1153        env.resize(ENVELOPE_LEN, 0);
1154
1155        // Write out.
1156        let mut f = BufWriter::new(File::create(path)?);
1157        f.write_all(&env)?;
1158        f.write_all(&header_json)?;
1159        f.write_all(&dir_bytes)?;
1160        let mut pos = dir_off + dir_bytes.len() as u64;
1161        f.write_all(&zeros((data_off - pos) as usize))?;
1162        pos = data_off;
1163        for (spec, entry) in tensors.iter().zip(&entries) {
1164            let target = data_off + entry.off;
1165            f.write_all(&zeros((target - pos) as usize))?;
1166            f.write_all(&spec.data)?;
1167            pos = target + spec.data.len() as u64;
1168        }
1169        debug_assert_eq!(pos, data_off + data_len);
1170        if let Some(mb) = &masks_bytes {
1171            f.write_all(mb)?;
1172        }
1173        if let Some(vb) = vocab {
1174            f.write_all(vb)?;
1175        }
1176        if let Some(ib) = &index_bytes {
1177            f.write_all(ib)?;
1178        }
1179        f.flush()?;
1180
1181        tracing::info!(
1182            "Wrote CMF v2: {} ({} tensors, {} masks, {:.1} MB)",
1183            path.display(),
1184            entries.len(),
1185            masks.map(|m| m.masks.len()).unwrap_or(0),
1186            std::fs::metadata(path)?.len() as f64 / 1e6
1187        );
1188        Ok(())
1189    }
1190
1191    pub(crate) fn encode_directory(entries: &[TensorEntry]) -> Vec<u8> {
1192        let mut pool = Vec::new();
1193        let mut name_offs = Vec::with_capacity(entries.len());
1194        for e in entries {
1195            name_offs.push((pool.len() as u32, e.name.len() as u16));
1196            pool.extend_from_slice(e.name.as_bytes());
1197        }
1198        let pool_off = 16 + entries.len() * DIR_RECORD_LEN;
1199
1200        let mut out = Vec::with_capacity(pool_off + pool.len());
1201        out.extend_from_slice(&(entries.len() as u64).to_le_bytes());
1202        out.extend_from_slice(&(pool_off as u64).to_le_bytes());
1203        for (e, (noff, nlen)) in entries.iter().zip(&name_offs) {
1204            out.extend_from_slice(&noff.to_le_bytes());
1205            out.extend_from_slice(&nlen.to_le_bytes());
1206            out.push(e.dtype.id());
1207            out.push(e.shape.len() as u8);
1208            for d in 0..DIR_MAX_NDIM {
1209                out.extend_from_slice(&(e.shape.get(d).copied().unwrap_or(0) as u32).to_le_bytes());
1210            }
1211            out.extend_from_slice(&e.off.to_le_bytes());
1212            out.extend_from_slice(&e.nbytes.to_le_bytes());
1213            out.extend_from_slice(&e.hash.to_le_bytes());
1214        }
1215        out.extend_from_slice(&pool);
1216        out
1217    }
1218}
1219
1220fn align_to(x: u64, a: u64) -> u64 {
1221    x.div_ceil(a) * a
1222}
1223
1224fn zeros(n: usize) -> Vec<u8> {
1225    vec![0u8; n]
1226}
1227
1228// ─────────────────── one-pass streaming writer (§8.4) ───────────────────
1229
1230/// Writes a CMF file in a single pass, payloads first.
1231///
1232/// [`CmfModel::write_ref`] needs every payload addressable at once, so a
1233/// converter has to hold the whole encoded model — RAM, or a spill file it
1234/// then copies into the output. For a 300B-class MoE that is ~120 GB written
1235/// twice. This writer instead reserves a gap at the head of the file, appends
1236/// each payload the moment it is encoded, and patches the envelope, header and
1237/// directory into that gap at the end. The bytes are written once and peak
1238/// disk cost is the finished file.
1239///
1240/// The gap is the one thing that can go wrong: the directory is not sized
1241/// until the last tensor arrives. [`CmfStreamWriter::finish`] therefore
1242/// refuses loudly if the head does not fit rather than truncating it, and
1243/// [`CmfStreamWriter::head_reserve_for`] gives callers a safe estimate.
1244/// What a resumed writer recovers from its manifest: the tensors already on
1245/// disk and any milestones the producer noted.
1246pub struct ResumeState {
1247    pub names: Vec<String>,
1248    pub marks: Vec<String>,
1249}
1250
1251pub struct CmfStreamWriter {
1252    file: BufWriter<File>,
1253    path: PathBuf,
1254    /// Absolute offset of the weight blob — also the size of the reserved gap.
1255    data_off: u64,
1256    /// Write cursor, relative to `data_off`.
1257    cursor: u64,
1258    entries: Vec<TensorEntry>,
1259    /// Append-only sidecar describing every payload already on disk. A Colab
1260    /// box can vanish mid-conversion; with this the finished payloads can be
1261    /// turned into a valid file instead of re-encoding for hours.
1262    manifest: Option<BufWriter<File>>,
1263}
1264
1265impl CmfStreamWriter {
1266    /// A gap that comfortably holds the head for `n_tensors` whose names run
1267    /// to `avg_name` bytes: the directory's fixed records, the name pool, the
1268    /// envelope, and a header JSON with room for arch metadata — then doubled,
1269    /// because being wrong here costs a whole re-run.
1270    pub fn head_reserve_for(n_tensors: usize, avg_name: usize) -> u64 {
1271        let dir = 16 + n_tensors * (DIR_RECORD_LEN + 6 + avg_name);
1272        let head = ENVELOPE_LEN + dir + (1 << 20);
1273        // Tripled, on top of a megabyte of slack that is already ~20x a
1274        // small model's directory. The asymmetry is deliberate: an
1275        // over-estimate costs zeros at the head of the file, an
1276        // under-estimate costs the entire conversion that produced it.
1277        align_to(3 * head as u64, DATA_ALIGNMENT).max(1 << 20)
1278    }
1279
1280    /// `gap` bytes are reserved for envelope + header + directory.
1281    pub fn new(path: impl AsRef<Path>, gap: u64) -> Result<Self, CmfError> {
1282        let path = path.as_ref().to_path_buf();
1283        let data_off = align_to(gap.max(ENVELOPE_LEN as u64 + 1), DATA_ALIGNMENT);
1284        let mut file = BufWriter::new(File::create(&path)?);
1285        file.write_all(&zeros(data_off as usize))?;
1286        Ok(Self {
1287            file,
1288            path,
1289            data_off,
1290            cursor: 0,
1291            entries: Vec::new(),
1292            manifest: None,
1293        })
1294    }
1295
1296    /// Append one tensor. The payload is consumed here, so the caller can drop
1297    /// it immediately — that is the entire point of this writer.
1298    pub fn push(
1299        &mut self,
1300        name: &str,
1301        dtype: TensorDtype,
1302        shape: &[usize],
1303        data: &[u8],
1304    ) -> Result<(), CmfError> {
1305        if shape.len() > DIR_MAX_NDIM {
1306            return Err(CmfError::Parse(format!(
1307                "tensor '{}': ndim {} > {}",
1308                name,
1309                shape.len(),
1310                DIR_MAX_NDIM
1311            )));
1312        }
1313        if let Some(expect) = expected_nbytes(dtype, shape) {
1314            if expect != data.len() {
1315                return Err(CmfError::Bounds(format!(
1316                    "tensor '{}': data {} bytes != expected {} for {:?}{:?}",
1317                    name,
1318                    data.len(),
1319                    expect,
1320                    dtype,
1321                    shape
1322                )));
1323            }
1324        }
1325        let align = if data.len() as u64 >= LARGE_TENSOR_MIN {
1326            LARGE_TENSOR_ALIGN
1327        } else {
1328            TENSOR_ALIGNMENT
1329        };
1330        let off = align_to(self.cursor, align);
1331        self.file.write_all(&zeros((off - self.cursor) as usize))?;
1332        self.file.write_all(data)?;
1333        self.entries.push(TensorEntry {
1334            name: name.to_string(),
1335            dtype,
1336            shape: shape.to_vec(),
1337            off,
1338            nbytes: data.len() as u64,
1339            shard: 0,
1340            hash: hash64(data),
1341        });
1342        self.cursor = off + data.len() as u64;
1343        if let Some(m) = self.manifest.as_mut() {
1344            let e = self.entries.last().unwrap();
1345            writeln!(
1346                m,
1347                "{{\"name\":{},\"dtype\":{},\"shape\":{:?},\"off\":{},\"nbytes\":{},\"hash\":{}}}",
1348                serde_json::to_string(&e.name).unwrap_or_else(|_| "\"?\"".into()),
1349                dtype.id(),
1350                e.shape,
1351                e.off,
1352                e.nbytes,
1353                e.hash
1354            )?;
1355            m.flush()?;
1356        }
1357        Ok(())
1358    }
1359
1360    /// Start recording a sidecar manifest at `path`. One JSON line per
1361    /// tensor, flushed as it goes, plus a first line pinning the gap size.
1362    pub fn with_manifest(mut self, path: impl AsRef<Path>) -> Result<Self, CmfError> {
1363        let mut f = BufWriter::new(File::create(path)?);
1364        writeln!(f, "{{\"data_off\":{}}}", self.data_off)?;
1365        f.flush()?;
1366        self.manifest = Some(f);
1367        Ok(self)
1368    }
1369
1370    /// Keep recording into an existing manifest — for a writer from
1371    /// [`CmfStreamWriter::resume`], whose earlier lines must survive.
1372    pub fn appending_manifest(mut self, path: impl AsRef<Path>) -> Result<Self, CmfError> {
1373        self.manifest = Some(BufWriter::new(
1374            std::fs::OpenOptions::new().append(true).open(path)?,
1375        ));
1376        Ok(self)
1377    }
1378
1379    /// Rebuild a writer over an output file whose payloads are already on
1380    /// disk, from the manifest that recorded them. The file is reopened for
1381    /// writing without truncation and the cursor is placed after the last
1382    /// recorded tensor, so `finish` can complete a conversion that died.
1383    pub fn resume(
1384        path: impl AsRef<Path>,
1385        manifest: impl AsRef<Path>,
1386    ) -> Result<(Self, ResumeState), CmfError> {
1387        let path = path.as_ref().to_path_buf();
1388        let text = std::fs::read_to_string(manifest.as_ref())?;
1389        let mut lines = text.lines();
1390        let first = lines
1391            .next()
1392            .ok_or_else(|| CmfError::Parse("manifest is empty".into()))?;
1393        let head: serde_json::Value = serde_json::from_str(first)
1394            .map_err(|e| CmfError::Parse(format!("manifest head: {e}")))?;
1395        let data_off = head["data_off"]
1396            .as_u64()
1397            .ok_or_else(|| CmfError::Parse("manifest head has no data_off".into()))?;
1398
1399        let mut entries = Vec::new();
1400        let mut names = Vec::new();
1401        let mut marks = Vec::new();
1402        let (mut safe_upto, mut safe_entries) = (0u64, 0usize);
1403        for (i, line) in lines.enumerate() {
1404            if line.trim().is_empty() {
1405                continue;
1406            }
1407            // A truncated last line is expected if the process was killed
1408            // mid-write; it is dropped, not an error.
1409            let Ok(v) = serde_json::from_str::<serde_json::Value>(line) else {
1410                tracing::warn!("manifest line {} is truncated — ignoring it", i + 2);
1411                break;
1412            };
1413            if let Some(mark) = v["mark"].as_str() {
1414                marks.push(mark.to_string());
1415                // Everything up to here is durable; anything the manifest
1416                // records after the LAST mark belongs to a shard that was
1417                // interrupted and will be redone, so it must not be kept —
1418                // otherwise the redo appends those tensors a second time.
1419                safe_upto = v["at"].as_u64().unwrap_or(0);
1420                safe_entries = entries.len();
1421                continue;
1422            }
1423            let dtype = TensorDtype::from_id(v["dtype"].as_u64().unwrap_or(0) as u8)
1424                .ok_or_else(|| CmfError::Parse(format!("manifest line {}: dtype", i + 2)))?;
1425            let name = v["name"].as_str().unwrap_or_default().to_string();
1426            names.push(name.clone());
1427            entries.push(TensorEntry {
1428                name,
1429                dtype,
1430                shape: v["shape"]
1431                    .as_array()
1432                    .map(|a| a.iter().filter_map(|x| x.as_u64()).map(|x| x as usize).collect())
1433                    .unwrap_or_default(),
1434                off: v["off"].as_u64().unwrap_or(0),
1435                nbytes: v["nbytes"].as_u64().unwrap_or(0),
1436                shard: 0,
1437                hash: v["hash"].as_u64().unwrap_or(0),
1438            });
1439        }
1440        entries.truncate(safe_entries);
1441        names.truncate(safe_entries);
1442        let cursor = safe_upto;
1443        debug_assert_eq!(
1444            entries.last().map(|e| e.off + e.nbytes).unwrap_or(0),
1445            cursor,
1446            "the last mark disagrees with the entries before it"
1447        );
1448        let on_disk = std::fs::metadata(&path)?.len();
1449        if on_disk < data_off + cursor {
1450            return Err(CmfError::Bounds(format!(
1451                "{} is {on_disk} bytes but its last checkpoint claims {} — \
1452                 the file is shorter than its own record",
1453                path.display(),
1454                data_off + cursor
1455            )));
1456        }
1457        let mut file = std::fs::OpenOptions::new().write(true).open(&path)?;
1458        file.seek(SeekFrom::Start(data_off + cursor))?;
1459        Ok((
1460            Self {
1461                file: BufWriter::new(file),
1462                path,
1463                data_off,
1464                cursor,
1465                entries,
1466                manifest: None,
1467            },
1468            ResumeState { names, marks },
1469        ))
1470    }
1471
1472    /// Note a milestone in the manifest — a source shard fully consumed, say.
1473    /// Resume reads these back, which is what lets a restart skip work whose
1474    /// payloads are already in the file rather than only skipping tensors it
1475    /// happens to recognise by name.
1476    pub fn mark(&mut self, note: &str) -> Result<(), CmfError> {
1477        // The payloads must be on disk BEFORE the mark claims they are.
1478        // Without this the manifest runs ahead of a buffered writer, and a
1479        // kill in between leaves a record of bytes that were never written.
1480        self.file.flush()?;
1481        if let Some(m) = self.manifest.as_mut() {
1482            writeln!(
1483                m,
1484                "{{\"mark\":{},\"at\":{}}}",
1485                serde_json::to_string(note).unwrap_or_default(),
1486                self.cursor
1487            )?;
1488            m.flush()?;
1489        }
1490        Ok(())
1491    }
1492
1493    pub fn tensor_count(&self) -> usize {
1494        self.entries.len()
1495    }
1496
1497    /// Bytes of weight blob written so far.
1498    pub fn data_len(&self) -> u64 {
1499        self.cursor
1500    }
1501
1502    /// Write the trailing sections, then patch the head into the reserved gap.
1503    pub fn finish(
1504        mut self,
1505        header: &CmfHeader,
1506        masks: Option<&MaskCatalog>,
1507        vocab: Option<&[u8]>,
1508    ) -> Result<(), CmfError> {
1509        let data_len = self.cursor;
1510
1511        let masks_bytes = match masks {
1512            Some(catalog) if !catalog.masks.is_empty() => {
1513                Some(encode_masks_section(catalog, &header.arch).map_err(CmfError::Parse)?)
1514            }
1515            _ => None,
1516        };
1517        let index_bytes = match masks {
1518            Some(catalog) if !catalog.masks.is_empty() => {
1519                Some(encode_sparse_index(&build_sparse_index(catalog, &header.arch)))
1520            }
1521            _ => None,
1522        };
1523        if let Some(mb) = &masks_bytes {
1524            self.file.write_all(mb)?;
1525        }
1526        if let Some(vb) = vocab {
1527            self.file.write_all(vb)?;
1528        }
1529        if let Some(ib) = &index_bytes {
1530            self.file.write_all(ib)?;
1531        }
1532        self.file.flush()?;
1533
1534        let dir_bytes = CmfModel::encode_directory(&self.entries);
1535
1536        let hex = |b: Option<&[u8]>| b.map(|b| format!("{:016x}", hash64(b)));
1537        let mut header = header.clone();
1538        if masks_bytes.is_some() || vocab.is_some() || index_bytes.is_some() {
1539            header.section_hashes = Some(SectionHashes {
1540                masks: hex(masks_bytes.as_deref()),
1541                vocab: hex(vocab),
1542                index: hex(index_bytes.as_deref()),
1543            });
1544        }
1545        let header_json =
1546            serde_json::to_vec(&header).map_err(|e| CmfError::Parse(format!("header: {e}")))?;
1547
1548        let mut required_features = features::TENSOR_DIR;
1549        if masks_bytes.is_some() {
1550            required_features |= features::BINARY_MASKS;
1551        }
1552        if self
1553            .entries
1554            .iter()
1555            .any(|t| matches!(t.dtype, TensorDtype::Q8_2f | TensorDtype::Vbit))
1556        {
1557            required_features |= features::QUANT_2F;
1558        }
1559
1560        let header_off = ENVELOPE_LEN as u64;
1561        let dir_off = header_off + header_json.len() as u64;
1562        let head_len = dir_off + dir_bytes.len() as u64;
1563        if head_len > self.data_off {
1564            return Err(CmfError::Parse(format!(
1565                "streamed head is {head_len} bytes but only {} were reserved — \
1566                 the payloads are already on disk at a fixed offset, so this \
1567                 file cannot be salvaged; re-run with a larger reserve",
1568                self.data_off
1569            )));
1570        }
1571        let data_off = self.data_off;
1572        let masks_off = data_off + data_len;
1573        let masks_len = masks_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1574        let vocab_off = masks_off + masks_len;
1575        let vocab_len = vocab.map(|b| b.len() as u64).unwrap_or(0);
1576        let index_off = vocab_off + vocab_len;
1577        let index_len = index_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1578
1579        let mut env = Vec::with_capacity(ENVELOPE_LEN);
1580        env.extend_from_slice(&CMF_MAGIC);
1581        env.extend_from_slice(&CMF_VERSION.to_le_bytes());
1582        env.extend_from_slice(&0u32.to_le_bytes());
1583        env.extend_from_slice(&required_features.to_le_bytes());
1584        for (off, len) in [
1585            (header_off, header_json.len() as u64),
1586            (dir_off, dir_bytes.len() as u64),
1587            (data_off, data_len),
1588            (if masks_len > 0 { masks_off } else { 0 }, masks_len),
1589            (if vocab_len > 0 { vocab_off } else { 0 }, vocab_len),
1590            (if index_len > 0 { index_off } else { 0 }, index_len),
1591        ] {
1592            env.extend_from_slice(&off.to_le_bytes());
1593            env.extend_from_slice(&len.to_le_bytes());
1594        }
1595        env.extend_from_slice(&hash64(&header_json).to_le_bytes());
1596        env.extend_from_slice(&hash64(&dir_bytes).to_le_bytes());
1597        env.resize(ENVELOPE_LEN, 0);
1598
1599        let mut f = self
1600            .file
1601            .into_inner()
1602            .map_err(|e| CmfError::Io(e.into_error()))?;
1603        f.seek(SeekFrom::Start(0))?;
1604        f.write_all(&env)?;
1605        f.write_all(&header_json)?;
1606        f.write_all(&dir_bytes)?;
1607        f.flush()?;
1608
1609        tracing::info!(
1610            "Wrote CMF v2 (streamed): {} ({} tensors, {:.1} MB)",
1611            self.path.display(),
1612            self.entries.len(),
1613            (data_off + data_len + masks_len + vocab_len + index_len) as f64 / 1e6
1614        );
1615        Ok(())
1616    }
1617}
1618
1619// ───────────────────── sparse index (§7 of the spec) ─────────────────────
1620
1621/// Build the sparse index from mask bitfields: a 32-neuron FFN group is
1622/// active if it contains at least one active bit.
1623pub fn build_sparse_index(catalog: &MaskCatalog, arch: &ModelArch) -> Vec<SparseIndexEntry> {
1624    let mut out = Vec::new();
1625    for m in &catalog.masks {
1626        for li in 0..arch.num_layers {
1627            if !m.layer_alive(li) {
1628                continue;
1629            }
1630            let mut groups = Vec::new();
1631            if let Some(bits) = m.ffn_masks.get(li) {
1632                let n_groups = arch.intermediate_size.div_ceil(32);
1633                for g in 0..n_groups {
1634                    // Group g covers bits [g*32, g*32+32) = bytes [g*4, g*4+4).
1635                    let active = bits[g * 4..(g * 4 + 4).min(bits.len())]
1636                        .iter()
1637                        .any(|&b| b != 0);
1638                    if active {
1639                        groups.push(g as u16);
1640                    }
1641                }
1642            }
1643            let mut heads = Vec::new();
1644            if let Some(bits) = m.head_masks.get(li) {
1645                for h in 0..arch.num_attention_heads {
1646                    if bits
1647                        .get(h / 8)
1648                        .map(|b| b & (1 << (h % 8)) != 0)
1649                        .unwrap_or(false)
1650                    {
1651                        heads.push(h as u8);
1652                    }
1653                }
1654            }
1655            out.push(SparseIndexEntry {
1656                task_id: m.task_id,
1657                layer_idx: li,
1658                active_ffn_groups: groups,
1659                active_heads: heads,
1660            });
1661        }
1662    }
1663    out
1664}
1665
1666/// `[u32 n_entries][u32 reserved]` then per entry:
1667/// `[u32 task][u32 layer][u32 n_groups][u32 n_heads][u16×g][u8×h][pad→4]`.
1668pub fn encode_sparse_index(entries: &[SparseIndexEntry]) -> Vec<u8> {
1669    let mut out = Vec::new();
1670    out.extend_from_slice(&(entries.len() as u32).to_le_bytes());
1671    out.extend_from_slice(&0u32.to_le_bytes());
1672    for e in entries {
1673        out.extend_from_slice(&e.task_id.to_le_bytes());
1674        out.extend_from_slice(&(e.layer_idx as u32).to_le_bytes());
1675        out.extend_from_slice(&(e.active_ffn_groups.len() as u32).to_le_bytes());
1676        out.extend_from_slice(&(e.active_heads.len() as u32).to_le_bytes());
1677        for g in &e.active_ffn_groups {
1678            out.extend_from_slice(&g.to_le_bytes());
1679        }
1680        out.extend_from_slice(&e.active_heads);
1681        while out.len() % 4 != 0 {
1682            out.push(0);
1683        }
1684    }
1685    out
1686}
1687
1688pub fn decode_sparse_index(bytes: &[u8]) -> Result<Vec<SparseIndexEntry>, CmfError> {
1689    let err = |msg: &str| CmfError::Parse(format!("sparse index: {msg}"));
1690    if bytes.len() < 8 {
1691        return Err(err("too short"));
1692    }
1693    let n = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
1694    let mut pos = 8usize;
1695    let mut out = Vec::with_capacity(n);
1696    for _ in 0..n {
1697        if pos + 16 > bytes.len() {
1698            return Err(err("entry header out of bounds"));
1699        }
1700        let task_id = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap());
1701        let layer_idx = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()) as usize;
1702        let n_groups = u32::from_le_bytes(bytes[pos + 8..pos + 12].try_into().unwrap()) as usize;
1703        let n_heads = u32::from_le_bytes(bytes[pos + 12..pos + 16].try_into().unwrap()) as usize;
1704        pos += 16;
1705        if pos + n_groups * 2 + n_heads > bytes.len() {
1706            return Err(err("entry data out of bounds"));
1707        }
1708        let mut groups = Vec::with_capacity(n_groups);
1709        for g in 0..n_groups {
1710            groups.push(u16::from_le_bytes(
1711                bytes[pos + g * 2..pos + g * 2 + 2].try_into().unwrap(),
1712            ));
1713        }
1714        pos += n_groups * 2;
1715        let heads = bytes[pos..pos + n_heads].to_vec();
1716        pos += n_heads;
1717        pos = pos.div_ceil(4) * 4;
1718        out.push(SparseIndexEntry {
1719            task_id,
1720            layer_idx,
1721            active_ffn_groups: groups,
1722            active_heads: heads,
1723        });
1724    }
1725    Ok(out)
1726}
1727
1728/// Errors from CMF operations. Every failure mode is loud.
1729#[derive(Debug, thiserror::Error)]
1730pub enum CmfError {
1731    #[error("File not found: {0}")]
1732    FileNotFound(String),
1733    #[error("Invalid CMF magic bytes")]
1734    InvalidMagic,
1735    #[error("Unsupported CMF version: {0}")]
1736    UnsupportedVersion(u32),
1737    #[error("File requires unsupported features (bits {0:#x})")]
1738    UnsupportedFeature(u32),
1739    #[error("Unknown tensor dtype id: {0}")]
1740    UnknownDtype(u8),
1741    #[error("Tensor not found: {0}")]
1742    MissingTensor(String),
1743    #[error("Bounds error: {0}")]
1744    Bounds(String),
1745    #[error("IO error: {0}")]
1746    Io(#[from] io::Error),
1747    #[error("Parse error: {0}")]
1748    Parse(String),
1749}