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