1use 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;
31pub const DATA_ALIGNMENT: u64 = 4096;
33pub const TENSOR_ALIGNMENT: u64 = 64;
35pub const LARGE_TENSOR_MIN: u64 = 16 * 1024;
43pub const LARGE_TENSOR_ALIGN: u64 = 4096;
48pub const DIR_RECORD_LEN: usize = 56;
50pub const DIR_MAX_NDIM: usize = 6;
51
52pub 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 pub const LOOP_MASKS: u32 = 1 << 5;
66 pub const SKILL_FILE: u32 = 1 << 6;
72 pub const PRISM_HADAMARD: u32 = 1 << 7;
74 pub const PRISM_AFFINE: u32 = 1 << 8;
78
79 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#[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 pub quant_type: QuantType,
137 #[serde(default, skip_serializing_if = "Option::is_none")]
138 pub provenance: Option<serde_json::Value>,
139 #[serde(default, skip_serializing_if = "Option::is_none")]
142 pub tokenizer_config: Option<TokenizerBundle>,
143 #[serde(default, skip_serializing_if = "Option::is_none")]
147 pub section_hashes: Option<SectionHashes>,
148 #[serde(default, skip_serializing_if = "Vec::is_empty")]
152 pub skills: Vec<SkillRecord>,
153 #[serde(default, skip_serializing_if = "Option::is_none")]
157 pub shard: Option<ShardInfo>,
158 #[serde(default, skip_serializing_if = "Option::is_none")]
164 pub calibration: Option<Calibration>,
165 #[serde(default, skip_serializing_if = "Option::is_none")]
170 pub routing: Option<RoutingCalibration>,
171}
172
173#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct RoutingCalibration {
180 pub temperature: f32,
181 pub novelty_theta: f32,
182 pub samples: usize,
184 pub target_fpr: f32,
186}
187
188#[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#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct ShardInfo {
203 pub no: usize,
204 pub count: usize,
205}
206
207#[derive(Debug, Clone, Serialize, Deserialize)]
210pub struct SelectionDescriptor {
211 pub metric: String,
213 pub phi_layer: usize,
215 pub mean: String,
217 pub basis: String,
219 pub rank: usize,
220 #[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 #[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#[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 #[serde(default)]
243 pub layers: Vec<usize>,
244 #[serde(default, skip_serializing_if = "Option::is_none")]
247 pub selection: Option<SelectionDescriptor>,
248 #[serde(default, skip_serializing_if = "Option::is_none")]
250 pub input_mask_task: Option<String>,
251 #[serde(default, skip_serializing_if = "Option::is_none")]
253 pub quality: Option<serde_json::Value>,
254
255 #[serde(default, skip_serializing_if = "Option::is_none")]
260 pub base_dir_hash: Option<String>,
261 #[serde(default, skip_serializing_if = "Option::is_none")]
264 pub base_arch: Option<String>,
265 #[serde(default, skip_serializing_if = "Option::is_none")]
267 pub task: Option<String>,
268 #[serde(default, skip_serializing_if = "Option::is_none")]
270 pub provenance: Option<serde_json::Value>,
271}
272
273#[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#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
287pub struct TokenizerBundle {
288 #[serde(default, skip_serializing_if = "Option::is_none")]
290 pub chat_template: Option<String>,
291 #[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#[derive(Debug, Clone, PartialEq, Eq)]
306pub struct TensorEntry {
307 pub name: String,
308 pub dtype: TensorDtype,
309 pub shape: Vec<usize>,
310 pub off: u64,
312 pub nbytes: u64,
313 pub shard: usize,
316 pub hash: u64,
318}
319
320impl TensorEntry {
321 pub fn n_elems(&self) -> usize {
322 self.shape.iter().product()
323 }
324}
325
326#[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
335pub struct TensorSpecRef<'a> {
337 pub name: String,
338 pub dtype: TensorDtype,
339 pub shape: Vec<usize>,
340 pub data: &'a [u8],
341}
342
343#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
345pub struct SparseIndexEntry {
346 pub task_id: u32,
347 pub layer_idx: usize,
348 pub active_ffn_groups: Vec<u16>,
350 pub active_heads: Vec<u8>,
352}
353
354#[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 header_hash: u64,
366 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
384pub struct CmfModel {
386 pub path: PathBuf,
387 uid: u64,
393 pub header: CmfHeader,
394 pub required_features: u32,
395 pub tensors: Vec<TensorEntry>,
396 by_name: HashMap<u64, u32>,
403 name_overflow: Vec<u32>,
404 pub masks: MaskCatalog,
405 pub sparse_index: Vec<SparseIndexEntry>,
406 pub vocab: Option<Vec<u8>>,
408 backing: Backing,
409 data_off: u64,
410 envelope: Envelope,
411 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
427static MODEL_UID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
429
430impl CmfModel {
431 pub fn uid(&self) -> u64 {
435 self.uid
436 }
437
438 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 #[cfg(unix)]
460 {
461 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 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 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 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 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 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), None => {
611 by_name.insert(h, i as u32);
612 }
613 }
614 }
615
616 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 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 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 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); };
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); 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 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 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 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 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 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 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 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 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 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 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 let _ = unsafe {
1043 m.unchecked_advise_range(memmap2::UncheckedAdvice::DontNeed, start, alen)
1044 };
1045 if let Some(f) = &file {
1046 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 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 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 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 pub fn verify(&self) -> Vec<String> {
1143 let mut problems = Vec::new();
1144
1145 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 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 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 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 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 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 if header.skills.iter().any(|s| s.base_dir_hash.is_some()) {
1362 required_features |= features::SKILL_FILE;
1363 }
1364
1365 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 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()); 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 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 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
1471fn 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
1490pub 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 data_off: u64,
1518 cursor: u64,
1520 entries: Vec<TensorEntry>,
1521 manifest: Option<BufWriter<File>>,
1525}
1526
1527impl CmfStreamWriter {
1528 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 align_to(3 * head as u64, DATA_ALIGNMENT).max(1 << 20)
1540 }
1541
1542 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 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 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 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 f.get_ref().sync_data()?;
1654 sync_parent_dir(&path)?;
1655 self.manifest = Some(f);
1656 Ok(self)
1657 }
1658
1659 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 file.sync_data()?;
1667 sync_parent_dir(&path)?;
1668 self.manifest = Some(BufWriter::new(file));
1669 Ok(self)
1670 }
1671
1672 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 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 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 pub fn mark(&mut self, note: &str) -> Result<(), CmfError> {
1775 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 sync_parent_dir(&self.path)?;
1793 Ok(())
1794 }
1795
1796 pub fn tensor_count(&self) -> usize {
1797 self.entries.len()
1798 }
1799
1800 pub fn data_len(&self) -> u64 {
1802 self.cursor
1803 }
1804
1805 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
1941pub 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
1961pub 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 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
2013pub 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#[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}