qualia_core_db/gguf_bridge/wasm_cpu/
model.rs

1use std::fmt;
2use std::sync::Arc;
3
4use crate::gguf_sharder::{GgufTensorIndex, GgufTokenizer};
5
6use super::{CPU_WASM_DEFAULT_CONTEXT, CPU_WASM_MAX_CONTEXT};
7
8const MAX_EMBEDDING: usize = 8192;
9const MAX_FFN: usize = 16_384;
10const MIN_REAL_VOCAB: usize = 1000;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub enum CpuWasmError {
14    InvalidModel(String),
15    UnsupportedModel(String),
16    ContextExceeded { position: usize, max_context: usize },
17    MissingTensor { layer: u32, role: &'static str },
18    KernelFailed { layer: u32, role: &'static str },
19}
20
21impl fmt::Display for CpuWasmError {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        match self {
24            Self::InvalidModel(reason) => write!(f, "invalid model: {reason}"),
25            Self::UnsupportedModel(reason) => write!(f, "unsupported CPU-WASM model: {reason}"),
26            Self::ContextExceeded {
27                position,
28                max_context,
29            } => write!(
30                f,
31                "token position {position} exceeds CPU-WASM context {max_context}"
32            ),
33            Self::MissingTensor { layer, role } => {
34                write!(f, "layer {layer} is missing required tensor {role}")
35            }
36            Self::KernelFailed { layer, role } => {
37                write!(f, "layer {layer} CPU kernel failed for {role}")
38            }
39        }
40    }
41}
42
43impl std::error::Error for CpuWasmError {}
44
45#[derive(Debug, Clone, Copy, PartialEq)]
46pub struct CpuWasmStep {
47    pub token_id: u32,
48    pub max_logit: f32,
49}
50
51/// Cold-built, allocation-stable CPU decode plan.
52pub struct CpuWasmEngine {
53    pub(super) model: Arc<[u8]>,
54    pub(super) index: GgufTensorIndex,
55    pub(super) tokenizer: GgufTokenizer,
56    pub(super) max_context: usize,
57    pub(super) n_embd: usize,
58    pub(super) n_ffn: usize,
59    pub(super) n_layer: usize,
60    pub(super) n_head: usize,
61    pub(super) n_kv_head: usize,
62    pub(super) head_dim: usize,
63    pub(super) kv_plane_elems: usize,
64    pub(super) kv: Vec<f32>,
65    pub(super) hidden: Vec<f32>,
66    pub(super) normed: Vec<f32>,
67    pub(super) norm_weight: Vec<f32>,
68    pub(super) q: Vec<f32>,
69    pub(super) k: Vec<f32>,
70    pub(super) v: Vec<f32>,
71    pub(super) attention: Vec<f32>,
72    pub(super) projection: Vec<f32>,
73    pub(super) gate: Vec<f32>,
74    pub(super) up: Vec<f32>,
75    pub(super) scores: Vec<f32>,
76    pub(super) logits: Vec<f32>,
77}
78
79impl CpuWasmEngine {
80    pub fn new(model: Arc<[u8]>) -> Result<Self, CpuWasmError> {
81        Self::new_with_context(model, CPU_WASM_DEFAULT_CONTEXT)
82    }
83
84    pub fn new_with_context(model: Arc<[u8]>, max_context: usize) -> Result<Self, CpuWasmError> {
85        if max_context == 0 || max_context > CPU_WASM_MAX_CONTEXT {
86            return Err(CpuWasmError::UnsupportedModel(format!(
87                "context {max_context} is outside 1..={CPU_WASM_MAX_CONTEXT}"
88            )));
89        }
90        let (tokenizer, index) = if crate::p64_weight::has_p64_magic(&model) {
91            let p64 = crate::p64_weight::P64TensorIndex::from_p64(&model)
92                .map_err(|e| CpuWasmError::InvalidModel(format!("P64 index: {e}")))?;
93            let tokenizer = GgufTokenizer::from_p64_section(p64.tokenizer_bytes(&model))
94                .ok_or_else(|| CpuWasmError::InvalidModel("P64 has no Q42T tokenizer".into()))?;
95            (tokenizer, p64.to_gguf_index())
96        } else {
97            (
98                GgufTokenizer::from_gguf(&model),
99                GgufTensorIndex::from_gguf(&model),
100            )
101        };
102
103        if (tokenizer.vocab_len() as usize) < MIN_REAL_VOCAB {
104            return Err(CpuWasmError::InvalidModel(format!(
105                "tokenizer vocabulary {} is fallback-only",
106                tokenizer.vocab_len()
107            )));
108        }
109        index
110            .hyperparams
111            .decode_supported()
112            .map_err(CpuWasmError::UnsupportedModel)?;
113
114        let n_embd = index.emb_dim();
115        let n_layer = index.hyperparams.n_layer as usize;
116        let n_head = index.hyperparams.n_head as usize;
117        let n_kv_head = index.hyperparams.effective_n_kv_head() as usize;
118        let head_dim = index.hyperparams.head_dim() as usize;
119        if n_embd == 0 || n_embd > MAX_EMBEDDING || n_layer == 0 || n_head == 0 {
120            return Err(CpuWasmError::UnsupportedModel(format!(
121                "dimensions emb={n_embd}, layers={n_layer}, heads={n_head}"
122            )));
123        }
124        if n_head * head_dim != n_embd || n_kv_head == 0 || n_head % n_kv_head != 0 {
125            return Err(CpuWasmError::UnsupportedModel(format!(
126                "non-Llama attention shape heads={n_head}, kv_heads={n_kv_head}, head_dim={head_dim}, emb={n_embd}"
127            )));
128        }
129        let layer0 = index.get_layer_tensors(0);
130        let n_ffn = layer0
131            .ffn_gate
132            .map(|t| t.dims[1] as usize)
133            .filter(|&n| n > 0 && n <= MAX_FFN)
134            .ok_or_else(|| CpuWasmError::UnsupportedModel("missing or oversized FFN".into()))?;
135        let vocab = index.vocab_dim();
136        if vocab == 0 || index.logits_projection_info().is_none() {
137            return Err(CpuWasmError::InvalidModel(
138                "missing output projection".into(),
139            ));
140        }
141
142        let per_token_kv = n_layer
143            .checked_mul(n_kv_head)
144            .and_then(|n| n.checked_mul(head_dim))
145            .and_then(|n| n.checked_mul(2))
146            .ok_or_else(|| CpuWasmError::UnsupportedModel("KV dimensions overflow".into()))?;
147        let kv_plane_elems = n_layer
148            .checked_mul(max_context)
149            .and_then(|n| n.checked_mul(n_kv_head))
150            .and_then(|n| n.checked_mul(head_dim))
151            .ok_or_else(|| CpuWasmError::UnsupportedModel("KV allocation overflow".into()))?;
152        let _kv_bytes = per_token_kv
153            .checked_mul(max_context)
154            .and_then(|n| n.checked_mul(core::mem::size_of::<f32>()))
155            .ok_or_else(|| CpuWasmError::UnsupportedModel("KV working-set overflow".into()))?;
156
157        Ok(Self {
158            model,
159            index,
160            tokenizer,
161            max_context,
162            n_embd,
163            n_ffn,
164            n_layer,
165            n_head,
166            n_kv_head,
167            head_dim,
168            kv_plane_elems,
169            kv: vec![0.0; kv_plane_elems * 2],
170            hidden: vec![0.0; n_embd],
171            normed: vec![0.0; n_embd],
172            norm_weight: vec![0.0; n_embd],
173            q: vec![0.0; n_embd],
174            k: vec![0.0; n_kv_head * head_dim],
175            v: vec![0.0; n_kv_head * head_dim],
176            attention: vec![0.0; n_embd],
177            projection: vec![0.0; n_embd.max(n_ffn)],
178            gate: vec![0.0; n_ffn],
179            up: vec![0.0; n_ffn],
180            scores: vec![0.0; max_context],
181            logits: vec![0.0; vocab],
182        })
183    }
184
185    pub fn tokenizer(&self) -> &GgufTokenizer {
186        &self.tokenizer
187    }
188
189    pub fn vocab_len(&self) -> u32 {
190        self.tokenizer.vocab_len()
191    }
192
193    pub fn max_context(&self) -> usize {
194        self.max_context
195    }
196
197    pub fn working_set_bytes(&self) -> usize {
198        (self.kv.capacity()
199            + self.hidden.capacity()
200            + self.normed.capacity()
201            + self.norm_weight.capacity()
202            + self.q.capacity()
203            + self.k.capacity()
204            + self.v.capacity()
205            + self.attention.capacity()
206            + self.projection.capacity()
207            + self.gate.capacity()
208            + self.up.capacity()
209            + self.scores.capacity()
210            + self.logits.capacity())
211            * core::mem::size_of::<f32>()
212    }
213
214    pub fn reset(&mut self) {
215        self.kv.fill(0.0);
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn llm_memory_domain_is_independent_from_slg_sentinel() {
225        assert_eq!(CPU_WASM_DEFAULT_CONTEXT, 512);
226        assert_eq!(CPU_WASM_MAX_CONTEXT, 4096);
227    }
228}
qualia_core_db/gguf_bridge/wasm_cpu/model.rs

qualia_core_db/gguf_bridge/wasm_cpu/
model.rs
