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
61 pub const SUPPORTED: u32 = TENSOR_DIR | BINARY_MASKS | QUANT_2F;
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct CmfHeader {
69 #[serde(default = "default_format")]
70 pub format: String,
71 pub version: u32,
72 pub arch: ModelArch,
73 pub quant_type: QuantType,
75 #[serde(default, skip_serializing_if = "Option::is_none")]
76 pub provenance: Option<serde_json::Value>,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub tokenizer_config: Option<TokenizerBundle>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub section_hashes: Option<SectionHashes>,
86 #[serde(default, skip_serializing_if = "Vec::is_empty")]
90 pub skills: Vec<SkillRecord>,
91 #[serde(default, skip_serializing_if = "Option::is_none")]
95 pub shard: Option<ShardInfo>,
96 #[serde(default, skip_serializing_if = "Option::is_none")]
102 pub calibration: Option<Calibration>,
103}
104
105#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct Calibration {
110 pub temperature: f32,
111 #[serde(default, skip_serializing_if = "Option::is_none")]
112 pub ece_before: Option<f32>,
113 #[serde(default, skip_serializing_if = "Option::is_none")]
114 pub ece_after: Option<f32>,
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct ShardInfo {
120 pub no: usize,
121 pub count: usize,
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
127pub struct SelectionDescriptor {
128 pub metric: String,
130 pub phi_layer: usize,
132 pub mean: String,
134 pub basis: String,
136 pub rank: usize,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct SkillRecord {
142 pub id: String,
143 #[serde(default, skip_serializing_if = "Option::is_none")]
144 pub name: Option<String>,
145 #[serde(default)]
147 pub layers: Vec<usize>,
148 #[serde(default, skip_serializing_if = "Option::is_none")]
151 pub selection: Option<SelectionDescriptor>,
152 #[serde(default, skip_serializing_if = "Option::is_none")]
154 pub input_mask_task: Option<String>,
155 #[serde(default, skip_serializing_if = "Option::is_none")]
157 pub quality: Option<serde_json::Value>,
158}
159
160#[derive(Debug, Clone, Default, Serialize, Deserialize)]
163pub struct SectionHashes {
164 #[serde(default, skip_serializing_if = "Option::is_none")]
165 pub masks: Option<String>,
166 #[serde(default, skip_serializing_if = "Option::is_none")]
167 pub vocab: Option<String>,
168 #[serde(default, skip_serializing_if = "Option::is_none")]
169 pub index: Option<String>,
170}
171
172#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
174pub struct TokenizerBundle {
175 #[serde(default, skip_serializing_if = "Option::is_none")]
177 pub chat_template: Option<String>,
178 #[serde(default)]
180 pub eos_token_ids: Vec<u32>,
181 #[serde(default, skip_serializing_if = "Option::is_none")]
182 pub bos_token_id: Option<u32>,
183 #[serde(default, skip_serializing_if = "Option::is_none")]
184 pub pad_token_id: Option<u32>,
185}
186
187fn default_format() -> String {
188 "cmf".to_string()
189}
190
191#[derive(Debug, Clone, PartialEq, Eq)]
193pub struct TensorEntry {
194 pub name: String,
195 pub dtype: TensorDtype,
196 pub shape: Vec<usize>,
197 pub off: u64,
199 pub nbytes: u64,
200 pub shard: usize,
203 pub hash: u64,
205}
206
207impl TensorEntry {
208 pub fn n_elems(&self) -> usize {
209 self.shape.iter().product()
210 }
211}
212
213#[derive(Debug, Clone)]
215pub struct TensorSpec {
216 pub name: String,
217 pub dtype: TensorDtype,
218 pub shape: Vec<usize>,
219 pub data: Vec<u8>,
220}
221
222pub struct TensorSpecRef<'a> {
224 pub name: String,
225 pub dtype: TensorDtype,
226 pub shape: Vec<usize>,
227 pub data: &'a [u8],
228}
229
230#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
232pub struct SparseIndexEntry {
233 pub task_id: u32,
234 pub layer_idx: usize,
235 pub active_ffn_groups: Vec<u16>,
237 pub active_heads: Vec<u8>,
239}
240
241#[derive(Debug, Clone, Copy, Default)]
243struct Envelope {
244 required_features: u32,
245 header: (u64, u64),
246 dir: (u64, u64),
247 data: (u64, u64),
248 masks: (u64, u64),
249 vocab: (u64, u64),
250 index: (u64, u64),
251 header_hash: u64,
253 dir_hash: u64,
255}
256
257enum Backing {
258 Mmap(memmap2::Mmap),
259 Owned(Vec<u8>),
260}
261
262impl Backing {
263 fn bytes(&self) -> &[u8] {
264 match self {
265 Backing::Mmap(m) => m,
266 Backing::Owned(v) => v,
267 }
268 }
269}
270
271pub struct CmfModel {
273 pub path: PathBuf,
274 uid: u64,
280 pub header: CmfHeader,
281 pub required_features: u32,
282 pub tensors: Vec<TensorEntry>,
283 by_name: HashMap<u64, u32>,
290 name_overflow: Vec<u32>,
291 pub masks: MaskCatalog,
292 pub sparse_index: Vec<SparseIndexEntry>,
293 pub vocab: Option<Vec<u8>>,
295 backing: Backing,
296 data_off: u64,
297 envelope: Envelope,
298 extra_shards: Vec<(Backing, u64)>,
301}
302
303impl std::fmt::Debug for CmfModel {
304 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305 f.debug_struct("CmfModel")
306 .field("path", &self.path)
307 .field("arch", &self.header.arch.arch_name)
308 .field("tensors", &self.tensors.len())
309 .field("masks", &self.masks.masks.len())
310 .finish()
311 }
312}
313
314static MODEL_UID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
316
317impl CmfModel {
318 pub fn uid(&self) -> u64 {
322 self.uid
323 }
324
325 pub fn open(path: impl AsRef<Path>) -> Result<Self, CmfError> {
328 let path = path.as_ref().to_path_buf();
329 if !path.exists() {
330 return Err(CmfError::FileNotFound(path.display().to_string()));
331 }
332 let file = File::open(&path)?;
333 let file_len = file.metadata()?.len();
334
335 let backing = match unsafe { memmap2::MmapOptions::new().map(&file) } {
336 Ok(m) => {
337 #[cfg(unix)]
347 {
348 let evicting = cfg!(target_os = "linux")
362 && std::env::var("CMF_GPU").is_ok_and(|v| v != "0" && v != "off")
363 && std::env::var("CMF_UPLOAD_EVICT")
364 .map(|v| v != "0")
365 .unwrap_or(true);
366 if !evicting
367 && std::env::var("CMF_MMAP_ADVISE")
368 .map(|v| v != "0")
369 .unwrap_or(true)
370 {
371 let _ = m.advise(memmap2::Advice::WillNeed);
372 }
373 if std::env::var("CMF_MLOCK")
374 .map(|v| v == "1")
375 .unwrap_or(false)
376 {
377 if let Err(e) = m.lock() {
378 tracing::warn!(
379 "CMF_MLOCK=1: mlock refused ({e}) — continuing unpinned"
380 );
381 }
382 }
383 }
384 Backing::Mmap(m)
385 }
386 Err(e) => {
387 tracing::warn!("mmap failed ({e}), reading file into memory");
388 Backing::Owned(std::fs::read(&path)?)
389 }
390 };
391
392 let env = Self::parse_envelope(backing.bytes(), file_len)?;
393
394 let bytes = backing.bytes();
395 let section = |off: u64, len: u64| -> &[u8] { &bytes[off as usize..(off + len) as usize] };
396
397 let header: CmfHeader = serde_json::from_slice(section(env.header.0, env.header.1))
399 .map_err(|e| CmfError::Parse(format!("header JSON: {e}")))?;
400
401 let tensors = Self::decode_directory(section(env.dir.0, env.dir.1))?;
403 for t in &tensors {
404 if t.off % TENSOR_ALIGNMENT != 0 {
405 return Err(CmfError::Bounds(format!(
406 "tensor '{}': offset {} not 64-aligned",
407 t.name, t.off
408 )));
409 }
410 let tensor_end = t.off.checked_add(t.nbytes).ok_or_else(|| {
411 CmfError::Bounds(format!("tensor '{}': offset + length overflows", t.name))
412 })?;
413 if tensor_end > env.data.1 {
414 return Err(CmfError::Bounds(format!(
415 "tensor '{}': [{}, {}) exceeds data section ({} bytes)",
416 t.name, t.off, tensor_end, env.data.1
417 )));
418 }
419 t.shape
420 .iter()
421 .try_fold(1usize, |n, &dim| n.checked_mul(dim))
422 .ok_or_else(|| {
423 CmfError::Bounds(format!(
424 "tensor '{}': shape product overflows usize",
425 t.name
426 ))
427 })?;
428 if let Some(expect) = expected_nbytes(t.dtype, &t.shape) {
429 if expect as u64 != t.nbytes {
430 return Err(CmfError::Bounds(format!(
431 "tensor '{}': nbytes {} != expected {} for {:?}{:?}",
432 t.name, t.nbytes, expect, t.dtype, t.shape
433 )));
434 }
435 }
436 if matches!(t.dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
439 let payload = section(env.data.0 + t.off, t.nbytes);
440 crate::quant::validate_payload(t.dtype, &t.shape, payload)
441 .map_err(|e| CmfError::Bounds(format!("tensor '{}': {e}", t.name)))?;
442 }
443 }
444 let mut by_name: HashMap<u64, u32> = HashMap::with_capacity(tensors.len());
448 let mut name_overflow: Vec<u32> = Vec::new();
449 for i in 0..tensors.len() {
450 let h = hash64(tensors[i].name.as_bytes());
451 match by_name.get(&h) {
452 Some(&j) if tensors[j as usize].name == tensors[i].name => {
453 return Err(CmfError::Parse(format!(
454 "duplicate tensor name '{}' in directory",
455 tensors[i].name
456 )));
457 }
458 Some(_) => name_overflow.push(i as u32), None => {
460 by_name.insert(h, i as u32);
461 }
462 }
463 }
464
465 let masks = if env.masks.1 > 0 {
467 decode_masks_section(section(env.masks.0, env.masks.1), &header.arch)
468 .map_err(CmfError::Parse)?
469 } else {
470 MaskCatalog::empty()
471 };
472
473 let vocab = if env.vocab.1 > 0 {
475 Some(section(env.vocab.0, env.vocab.1).to_vec())
476 } else {
477 None
478 };
479
480 let sparse_index = if env.index.1 > 0 {
482 decode_sparse_index(section(env.index.0, env.index.1))?
483 } else {
484 vec![]
485 };
486
487 tracing::info!(
488 "Opened CMF v2: {} | {} tensors | {} masks | vocab {} | {:.1} MB",
489 header.arch.arch_name,
490 tensors.len(),
491 masks.masks.len(),
492 if vocab.is_some() { "embedded" } else { "none" },
493 file_len as f64 / 1e6
494 );
495
496 Ok(Self {
497 uid: MODEL_UID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
498 path,
499 header,
500 required_features: env.required_features,
501 tensors,
502 by_name,
503 name_overflow,
504 masks,
505 sparse_index,
506 vocab,
507 backing,
508 data_off: env.data.0,
509 envelope: env,
510 extra_shards: Vec::new(),
511 })
512 }
513
514 pub fn open_sharded(path: impl AsRef<Path>) -> Result<Self, CmfError> {
518 let path = path.as_ref();
519 let mut first = Self::open(path)?;
520 let Some(info) = first.header.shard.clone() else {
521 return Ok(first); };
523 if info.no != 1 {
524 return Err(CmfError::Parse(format!(
525 "open shard 1, not {} (of {})",
526 info.no, info.count
527 )));
528 }
529 let name = path
530 .file_name()
531 .and_then(|n| n.to_str())
532 .ok_or_else(|| CmfError::Parse("bad shard path".into()))?;
533 let tag1 = format!("-{:05}-of-{:05}.cmf", 1, info.count);
534 if !name.ends_with(&tag1) {
535 return Err(CmfError::Parse(format!(
536 "shard file must end with '{tag1}' (got '{name}')"
537 )));
538 }
539 let stem = &name[..name.len() - tag1.len()];
540 for no in 2..=info.count {
541 let sib = path.with_file_name(format!("{stem}-{:05}-of-{:05}.cmf", no, info.count));
542 let sh = Self::open(&sib)?;
543 match &sh.header.shard {
544 Some(si) if si.no == no && si.count == info.count => {}
545 other => {
546 return Err(CmfError::Parse(format!(
547 "{}: wrong shard coords {other:?}",
548 sib.display()
549 )));
550 }
551 }
552 let shard_idx = first.extra_shards.len() + 1;
553 first.extra_shards.push((sh.backing, sh.envelope.data.0));
554 for mut t in sh.tensors {
555 t.shard = shard_idx;
556 let idx = first.tensors.len() as u32;
557 let h = hash64(t.name.as_bytes());
558 match first.by_name.get(&h) {
559 Some(&j) if first.tensors[j as usize].name == t.name => {
560 return Err(CmfError::Parse(format!(
561 "duplicate tensor name '{}' across shards",
562 t.name
563 )));
564 }
565 Some(_) => first.name_overflow.push(idx),
566 None => {
567 first.by_name.insert(h, idx);
568 }
569 }
570 first.tensors.push(t);
571 }
572 }
573 tracing::info!(
574 "sharded model: {} files, {} tensors total",
575 info.count,
576 first.tensors.len()
577 );
578 Ok(first)
579 }
580
581 fn parse_envelope(bytes: &[u8], file_len: u64) -> Result<Envelope, CmfError> {
582 if bytes.len() < ENVELOPE_LEN {
583 return Err(CmfError::Bounds(format!(
584 "file too small for CMF envelope: {} bytes",
585 bytes.len()
586 )));
587 }
588 if bytes[0..4] != CMF_MAGIC {
589 return Err(CmfError::InvalidMagic);
590 }
591 let u32le = |o: usize| u32::from_le_bytes(bytes[o..o + 4].try_into().unwrap());
592 let u64le = |o: usize| u64::from_le_bytes(bytes[o..o + 8].try_into().unwrap());
593
594 let version = u32le(4);
595 if version != CMF_VERSION {
596 return Err(CmfError::UnsupportedVersion(version));
597 }
598 let _flags = u32le(8); let required_features = u32le(12);
600 let unknown = required_features & !features::SUPPORTED;
601 if unknown != 0 {
602 return Err(CmfError::UnsupportedFeature(unknown));
603 }
604
605 let env = Envelope {
606 required_features,
607 header: (u64le(0x10), u64le(0x18)),
608 dir: (u64le(0x20), u64le(0x28)),
609 data: (u64le(0x30), u64le(0x38)),
610 masks: (u64le(0x40), u64le(0x48)),
611 vocab: (u64le(0x50), u64le(0x58)),
612 index: (u64le(0x60), u64le(0x68)),
613 header_hash: u64le(0x70),
614 dir_hash: u64le(0x78),
615 };
616
617 for (name, (off, len), required) in [
618 ("header", env.header, true),
619 ("dir", env.dir, true),
620 ("data", env.data, false),
621 ("masks", env.masks, false),
622 ("vocab", env.vocab, false),
623 ("index", env.index, false),
624 ] {
625 if required && len == 0 {
626 return Err(CmfError::Bounds(format!("section '{name}' is required")));
627 }
628 if len > 0
629 && off
630 .checked_add(len)
631 .map(|end| end > file_len)
632 .unwrap_or(true)
633 {
634 return Err(CmfError::Bounds(format!(
635 "section '{name}' [{off}, {}) exceeds file ({file_len} bytes)",
636 off.saturating_add(len)
637 )));
638 }
639 if len > 0
640 && (usize::try_from(off).is_err()
641 || usize::try_from(len).is_err()
642 || usize::try_from(off + len).is_err())
643 {
644 return Err(CmfError::Bounds(format!(
645 "section '{name}' cannot be addressed on this platform"
646 )));
647 }
648 }
649 if env.data.1 > 0 && env.data.0 % DATA_ALIGNMENT != 0 {
650 return Err(CmfError::Bounds(format!(
651 "data section offset {} not {}-aligned",
652 env.data.0, DATA_ALIGNMENT
653 )));
654 }
655 Ok(env)
656 }
657
658 fn decode_directory(bytes: &[u8]) -> Result<Vec<TensorEntry>, CmfError> {
659 if bytes.len() < 16 {
660 return Err(CmfError::Parse("tensor directory too short".into()));
661 }
662 let count = u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
663 let pool_off = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
664 let records_len = count
665 .checked_mul(DIR_RECORD_LEN)
666 .ok_or_else(|| CmfError::Parse("tensor directory record count overflows".into()))?;
667 let records_end = 16usize
668 .checked_add(records_len)
669 .ok_or_else(|| CmfError::Parse("tensor directory size overflows".into()))?;
670 if records_end > bytes.len() || pool_off > bytes.len() || pool_off < records_end {
671 return Err(CmfError::Parse(format!(
672 "tensor directory malformed: count={count}, pool_off={pool_off}, len={}",
673 bytes.len()
674 )));
675 }
676 let pool = &bytes[pool_off..];
677
678 let mut out = Vec::with_capacity(count);
679 for i in 0..count {
680 let r = &bytes[16 + i * DIR_RECORD_LEN..16 + (i + 1) * DIR_RECORD_LEN];
681 let name_off = u32::from_le_bytes(r[0..4].try_into().unwrap()) as usize;
682 let name_len = u16::from_le_bytes(r[4..6].try_into().unwrap()) as usize;
683 let dtype_id = r[6];
684 let ndim = r[7] as usize;
685 if ndim > DIR_MAX_NDIM {
686 return Err(CmfError::Parse(format!("tensor #{i}: ndim {ndim} > 6")));
687 }
688 let mut shape = Vec::with_capacity(ndim);
689 for d in 0..ndim {
690 shape.push(
691 u32::from_le_bytes(r[8 + d * 4..12 + d * 4].try_into().unwrap()) as usize,
692 );
693 }
694 let off = u64::from_le_bytes(r[32..40].try_into().unwrap());
695 let nbytes = u64::from_le_bytes(r[40..48].try_into().unwrap());
696 let hash = u64::from_le_bytes(r[48..56].try_into().unwrap());
697
698 let name_end = name_off
699 .checked_add(name_len)
700 .ok_or_else(|| CmfError::Parse(format!("tensor #{i}: name range overflows")))?;
701 if name_end > pool.len() {
702 return Err(CmfError::Parse(format!("tensor #{i}: name out of pool")));
703 }
704 let name = std::str::from_utf8(&pool[name_off..name_end])
705 .map_err(|_| CmfError::Parse(format!("tensor #{i}: name is not UTF-8")))?
706 .to_string();
707 let dtype = TensorDtype::from_id(dtype_id).ok_or(CmfError::UnknownDtype(dtype_id))?;
708
709 out.push(TensorEntry {
710 name,
711 dtype,
712 shape,
713 off,
714 nbytes,
715 shard: 0,
716 hash,
717 });
718 }
719 Ok(out)
720 }
721
722 pub fn arch(&self) -> &ModelArch {
725 &self.header.arch
726 }
727
728 pub fn tensor(&self, name: &str) -> Option<&TensorEntry> {
729 self.tensor_index(name).map(|i| &self.tensors[i])
730 }
731
732 pub fn tensor_index(&self, name: &str) -> Option<usize> {
738 let h = hash64(name.as_bytes());
739 if let Some(&i) = self.by_name.get(&h) {
740 if self.tensors[i as usize].name == name {
741 return Some(i as usize);
742 }
743 }
744 self.name_overflow
745 .iter()
746 .copied()
747 .find(|&i| self.tensors[i as usize].name == name)
748 .map(|i| i as usize)
749 }
750
751 pub fn resolve_tensor(&self, name: &str, skill: Option<&str>) -> Option<&TensorEntry> {
755 if let Some(sid) = skill {
756 if let Some(t) = self.tensor(&format!("skill.{sid}.{name}")) {
757 return Some(t);
758 }
759 }
760 self.tensor(name)
761 }
762
763 pub fn skill_tensors(&self, skill_id: &str) -> impl Iterator<Item = &TensorEntry> {
766 let prefix = format!("skill.{skill_id}.");
767 self.tensors
768 .iter()
769 .filter(move |t| t.name.starts_with(&prefix))
770 }
771
772 pub fn tensor_bytes(&self, name: &str) -> Result<&[u8], CmfError> {
774 let entry = self
775 .tensor(name)
776 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
777 Ok(self.entry_bytes(entry))
778 }
779
780 pub fn primary_bytes(&self) -> &[u8] {
783 self.backing.bytes()
784 }
785
786 pub fn entry_abs_offset(&self, entry: &TensorEntry) -> Option<usize> {
789 (entry.shard == 0).then(|| (self.data_off + entry.off) as usize)
790 }
791
792 pub fn entry_bytes(&self, entry: &TensorEntry) -> &[u8] {
793 let (bytes, data_off) = if entry.shard == 0 {
794 (self.backing.bytes(), self.data_off)
795 } else {
796 let (b, o) = &self.extra_shards[entry.shard - 1];
797 (b.bytes(), *o)
798 };
799 let start = (data_off + entry.off) as usize;
800 &bytes[start..start + entry.nbytes as usize]
801 }
802
803 pub fn evict_ranges(&self, ranges: &[(usize, usize)]) {
812 #[cfg(target_os = "linux")]
813 {
814 let Backing::Mmap(m) = &self.backing else {
815 return;
816 };
817 let page = 4096usize;
818 use std::os::unix::io::AsRawFd;
821 let file = File::open(&self.path).ok();
822 for &(off, len) in ranges {
823 if len == 0 || off.saturating_add(len) > m.len() {
824 continue;
825 }
826 let start = off & !(page - 1);
827 let end = off + len;
828 let alen = end.next_multiple_of(page).min(m.len()) - start;
829 let _ = unsafe {
832 m.unchecked_advise_range(memmap2::UncheckedAdvice::DontNeed, start, alen)
833 };
834 if let Some(f) = &file {
835 unsafe {
837 libc::posix_fadvise(
838 f.as_raw_fd(),
839 start as libc::off_t,
840 alen as libc::off_t,
841 libc::POSIX_FADV_DONTNEED,
842 );
843 }
844 }
845 }
846 }
847 #[cfg(not(target_os = "linux"))]
848 let _ = ranges;
849 }
850
851 pub fn layer_tensors(&self, layer_idx: usize) -> Vec<&TensorEntry> {
853 let prefix = format!("model.layers.{layer_idx}.");
854 self.tensors
855 .iter()
856 .filter(|t| t.name.starts_with(&prefix))
857 .collect()
858 }
859
860 pub fn total_param_count(&self) -> u64 {
862 self.tensors
863 .iter()
864 .filter(|t| t.shape.len() >= 2)
865 .map(|t| t.n_elems() as u64)
866 .sum()
867 }
868
869 pub fn recode_entries_in_place(
880 path: &str,
881 patches: &[(usize, TensorDtype, Vec<u8>)],
882 ) -> Result<(), CmfError> {
883 use std::io::{Read, Seek, SeekFrom, Write};
884 let mut f = std::fs::OpenOptions::new().read(true).write(true).open(path)?;
885 let file_len = f.metadata()?.len();
886 let mut head = vec![0u8; ENVELOPE_LEN];
887 f.read_exact(&mut head)?;
888 let env = Self::parse_envelope(&head, file_len)?;
889
890 let mut dir = vec![0u8; env.dir.1 as usize];
891 f.seek(SeekFrom::Start(env.dir.0))?;
892 f.read_exact(&mut dir)?;
893 let count = u64::from_le_bytes(dir[0..8].try_into().unwrap()) as usize;
894
895 for (i, dtype, data) in patches {
896 if *i >= count {
897 return Err(CmfError::Bounds(format!(
898 "recode: tensor #{i} out of directory ({count} entries)"
899 )));
900 }
901 let rb = 16 + i * DIR_RECORD_LEN;
902 let rec = &mut dir[rb..rb + DIR_RECORD_LEN];
903 let off = u64::from_le_bytes(rec[32..40].try_into().unwrap());
904 let old_n = u64::from_le_bytes(rec[40..48].try_into().unwrap());
905 if data.len() as u64 > old_n {
906 return Err(CmfError::Bounds(format!(
907 "recode: tensor #{i} payload {} > slot {old_n}",
908 data.len()
909 )));
910 }
911 f.seek(SeekFrom::Start(env.data.0 + off))?;
912 f.write_all(data)?;
913 rec[6] = dtype.id();
914 rec[40..48].copy_from_slice(&(data.len() as u64).to_le_bytes());
915 rec[48..56].copy_from_slice(&hash64(data).to_le_bytes());
916 }
917
918 f.seek(SeekFrom::Start(env.dir.0))?;
919 f.write_all(&dir)?;
920 f.seek(SeekFrom::Start(0x78))?;
921 f.write_all(&hash64(&dir).to_le_bytes())?;
922 f.sync_all()?;
923 Ok(())
924 }
925
926 pub fn verify(&self) -> Vec<String> {
929 let mut problems = Vec::new();
930
931 let bytes = self.backing.bytes();
933 let env = &self.envelope;
934 let sect = |(off, len): (u64, u64)| &bytes[off as usize..(off + len) as usize];
935 let check = |name: &str, stored: u64, span: (u64, u64)| -> Option<String> {
936 if stored != 0 && span.1 > 0 {
937 let actual = hash64(sect(span));
938 if actual != stored {
939 return Some(format!(
940 "section '{name}': hash mismatch (stored {stored:016x}, \
941 actual {actual:016x})"
942 ));
943 }
944 }
945 None
946 };
947 problems.extend(check("header", env.header_hash, env.header));
948 problems.extend(check("dir", env.dir_hash, env.dir));
949 if let Some(sh) = &self.header.section_hashes {
950 for (name, hex, span) in [
951 ("masks", &sh.masks, env.masks),
952 ("vocab", &sh.vocab, env.vocab),
953 ("index", &sh.index, env.index),
954 ] {
955 if let Some(hex) = hex {
956 match u64::from_str_radix(hex, 16) {
957 Ok(stored) => problems.extend(check(name, stored, span)),
958 Err(_) => {
959 problems.push(format!("section '{name}': malformed hash '{hex}'"))
960 }
961 }
962 }
963 }
964 }
965
966 for t in &self.tensors {
967 let actual = hash64(self.entry_bytes(t));
968 if actual != t.hash {
969 problems.push(format!(
970 "tensor '{}': hash mismatch (stored {:016x}, actual {:016x})",
971 t.name, t.hash, actual
972 ));
973 }
974 }
975 problems
976 }
977
978 pub fn compute_active_size(&self, mask: &TaskMask) -> u64 {
981 let arch = &self.header.arch;
982 let mut total = 0u64;
983 for li in 0..arch.num_layers {
984 if !mask.layer_alive(li) {
985 continue;
986 }
987 let ffn_frac = mask.ffn_active_count(li) as f64 / arch.intermediate_size.max(1) as f64;
988 let head_frac =
989 mask.active_head_count(li) as f64 / arch.num_attention_heads.max(1) as f64;
990 for t in self.layer_tensors(li) {
991 let frac = if t.name.contains(".mlp.") {
992 ffn_frac
993 } else if t.name.contains(".self_attn.") {
994 head_frac
995 } else {
996 1.0
997 };
998 total += (t.nbytes as f64 * frac) as u64;
999 }
1000 }
1001 total
1002 }
1003
1004 pub fn write(
1009 path: impl AsRef<Path>,
1010 header: &CmfHeader,
1011 tensors: &[TensorSpec],
1012 masks: Option<&MaskCatalog>,
1013 vocab: Option<&[u8]>,
1014 ) -> Result<(), CmfError> {
1015 let refs: Vec<TensorSpecRef> = tensors
1016 .iter()
1017 .map(|t| TensorSpecRef {
1018 name: t.name.clone(),
1019 dtype: t.dtype,
1020 shape: t.shape.clone(),
1021 data: &t.data,
1022 })
1023 .collect();
1024 Self::write_ref(path, header, &refs, masks, vocab)
1025 }
1026
1027 pub fn write_ref(
1031 path: impl AsRef<Path>,
1032 header: &CmfHeader,
1033 tensors: &[TensorSpecRef],
1034 masks: Option<&MaskCatalog>,
1035 vocab: Option<&[u8]>,
1036 ) -> Result<(), CmfError> {
1037 let path = path.as_ref();
1038
1039 let mut entries = Vec::with_capacity(tensors.len());
1041 let mut data_cursor = 0u64;
1042 for t in tensors {
1043 if t.shape.len() > DIR_MAX_NDIM {
1044 return Err(CmfError::Parse(format!(
1045 "tensor '{}': ndim {} > 6",
1046 t.name,
1047 t.shape.len()
1048 )));
1049 }
1050 if let Some(expect) = expected_nbytes(t.dtype, &t.shape) {
1051 if expect != t.data.len() {
1052 return Err(CmfError::Bounds(format!(
1053 "tensor '{}': data {} bytes != expected {} for {:?}{:?}",
1054 t.name,
1055 t.data.len(),
1056 expect,
1057 t.dtype,
1058 t.shape
1059 )));
1060 }
1061 }
1062 let align = if t.data.len() as u64 >= LARGE_TENSOR_MIN {
1063 LARGE_TENSOR_ALIGN
1064 } else {
1065 TENSOR_ALIGNMENT
1066 };
1067 data_cursor = align_to(data_cursor, align);
1068 entries.push(TensorEntry {
1069 name: t.name.clone(),
1070 dtype: t.dtype,
1071 shape: t.shape.clone(),
1072 off: data_cursor,
1073 nbytes: t.data.len() as u64,
1074 shard: 0,
1075 hash: hash64(&t.data),
1076 });
1077 data_cursor += t.data.len() as u64;
1078 }
1079 let data_len = data_cursor;
1080
1081 let dir_bytes = Self::encode_directory(&entries);
1082
1083 let masks_bytes = match masks {
1084 Some(catalog) if !catalog.masks.is_empty() => {
1085 Some(encode_masks_section(catalog, &header.arch).map_err(CmfError::Parse)?)
1086 }
1087 _ => None,
1088 };
1089 let index_bytes = match masks {
1090 Some(catalog) if !catalog.masks.is_empty() => {
1091 let idx = build_sparse_index(catalog, &header.arch);
1092 Some(encode_sparse_index(&idx))
1093 }
1094 _ => None,
1095 };
1096
1097 let hex = |b: Option<&[u8]>| b.map(|b| format!("{:016x}", hash64(b)));
1100 let mut header = header.clone();
1101 if masks_bytes.is_some() || vocab.is_some() || index_bytes.is_some() {
1102 header.section_hashes = Some(SectionHashes {
1103 masks: hex(masks_bytes.as_deref()),
1104 vocab: hex(vocab),
1105 index: hex(index_bytes.as_deref()),
1106 });
1107 }
1108 let header_json =
1109 serde_json::to_vec(&header).map_err(|e| CmfError::Parse(format!("header: {e}")))?;
1110
1111 let mut required_features = features::TENSOR_DIR;
1112 if masks_bytes.is_some() {
1113 required_features |= features::BINARY_MASKS;
1114 }
1115 if entries
1116 .iter()
1117 .any(|t| matches!(t.dtype, TensorDtype::Q8_2f | TensorDtype::Vbit))
1118 {
1119 required_features |= features::QUANT_2F;
1120 }
1121
1122 let header_off = ENVELOPE_LEN as u64;
1124 let dir_off = header_off + header_json.len() as u64;
1125 let data_off = align_to(dir_off + dir_bytes.len() as u64, DATA_ALIGNMENT);
1126 let masks_off = data_off + data_len;
1127 let masks_len = masks_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1128 let vocab_off = masks_off + masks_len;
1129 let vocab_len = vocab.map(|b| b.len() as u64).unwrap_or(0);
1130 let index_off = vocab_off + vocab_len;
1131 let index_len = index_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1132
1133 let mut env = Vec::with_capacity(ENVELOPE_LEN);
1135 env.extend_from_slice(&CMF_MAGIC);
1136 env.extend_from_slice(&CMF_VERSION.to_le_bytes());
1137 env.extend_from_slice(&0u32.to_le_bytes()); env.extend_from_slice(&required_features.to_le_bytes());
1139 for (off, len) in [
1140 (header_off, header_json.len() as u64),
1141 (dir_off, dir_bytes.len() as u64),
1142 (data_off, data_len),
1143 (if masks_len > 0 { masks_off } else { 0 }, masks_len),
1144 (if vocab_len > 0 { vocab_off } else { 0 }, vocab_len),
1145 (if index_len > 0 { index_off } else { 0 }, index_len),
1146 ] {
1147 env.extend_from_slice(&off.to_le_bytes());
1148 env.extend_from_slice(&len.to_le_bytes());
1149 }
1150 env.extend_from_slice(&hash64(&header_json).to_le_bytes());
1152 env.extend_from_slice(&hash64(&dir_bytes).to_le_bytes());
1153 env.resize(ENVELOPE_LEN, 0);
1154
1155 let mut f = BufWriter::new(File::create(path)?);
1157 f.write_all(&env)?;
1158 f.write_all(&header_json)?;
1159 f.write_all(&dir_bytes)?;
1160 let mut pos = dir_off + dir_bytes.len() as u64;
1161 f.write_all(&zeros((data_off - pos) as usize))?;
1162 pos = data_off;
1163 for (spec, entry) in tensors.iter().zip(&entries) {
1164 let target = data_off + entry.off;
1165 f.write_all(&zeros((target - pos) as usize))?;
1166 f.write_all(&spec.data)?;
1167 pos = target + spec.data.len() as u64;
1168 }
1169 debug_assert_eq!(pos, data_off + data_len);
1170 if let Some(mb) = &masks_bytes {
1171 f.write_all(mb)?;
1172 }
1173 if let Some(vb) = vocab {
1174 f.write_all(vb)?;
1175 }
1176 if let Some(ib) = &index_bytes {
1177 f.write_all(ib)?;
1178 }
1179 f.flush()?;
1180
1181 tracing::info!(
1182 "Wrote CMF v2: {} ({} tensors, {} masks, {:.1} MB)",
1183 path.display(),
1184 entries.len(),
1185 masks.map(|m| m.masks.len()).unwrap_or(0),
1186 std::fs::metadata(path)?.len() as f64 / 1e6
1187 );
1188 Ok(())
1189 }
1190
1191 pub(crate) fn encode_directory(entries: &[TensorEntry]) -> Vec<u8> {
1192 let mut pool = Vec::new();
1193 let mut name_offs = Vec::with_capacity(entries.len());
1194 for e in entries {
1195 name_offs.push((pool.len() as u32, e.name.len() as u16));
1196 pool.extend_from_slice(e.name.as_bytes());
1197 }
1198 let pool_off = 16 + entries.len() * DIR_RECORD_LEN;
1199
1200 let mut out = Vec::with_capacity(pool_off + pool.len());
1201 out.extend_from_slice(&(entries.len() as u64).to_le_bytes());
1202 out.extend_from_slice(&(pool_off as u64).to_le_bytes());
1203 for (e, (noff, nlen)) in entries.iter().zip(&name_offs) {
1204 out.extend_from_slice(&noff.to_le_bytes());
1205 out.extend_from_slice(&nlen.to_le_bytes());
1206 out.push(e.dtype.id());
1207 out.push(e.shape.len() as u8);
1208 for d in 0..DIR_MAX_NDIM {
1209 out.extend_from_slice(&(e.shape.get(d).copied().unwrap_or(0) as u32).to_le_bytes());
1210 }
1211 out.extend_from_slice(&e.off.to_le_bytes());
1212 out.extend_from_slice(&e.nbytes.to_le_bytes());
1213 out.extend_from_slice(&e.hash.to_le_bytes());
1214 }
1215 out.extend_from_slice(&pool);
1216 out
1217 }
1218}
1219
1220fn align_to(x: u64, a: u64) -> u64 {
1221 x.div_ceil(a) * a
1222}
1223
1224fn zeros(n: usize) -> Vec<u8> {
1225 vec![0u8; n]
1226}
1227
1228pub struct ResumeState {
1247 pub names: Vec<String>,
1248 pub marks: Vec<String>,
1249}
1250
1251pub struct CmfStreamWriter {
1252 file: BufWriter<File>,
1253 path: PathBuf,
1254 data_off: u64,
1256 cursor: u64,
1258 entries: Vec<TensorEntry>,
1259 manifest: Option<BufWriter<File>>,
1263}
1264
1265impl CmfStreamWriter {
1266 pub fn head_reserve_for(n_tensors: usize, avg_name: usize) -> u64 {
1271 let dir = 16 + n_tensors * (DIR_RECORD_LEN + 6 + avg_name);
1272 let head = ENVELOPE_LEN + dir + (1 << 20);
1273 align_to(3 * head as u64, DATA_ALIGNMENT).max(1 << 20)
1278 }
1279
1280 pub fn new(path: impl AsRef<Path>, gap: u64) -> Result<Self, CmfError> {
1282 let path = path.as_ref().to_path_buf();
1283 let data_off = align_to(gap.max(ENVELOPE_LEN as u64 + 1), DATA_ALIGNMENT);
1284 let mut file = BufWriter::new(File::create(&path)?);
1285 file.write_all(&zeros(data_off as usize))?;
1286 Ok(Self {
1287 file,
1288 path,
1289 data_off,
1290 cursor: 0,
1291 entries: Vec::new(),
1292 manifest: None,
1293 })
1294 }
1295
1296 pub fn push(
1299 &mut self,
1300 name: &str,
1301 dtype: TensorDtype,
1302 shape: &[usize],
1303 data: &[u8],
1304 ) -> Result<(), CmfError> {
1305 if shape.len() > DIR_MAX_NDIM {
1306 return Err(CmfError::Parse(format!(
1307 "tensor '{}': ndim {} > {}",
1308 name,
1309 shape.len(),
1310 DIR_MAX_NDIM
1311 )));
1312 }
1313 if let Some(expect) = expected_nbytes(dtype, shape) {
1314 if expect != data.len() {
1315 return Err(CmfError::Bounds(format!(
1316 "tensor '{}': data {} bytes != expected {} for {:?}{:?}",
1317 name,
1318 data.len(),
1319 expect,
1320 dtype,
1321 shape
1322 )));
1323 }
1324 }
1325 let align = if data.len() as u64 >= LARGE_TENSOR_MIN {
1326 LARGE_TENSOR_ALIGN
1327 } else {
1328 TENSOR_ALIGNMENT
1329 };
1330 let off = align_to(self.cursor, align);
1331 self.file.write_all(&zeros((off - self.cursor) as usize))?;
1332 self.file.write_all(data)?;
1333 self.entries.push(TensorEntry {
1334 name: name.to_string(),
1335 dtype,
1336 shape: shape.to_vec(),
1337 off,
1338 nbytes: data.len() as u64,
1339 shard: 0,
1340 hash: hash64(data),
1341 });
1342 self.cursor = off + data.len() as u64;
1343 if let Some(m) = self.manifest.as_mut() {
1344 let e = self.entries.last().unwrap();
1345 writeln!(
1346 m,
1347 "{{\"name\":{},\"dtype\":{},\"shape\":{:?},\"off\":{},\"nbytes\":{},\"hash\":{}}}",
1348 serde_json::to_string(&e.name).unwrap_or_else(|_| "\"?\"".into()),
1349 dtype.id(),
1350 e.shape,
1351 e.off,
1352 e.nbytes,
1353 e.hash
1354 )?;
1355 m.flush()?;
1356 }
1357 Ok(())
1358 }
1359
1360 pub fn with_manifest(mut self, path: impl AsRef<Path>) -> Result<Self, CmfError> {
1363 let mut f = BufWriter::new(File::create(path)?);
1364 writeln!(f, "{{\"data_off\":{}}}", self.data_off)?;
1365 f.flush()?;
1366 self.manifest = Some(f);
1367 Ok(self)
1368 }
1369
1370 pub fn appending_manifest(mut self, path: impl AsRef<Path>) -> Result<Self, CmfError> {
1373 self.manifest = Some(BufWriter::new(
1374 std::fs::OpenOptions::new().append(true).open(path)?,
1375 ));
1376 Ok(self)
1377 }
1378
1379 pub fn resume(
1384 path: impl AsRef<Path>,
1385 manifest: impl AsRef<Path>,
1386 ) -> Result<(Self, ResumeState), CmfError> {
1387 let path = path.as_ref().to_path_buf();
1388 let text = std::fs::read_to_string(manifest.as_ref())?;
1389 let mut lines = text.lines();
1390 let first = lines
1391 .next()
1392 .ok_or_else(|| CmfError::Parse("manifest is empty".into()))?;
1393 let head: serde_json::Value = serde_json::from_str(first)
1394 .map_err(|e| CmfError::Parse(format!("manifest head: {e}")))?;
1395 let data_off = head["data_off"]
1396 .as_u64()
1397 .ok_or_else(|| CmfError::Parse("manifest head has no data_off".into()))?;
1398
1399 let mut entries = Vec::new();
1400 let mut names = Vec::new();
1401 let mut marks = Vec::new();
1402 let (mut safe_upto, mut safe_entries) = (0u64, 0usize);
1403 for (i, line) in lines.enumerate() {
1404 if line.trim().is_empty() {
1405 continue;
1406 }
1407 let Ok(v) = serde_json::from_str::<serde_json::Value>(line) else {
1410 tracing::warn!("manifest line {} is truncated — ignoring it", i + 2);
1411 break;
1412 };
1413 if let Some(mark) = v["mark"].as_str() {
1414 marks.push(mark.to_string());
1415 safe_upto = v["at"].as_u64().unwrap_or(0);
1420 safe_entries = entries.len();
1421 continue;
1422 }
1423 let dtype = TensorDtype::from_id(v["dtype"].as_u64().unwrap_or(0) as u8)
1424 .ok_or_else(|| CmfError::Parse(format!("manifest line {}: dtype", i + 2)))?;
1425 let name = v["name"].as_str().unwrap_or_default().to_string();
1426 names.push(name.clone());
1427 entries.push(TensorEntry {
1428 name,
1429 dtype,
1430 shape: v["shape"]
1431 .as_array()
1432 .map(|a| a.iter().filter_map(|x| x.as_u64()).map(|x| x as usize).collect())
1433 .unwrap_or_default(),
1434 off: v["off"].as_u64().unwrap_or(0),
1435 nbytes: v["nbytes"].as_u64().unwrap_or(0),
1436 shard: 0,
1437 hash: v["hash"].as_u64().unwrap_or(0),
1438 });
1439 }
1440 entries.truncate(safe_entries);
1441 names.truncate(safe_entries);
1442 let cursor = safe_upto;
1443 debug_assert_eq!(
1444 entries.last().map(|e| e.off + e.nbytes).unwrap_or(0),
1445 cursor,
1446 "the last mark disagrees with the entries before it"
1447 );
1448 let on_disk = std::fs::metadata(&path)?.len();
1449 if on_disk < data_off + cursor {
1450 return Err(CmfError::Bounds(format!(
1451 "{} is {on_disk} bytes but its last checkpoint claims {} — \
1452 the file is shorter than its own record",
1453 path.display(),
1454 data_off + cursor
1455 )));
1456 }
1457 let mut file = std::fs::OpenOptions::new().write(true).open(&path)?;
1458 file.seek(SeekFrom::Start(data_off + cursor))?;
1459 Ok((
1460 Self {
1461 file: BufWriter::new(file),
1462 path,
1463 data_off,
1464 cursor,
1465 entries,
1466 manifest: None,
1467 },
1468 ResumeState { names, marks },
1469 ))
1470 }
1471
1472 pub fn mark(&mut self, note: &str) -> Result<(), CmfError> {
1477 self.file.flush()?;
1481 if let Some(m) = self.manifest.as_mut() {
1482 writeln!(
1483 m,
1484 "{{\"mark\":{},\"at\":{}}}",
1485 serde_json::to_string(note).unwrap_or_default(),
1486 self.cursor
1487 )?;
1488 m.flush()?;
1489 }
1490 Ok(())
1491 }
1492
1493 pub fn tensor_count(&self) -> usize {
1494 self.entries.len()
1495 }
1496
1497 pub fn data_len(&self) -> u64 {
1499 self.cursor
1500 }
1501
1502 pub fn finish(
1504 mut self,
1505 header: &CmfHeader,
1506 masks: Option<&MaskCatalog>,
1507 vocab: Option<&[u8]>,
1508 ) -> Result<(), CmfError> {
1509 let data_len = self.cursor;
1510
1511 let masks_bytes = match masks {
1512 Some(catalog) if !catalog.masks.is_empty() => {
1513 Some(encode_masks_section(catalog, &header.arch).map_err(CmfError::Parse)?)
1514 }
1515 _ => None,
1516 };
1517 let index_bytes = match masks {
1518 Some(catalog) if !catalog.masks.is_empty() => {
1519 Some(encode_sparse_index(&build_sparse_index(catalog, &header.arch)))
1520 }
1521 _ => None,
1522 };
1523 if let Some(mb) = &masks_bytes {
1524 self.file.write_all(mb)?;
1525 }
1526 if let Some(vb) = vocab {
1527 self.file.write_all(vb)?;
1528 }
1529 if let Some(ib) = &index_bytes {
1530 self.file.write_all(ib)?;
1531 }
1532 self.file.flush()?;
1533
1534 let dir_bytes = CmfModel::encode_directory(&self.entries);
1535
1536 let hex = |b: Option<&[u8]>| b.map(|b| format!("{:016x}", hash64(b)));
1537 let mut header = header.clone();
1538 if masks_bytes.is_some() || vocab.is_some() || index_bytes.is_some() {
1539 header.section_hashes = Some(SectionHashes {
1540 masks: hex(masks_bytes.as_deref()),
1541 vocab: hex(vocab),
1542 index: hex(index_bytes.as_deref()),
1543 });
1544 }
1545 let header_json =
1546 serde_json::to_vec(&header).map_err(|e| CmfError::Parse(format!("header: {e}")))?;
1547
1548 let mut required_features = features::TENSOR_DIR;
1549 if masks_bytes.is_some() {
1550 required_features |= features::BINARY_MASKS;
1551 }
1552 if self
1553 .entries
1554 .iter()
1555 .any(|t| matches!(t.dtype, TensorDtype::Q8_2f | TensorDtype::Vbit))
1556 {
1557 required_features |= features::QUANT_2F;
1558 }
1559
1560 let header_off = ENVELOPE_LEN as u64;
1561 let dir_off = header_off + header_json.len() as u64;
1562 let head_len = dir_off + dir_bytes.len() as u64;
1563 if head_len > self.data_off {
1564 return Err(CmfError::Parse(format!(
1565 "streamed head is {head_len} bytes but only {} were reserved — \
1566 the payloads are already on disk at a fixed offset, so this \
1567 file cannot be salvaged; re-run with a larger reserve",
1568 self.data_off
1569 )));
1570 }
1571 let data_off = self.data_off;
1572 let masks_off = data_off + data_len;
1573 let masks_len = masks_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1574 let vocab_off = masks_off + masks_len;
1575 let vocab_len = vocab.map(|b| b.len() as u64).unwrap_or(0);
1576 let index_off = vocab_off + vocab_len;
1577 let index_len = index_bytes.as_ref().map(|b| b.len() as u64).unwrap_or(0);
1578
1579 let mut env = Vec::with_capacity(ENVELOPE_LEN);
1580 env.extend_from_slice(&CMF_MAGIC);
1581 env.extend_from_slice(&CMF_VERSION.to_le_bytes());
1582 env.extend_from_slice(&0u32.to_le_bytes());
1583 env.extend_from_slice(&required_features.to_le_bytes());
1584 for (off, len) in [
1585 (header_off, header_json.len() as u64),
1586 (dir_off, dir_bytes.len() as u64),
1587 (data_off, data_len),
1588 (if masks_len > 0 { masks_off } else { 0 }, masks_len),
1589 (if vocab_len > 0 { vocab_off } else { 0 }, vocab_len),
1590 (if index_len > 0 { index_off } else { 0 }, index_len),
1591 ] {
1592 env.extend_from_slice(&off.to_le_bytes());
1593 env.extend_from_slice(&len.to_le_bytes());
1594 }
1595 env.extend_from_slice(&hash64(&header_json).to_le_bytes());
1596 env.extend_from_slice(&hash64(&dir_bytes).to_le_bytes());
1597 env.resize(ENVELOPE_LEN, 0);
1598
1599 let mut f = self
1600 .file
1601 .into_inner()
1602 .map_err(|e| CmfError::Io(e.into_error()))?;
1603 f.seek(SeekFrom::Start(0))?;
1604 f.write_all(&env)?;
1605 f.write_all(&header_json)?;
1606 f.write_all(&dir_bytes)?;
1607 f.flush()?;
1608
1609 tracing::info!(
1610 "Wrote CMF v2 (streamed): {} ({} tensors, {:.1} MB)",
1611 self.path.display(),
1612 self.entries.len(),
1613 (data_off + data_len + masks_len + vocab_len + index_len) as f64 / 1e6
1614 );
1615 Ok(())
1616 }
1617}
1618
1619pub fn build_sparse_index(catalog: &MaskCatalog, arch: &ModelArch) -> Vec<SparseIndexEntry> {
1624 let mut out = Vec::new();
1625 for m in &catalog.masks {
1626 for li in 0..arch.num_layers {
1627 if !m.layer_alive(li) {
1628 continue;
1629 }
1630 let mut groups = Vec::new();
1631 if let Some(bits) = m.ffn_masks.get(li) {
1632 let n_groups = arch.intermediate_size.div_ceil(32);
1633 for g in 0..n_groups {
1634 let active = bits[g * 4..(g * 4 + 4).min(bits.len())]
1636 .iter()
1637 .any(|&b| b != 0);
1638 if active {
1639 groups.push(g as u16);
1640 }
1641 }
1642 }
1643 let mut heads = Vec::new();
1644 if let Some(bits) = m.head_masks.get(li) {
1645 for h in 0..arch.num_attention_heads {
1646 if bits
1647 .get(h / 8)
1648 .map(|b| b & (1 << (h % 8)) != 0)
1649 .unwrap_or(false)
1650 {
1651 heads.push(h as u8);
1652 }
1653 }
1654 }
1655 out.push(SparseIndexEntry {
1656 task_id: m.task_id,
1657 layer_idx: li,
1658 active_ffn_groups: groups,
1659 active_heads: heads,
1660 });
1661 }
1662 }
1663 out
1664}
1665
1666pub fn encode_sparse_index(entries: &[SparseIndexEntry]) -> Vec<u8> {
1669 let mut out = Vec::new();
1670 out.extend_from_slice(&(entries.len() as u32).to_le_bytes());
1671 out.extend_from_slice(&0u32.to_le_bytes());
1672 for e in entries {
1673 out.extend_from_slice(&e.task_id.to_le_bytes());
1674 out.extend_from_slice(&(e.layer_idx as u32).to_le_bytes());
1675 out.extend_from_slice(&(e.active_ffn_groups.len() as u32).to_le_bytes());
1676 out.extend_from_slice(&(e.active_heads.len() as u32).to_le_bytes());
1677 for g in &e.active_ffn_groups {
1678 out.extend_from_slice(&g.to_le_bytes());
1679 }
1680 out.extend_from_slice(&e.active_heads);
1681 while out.len() % 4 != 0 {
1682 out.push(0);
1683 }
1684 }
1685 out
1686}
1687
1688pub fn decode_sparse_index(bytes: &[u8]) -> Result<Vec<SparseIndexEntry>, CmfError> {
1689 let err = |msg: &str| CmfError::Parse(format!("sparse index: {msg}"));
1690 if bytes.len() < 8 {
1691 return Err(err("too short"));
1692 }
1693 let n = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
1694 let mut pos = 8usize;
1695 let mut out = Vec::with_capacity(n);
1696 for _ in 0..n {
1697 if pos + 16 > bytes.len() {
1698 return Err(err("entry header out of bounds"));
1699 }
1700 let task_id = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap());
1701 let layer_idx = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()) as usize;
1702 let n_groups = u32::from_le_bytes(bytes[pos + 8..pos + 12].try_into().unwrap()) as usize;
1703 let n_heads = u32::from_le_bytes(bytes[pos + 12..pos + 16].try_into().unwrap()) as usize;
1704 pos += 16;
1705 if pos + n_groups * 2 + n_heads > bytes.len() {
1706 return Err(err("entry data out of bounds"));
1707 }
1708 let mut groups = Vec::with_capacity(n_groups);
1709 for g in 0..n_groups {
1710 groups.push(u16::from_le_bytes(
1711 bytes[pos + g * 2..pos + g * 2 + 2].try_into().unwrap(),
1712 ));
1713 }
1714 pos += n_groups * 2;
1715 let heads = bytes[pos..pos + n_heads].to_vec();
1716 pos += n_heads;
1717 pos = pos.div_ceil(4) * 4;
1718 out.push(SparseIndexEntry {
1719 task_id,
1720 layer_idx,
1721 active_ffn_groups: groups,
1722 active_heads: heads,
1723 });
1724 }
1725 Ok(out)
1726}
1727
1728#[derive(Debug, thiserror::Error)]
1730pub enum CmfError {
1731 #[error("File not found: {0}")]
1732 FileNotFound(String),
1733 #[error("Invalid CMF magic bytes")]
1734 InvalidMagic,
1735 #[error("Unsupported CMF version: {0}")]
1736 UnsupportedVersion(u32),
1737 #[error("File requires unsupported features (bits {0:#x})")]
1738 UnsupportedFeature(u32),
1739 #[error("Unknown tensor dtype id: {0}")]
1740 UnknownDtype(u8),
1741 #[error("Tensor not found: {0}")]
1742 MissingTensor(String),
1743 #[error("Bounds error: {0}")]
1744 Bounds(String),
1745 #[error("IO error: {0}")]
1746 Io(#[from] io::Error),
1747 #[error("Parse error: {0}")]
1748 Parse(String),
1749}