1use super::attn::freqs_cis;
12use super::compress::{
13 CompAttnCfg, CompAttnRun, IndexerCfg, LayerAttnState, SharedRuntime, yarn_freqs_cis,
14};
15use super::engram::{EngramTables, NgramHashState, regen_bf16_matrix, regen_bf16_qk, regen_table};
16use super::moe::MoeCfg;
17use super::{Golden, regen_param};
18
19mod forward;
20
21pub use forward::forward;
22
23pub struct ModelCfg {
24 pub dim: usize,
25 pub hc: usize,
26 pub iters: usize,
27 pub hc_eps: f32,
28 pub eps: f32,
29 pub n_layers: usize,
30 pub vocab: usize,
31 pub max_seq: usize,
32 pub n_heads: usize,
33 pub head_dim: usize,
34 pub rope_dim: usize,
35 pub q_rank: usize,
36 pub o_rank: usize,
37 pub groups: usize,
38 pub window: usize,
39 pub ratios: Vec<usize>,
40 pub kv_sources: Vec<usize>,
41 pub index_sources: Vec<usize>,
42 pub cand_src: i64,
43 pub cand_topk_blocks: usize,
44 pub cand_block: usize,
45 pub index_heads: usize,
46 pub index_hd: usize,
47 pub index_topk: usize,
48 pub rope_theta: f32,
49 pub compress_rope_theta: f32,
50 pub rope_factor: f32,
51 pub orig_seq: usize,
52 pub beta_fast: f32,
53 pub beta_slow: f32,
54 pub moe: MoeCfg,
55}
56
57impl ModelCfg {
58 pub fn from_golden(g: &Golden) -> Self {
59 let u = |k: &str| g.fixture_u64(k) as usize;
60 let f = |k: &str| g.fixture_f64(k) as f32;
61 let list = |k: &str| g.fixture_usize_list(k);
62 let dim = u("dim");
63 ModelCfg {
64 dim,
65 hc: u("hc_mult"),
66 iters: u("hc_sinkhorn_iters"),
67 hc_eps: f("hc_eps"),
68 eps: f("norm_eps"),
69 n_layers: u("n_layers"),
70 vocab: u("vocab_size"),
71 max_seq: u("max_seq_len"),
72 n_heads: u("n_heads"),
73 head_dim: u("head_dim"),
74 rope_dim: u("rope_head_dim"),
75 q_rank: u("q_lora_rank"),
76 o_rank: u("o_lora_rank"),
77 groups: u("o_groups"),
78 window: u("window_size"),
79 ratios: list("compress_ratios"),
80 kv_sources: list("kv_source_layers"),
81 index_sources: list("index_source_layers"),
82 cand_src: g.fixture_i64("candidate_source_layer"),
83 cand_topk_blocks: u("candidate_topk_blocks"),
84 cand_block: u("candidate_block_size"),
85 index_heads: u("index_n_heads"),
86 index_hd: u("index_head_dim"),
87 index_topk: u("index_topk"),
88 rope_theta: f("rope_theta"),
89 compress_rope_theta: f("compress_rope_theta"),
90 rope_factor: f("rope_factor"),
91 orig_seq: u("original_seq_len"),
92 beta_fast: f("beta_fast"),
93 beta_slow: f("beta_slow"),
94 moe: MoeCfg {
95 dim,
96 inter: u("moe_inter_dim"),
97 n_routed: u("n_routed_experts"),
98 topk: u("n_activated_experts"),
99 gate_temp: f("gate_temp"),
100 norm_topk_prob: true,
101 route_scale: f("route_scale"),
102 swiglu_limit: f("swiglu_limit"),
103 },
104 }
105 }
106
107 pub fn attn_cfg(&self, layer: usize) -> CompAttnCfg {
108 let ratio = self.ratios[layer];
109 CompAttnCfg {
110 dim: self.dim,
111 n_heads: self.n_heads,
112 head_dim: self.head_dim,
113 rope_dim: self.rope_dim,
114 q_rank: self.q_rank,
115 o_rank: self.o_rank,
116 groups: self.groups,
117 window: self.window,
118 eps: self.eps,
119 ratio,
120 is_kv_source: self.kv_sources.contains(&layer),
121 is_index_source: self.index_sources.contains(&layer),
122 is_candidate_source: self.cand_src == layer as i64,
123 uses_candidates: self.cand_src >= 0 && (self.cand_src as usize) < layer,
124 }
125 }
126
127 pub fn indexer_cfg(&self) -> IndexerCfg {
128 IndexerCfg {
129 n_heads: self.index_heads,
130 index_hd: self.index_hd,
131 rope_dim: self.rope_dim,
132 q_rank: self.q_rank,
133 dim: self.dim,
134 hd: self.head_dim,
135 index_topk: self.index_topk,
136 cand_topk_blocks: self.cand_topk_blocks,
137 cand_block: self.cand_block,
138 eps: self.eps,
139 }
140 }
141}
142
143pub struct EngramWeights {
144 pub table: Vec<f32>,
145 pub wkv: Vec<f32>,
146 pub q: Vec<f32>,
147 pub k: Vec<f32>,
148 pub hash_index: usize,
149}
150
151pub struct LayerWeights {
152 pub hc_attn_fn: Vec<f32>,
153 pub hc_attn_scale: Vec<f32>,
154 pub hc_attn_base: Vec<f32>,
155 pub hc_ffn_fn: Vec<f32>,
156 pub hc_ffn_scale: Vec<f32>,
157 pub hc_ffn_base: Vec<f32>,
158 pub attn_norm: Vec<f32>,
159 pub ffn_norm: Vec<f32>,
160 pub sink: Vec<f32>,
161 pub wq_a: Vec<f32>,
162 pub q_norm: Vec<f32>,
163 pub wq_b: Vec<f32>,
164 pub wkv: Vec<f32>,
165 pub kv_norm: Vec<f32>,
166 pub wo_a: Vec<f32>,
167 pub wo_b: Vec<f32>,
168 pub comp_wkv: Option<Vec<f32>>,
169 pub comp_wgate: Option<Vec<f32>>,
170 pub comp_norm: Option<Vec<f32>>,
171 pub idx_wq_b: Option<Vec<f32>>,
172 pub idx_weights_proj: Option<Vec<f32>>,
173 pub idx_wk: Option<Vec<f32>>,
174 pub idx_k_norm: Option<Vec<f32>>,
175 pub gate_w: Vec<f32>,
176 pub gate_bias: Vec<f32>,
177 pub experts: Vec<(Vec<f32>, Vec<f32>, Vec<f32>)>,
178 pub shared: (Vec<f32>, Vec<f32>, Vec<f32>),
179 pub engram: Option<EngramWeights>,
180}
181
182pub struct ModelWeights {
183 pub embed: Vec<f32>,
184 pub norm: Vec<f32>,
185 pub head: Vec<f32>,
186 pub layers: Vec<LayerWeights>,
187}
188
189impl ModelWeights {
190 pub fn from_golden(g: &Golden, c: &ModelCfg, t: &EngramTables) -> Self {
191 let layers = (0..c.n_layers)
192 .map(|l| {
193 let p = |n: &str| regen_param(g, &format!("layers.{l}.{n}"));
194 let opt = |n: &str, on: bool| if on { Some(p(n)) } else { None };
195 let ac = c.attn_cfg(l);
196 let has_comp = ac.is_kv_source;
197 let has_idx = ac.is_index_source;
198 let engram = t
199 .layer_ids
200 .iter()
201 .position(|&x| x == l)
202 .map(|hi| EngramWeights {
203 table: regen_table(t, hi),
204 wkv: regen_bf16_matrix(
205 &format!("layers.{l}.engram.wkv.weight"),
206 c.dim * (c.hc + 1),
207 t.n_hash_cols() * t.head_dim,
208 ),
209 q: regen_bf16_qk(&format!("layers.{l}.engram.q_weight"), c.hc * c.dim),
210 k: regen_bf16_qk(&format!("layers.{l}.engram.k_weight"), c.hc * c.dim),
211 hash_index: hi,
212 });
213 LayerWeights {
214 hc_attn_fn: p("hc_attn_fn"),
215 hc_attn_scale: p("hc_attn_scale"),
216 hc_attn_base: p("hc_attn_base"),
217 hc_ffn_fn: p("hc_ffn_fn"),
218 hc_ffn_scale: p("hc_ffn_scale"),
219 hc_ffn_base: p("hc_ffn_base"),
220 attn_norm: p("attn_norm.weight"),
221 ffn_norm: p("ffn_norm.weight"),
222 sink: p("attn.attn_sink"),
223 wq_a: p("attn.wq_a.weight"),
224 q_norm: p("attn.q_norm.weight"),
225 wq_b: p("attn.wq_b.weight"),
226 wkv: p("attn.wkv.weight"),
227 kv_norm: p("attn.kv_norm.weight"),
228 wo_a: p("attn.wo_a.weight"),
229 wo_b: p("attn.wo_b.weight"),
230 comp_wkv: opt("attn.compressor.wkv.weight", has_comp),
231 comp_wgate: opt("attn.compressor.wgate.weight", has_comp && ac.ratio > 1),
232 comp_norm: opt("attn.compressor.norm.weight", has_comp),
233 idx_wq_b: opt("attn.indexer.wq_b.weight", has_idx),
234 idx_weights_proj: opt("attn.indexer.weights_proj.weight", has_idx),
235 idx_wk: opt("attn.indexer.wk.weight", has_idx && has_comp),
236 idx_k_norm: opt("attn.indexer.k_norm.weight", has_idx && has_comp),
237 gate_w: p("ffn.gate.weight"),
238 gate_bias: p("ffn.gate.bias"),
239 experts: (0..c.moe.n_routed)
240 .map(|e| {
241 (
242 p(&format!("ffn.experts.{e}.w1.weight")),
243 p(&format!("ffn.experts.{e}.w2.weight")),
244 p(&format!("ffn.experts.{e}.w3.weight")),
245 )
246 })
247 .collect(),
248 shared: (
249 p("ffn.shared_experts.w1.weight"),
250 p("ffn.shared_experts.w2.weight"),
251 p("ffn.shared_experts.w3.weight"),
252 ),
253 engram,
254 }
255 })
256 .collect();
257 ModelWeights {
258 embed: regen_param(g, "embed.weight"),
259 norm: regen_param(g, "norm.weight"),
260 head: regen_param(g, "head.weight"),
261 layers,
262 }
263 }
264}
265
266pub struct ModelState<'a> {
267 pub hash: NgramHashState<'a>,
268 pub layers: Vec<LayerAttnState>,
269 pub shared: SharedRuntime,
270 pub fcs: Vec<Vec<(f32, f32)>>,
271}
272
273impl<'a> ModelState<'a> {
274 pub fn new(c: &ModelCfg, t: &'a EngramTables) -> Self {
275 ModelState {
276 hash: NgramHashState::new(t, 1, c.max_seq),
277 layers: (0..c.n_layers)
278 .map(|l| LayerAttnState::new(&c.attn_cfg(l), c.max_seq, c.index_hd))
279 .collect(),
280 shared: SharedRuntime::default(),
281 fcs: (0..c.n_layers)
282 .map(|l| {
283 if c.ratios[l] == 0 {
284 freqs_cis(c.rope_dim, c.max_seq, c.rope_theta)
285 } else {
286 yarn_freqs_cis(
287 c.rope_dim,
288 c.max_seq,
289 c.orig_seq,
290 c.compress_rope_theta,
291 c.rope_factor,
292 c.beta_fast,
293 c.beta_slow,
294 )
295 }
296 })
297 .collect(),
298 }
299 }
300}
301
302pub struct LayerTrace {
303 pub engram_out: Option<Vec<f32>>,
304 pub h_in: Vec<f32>,
305 pub pre_mix_in: Vec<f32>,
306 pub attn_pre: Vec<f32>,
307 pub attn_post: Vec<f32>,
308 pub attn_comb: Vec<f32>,
309 pub attn_in: Vec<f32>,
310 pub attn: CompAttnRun,
311 pub ffn_pre: Vec<f32>,
312 pub ffn_post: Vec<f32>,
313 pub ffn_comb: Vec<f32>,
314 pub ffn_in: Vec<f32>,
315 pub moe_weights: Vec<f32>,
316 pub moe_indices: Vec<usize>,
317 pub ffn_out: Vec<f32>,
318 pub h_out: Vec<f32>,
319 pub shared_compress_kv: Vec<f32>,
321 pub shared_index_k: Vec<f32>,
322 pub shared_topk: Vec<i32>,
323 pub shared_topk_w: usize,
324 pub shared_cand: Vec<bool>,
325 pub shared_cand_w: usize,
326}
327
328pub struct StepTrace {
329 pub tokens: usize,
330 pub embed: Vec<f32>,
331 pub layers: Vec<LayerTrace>,
332 pub h_final: Vec<f32>,
333 pub head_in: Vec<f32>,
334 pub logits: Vec<f32>,
335}
336
337#[cfg(test)]
338#[path = "model_tests.rs"]
339mod tests;