cherenkov/qwen4_exp/
manifest.rs

1//! On-disk packed tensor layouts, independent of runtime sequence state.
2
3use anyhow::{Context, Result};
4use serde::{Deserialize, Serialize};
5use std::path::Path;
6
7/// Where a dense tensor lives inside `dense.bin`.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct DenseEntry {
10    pub name: String,
11    pub dtype: String,
12    pub shape: Vec<usize>,
13    pub offset: u64,
14    pub nbytes: u64,
15}
16
17/// Expert record layout inside `experts.bin`. One record per (layer, expert):
18/// three packed 4-bit weight matrices followed by their scales and biases,
19/// padded to a page multiple so each record can be its own Metal buffer.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct ExpertLayout {
22    pub layers: usize,
23    pub experts: usize,
24    pub inter: usize,
25    pub hidden: usize,
26    pub group: usize,
27    pub record_bytes: u64,
28    pub record_stride: u64,
29    /// Tensor prefixes in record order (main layers, then MTP layers).
30    pub layer_prefixes: Vec<String>,
31    pub gate_w: u64,
32    pub up_w: u64,
33    pub down_w: u64,
34    pub gate_s: u64,
35    pub gate_b: u64,
36    pub up_s: u64,
37    pub up_b: u64,
38    pub down_s: u64,
39    pub down_b: u64,
40}
41
42/// N-gram table layout inside `ngram.bin`: one `row_bytes` record per hashed
43/// id holding the packed 4-bit row, then its scales, then its biases.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct NgramLayout {
46    pub rows: u64,
47    pub row_bytes: u64,
48    pub dim: usize,
49    pub group: usize,
50    pub weight_bytes: u64,
51    pub scale_bytes: u64,
52    pub head_offsets: Vec<u64>,
53    pub head_vocab_sizes: Vec<u64>,
54    pub layer_multipliers: Vec<i64>,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct Manifest {
59    pub version: u32,
60    pub source_dir: String,
61    pub dense: Vec<DenseEntry>,
62    pub experts: ExpertLayout,
63    pub ngram: NgramLayout,
64}
65
66impl Manifest {
67    pub fn load(packed_dir: &Path) -> Result<Self> {
68        let path = packed_dir.join("manifest.json");
69        let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
70        let manifest: Self = serde_json::from_slice(&bytes).context("manifest.json")?;
71
72        anyhow::ensure!(
73            manifest.version == 1,
74            "unsupported packed layout version {}",
75            manifest.version
76        );
77
78        manifest.experts.validate_dimensions()?;
79
80        Ok(manifest)
81    }
82
83    pub fn dense(&self, name: &str) -> Result<&DenseEntry> {
84        self.dense
85            .iter()
86            .find(|e| e.name == name)
87            .with_context(|| format!("dense tensor {name:?} not in manifest"))
88    }
89}
90
91pub const PAGE: u64 = 16384;
92
93impl ExpertLayout {
94    pub(crate) fn validate_dimensions(&self) -> Result<()> {
95        anyhow::ensure!(
96            self.layers > 0 && self.experts > 0,
97            "expert layer dimensions must be nonzero"
98        );
99        anyhow::ensure!(
100            self.layer_prefixes.len() == self.layers,
101            "expert layer count does not match layer prefixes"
102        );
103        anyhow::ensure!(
104            self.layers.checked_mul(self.experts).is_some(),
105            "expert record count overflows"
106        );
107
108        Ok(())
109    }
110}
cherenkov/qwen4_exp/manifest.rs

cherenkov/qwen4_exp/
manifest.rs
