1use std::path::Path;
12
13use anyhow::{Context, Result};
14use burn::tensor::{activation, Device, Tensor, TensorData};
15use zsfm_burn::linear::{burn_linear_bias, burn_linear_nobias};
16use zsfm_burn::norm::burn_rms_norm_nd;
17use zsfm_burn::rope::apply_interleaved_rope;
18use zsfm_burn::{open_gguf_weights, GgufWeights};
19
20use crate::config::Moirai2Config;
21
22fn device() -> Device {
23 Device::flex()
24}
25
26fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
27 weights.load_2d(name)
28}
29
30fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
31 weights.load_1d(name)
32}
33
34fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
35 t.to_data()
36 .try_to_vec()
37 .map_err(|e| anyhow::anyhow!("{e:?}"))
38}
39
40struct BurnResidualBlockW {
41 hidden_w: Tensor<2>,
42 hidden_b: Tensor<1>,
43 output_w: Tensor<2>,
44 output_b: Tensor<1>,
45 residual_w: Tensor<2>,
46 residual_b: Tensor<1>,
47}
48
49struct BurnEncoderBlock {
50 norm1_w: Tensor<1>,
51 norm2_w: Tensor<1>,
52 attn_qkv_w: Tensor<2>,
53 attn_o_w: Tensor<2>,
54 attn_qn_w: Tensor<1>,
55 attn_kn_w: Tensor<1>,
56 attn_vbias_t: Tensor<3>,
57 ffn_fc1_w: Tensor<2>,
58 ffn_fc2_w: Tensor<2>,
59 ffn_gate_w: Tensor<2>,
60}
61
62pub struct BurnMoirai2Model {
63 config: Moirai2Config,
64 in_proj: BurnResidualBlockW,
65 blocks: Vec<BurnEncoderBlock>,
66 norm_f_w: Tensor<1>,
67 out_proj: BurnResidualBlockW,
68 rope_cos: Vec<f32>,
69 rope_sin: Vec<f32>,
70}
71
72fn load_residual_block(weights: &GgufWeights, prefix: &str) -> Result<BurnResidualBlockW> {
73 let p = |s: &str| format!("{prefix}.{s}");
74 Ok(BurnResidualBlockW {
75 hidden_w: load_t(weights, &p("hidden.weight"))?,
76 hidden_b: load_v(weights, &p("hidden.bias"))?,
77 output_w: load_t(weights, &p("output.weight"))?,
78 output_b: load_v(weights, &p("output.bias"))?,
79 residual_w: load_t(weights, &p("residual.weight"))?,
80 residual_b: load_v(weights, &p("residual.bias"))?,
81 })
82}
83
84fn residual_block_fwd(x: Tensor<2>, w: &BurnResidualBlockW) -> Tensor<2> {
85 let hidden = activation::silu(burn_linear_bias(
86 x.clone(),
87 w.hidden_w.clone(),
88 w.hidden_b.clone(),
89 ));
90 let output = burn_linear_bias(hidden, w.output_w.clone(), w.output_b.clone());
91 let residual = burn_linear_bias(x, w.residual_w.clone(), w.residual_b.clone());
92 output + residual
93}
94
95fn qk_norm_heads(
96 x: Tensor<3>,
97 weight: Tensor<1>,
98 seq_len: usize,
99 n_heads: usize,
100 head_dim: usize,
101) -> Tensor<3> {
102 burn_rms_norm_nd(x.reshape([seq_len * n_heads, head_dim]), weight, 1e-6)
103 .reshape([seq_len, n_heads, head_dim])
104}
105
106fn make_causal_mask_t(seq_len: usize) -> Tensor<3> {
107 let dev = device();
108 let mut mask = vec![0.0f32; seq_len * seq_len];
109 for i in 0..seq_len {
110 for j in (i + 1)..seq_len {
111 mask[i * seq_len + j] = f32::NEG_INFINITY;
112 }
113 }
114 Tensor::<3>::from_data(TensorData::new(mask, [1, seq_len, seq_len]), &dev)
115}
116
117fn add_decode_causal_mask(scores: Tensor<3>, new_len: usize, cached_len: usize) -> Tensor<3> {
118 let dev = device();
119 let total = cached_len + new_len;
120 let mut mask = vec![0.0f32; new_len * total];
121 for q_rel in 0..new_len {
122 for k_abs in (cached_len + q_rel + 1)..total {
123 mask[q_rel * total + k_abs] = f32::NEG_INFINITY;
124 }
125 }
126 let mask_t = Tensor::<3>::from_data(TensorData::new(mask, [1, new_len, total]), &dev);
127 scores + mask_t
128}
129
130fn swiglu_ffn(x: Tensor<2>, fc1_w: Tensor<2>, fc2_w: Tensor<2>, gate_w: Tensor<2>) -> Tensor<2> {
131 let content = activation::silu(burn_linear_nobias(x.clone(), fc1_w));
132 let gate = burn_linear_nobias(x, gate_w);
133 burn_linear_nobias(content * gate, fc2_w)
134}
135
136impl BurnMoirai2Model {
137 pub fn load(gguf_path: &Path, config: Moirai2Config) -> Result<Self> {
138 let weights = open_gguf_weights(gguf_path)
139 .with_context(|| format!("open {}", gguf_path.display()))?;
140 let dev = device();
141
142 let in_proj = load_residual_block(&weights, "in_proj")?;
143 let mut blocks = Vec::with_capacity(config.n_layers);
144 for n in 0..config.n_layers {
145 let p = |s: &str| format!("blk.{n}.{s}");
146 let qw = load_t(&weights, &p("attn_q.weight"))?;
147 let kw = load_t(&weights, &p("attn_k.weight"))?;
148 let vw = load_t(&weights, &p("attn_v.weight"))?;
149 let attn_qkv_w = Tensor::cat(vec![qw, kw, vw], 0);
150 let vbias_raw: Vec<f32> = weights.load_flat(&p("attn_vbias.weight"))?;
151 let n_heads = config.n_heads;
152 let same_var_bias: Vec<f32> = (0..n_heads).map(|h| vbias_raw[n_heads + h]).collect();
153 let attn_vbias_t =
154 Tensor::<3>::from_data(TensorData::new(same_var_bias, [n_heads, 1, 1]), &dev);
155 blocks.push(BurnEncoderBlock {
156 norm1_w: load_v(&weights, &p("norm1.weight"))?,
157 norm2_w: load_v(&weights, &p("norm2.weight"))?,
158 attn_qkv_w,
159 attn_o_w: load_t(&weights, &p("attn_o.weight"))?,
160 attn_qn_w: load_v(&weights, &p("attn_qn.weight"))?,
161 attn_kn_w: load_v(&weights, &p("attn_kn.weight"))?,
162 attn_vbias_t,
163 ffn_fc1_w: load_t(&weights, &p("ffn_fc1.weight"))?,
164 ffn_fc2_w: load_t(&weights, &p("ffn_fc2.weight"))?,
165 ffn_gate_w: load_t(&weights, &p("ffn_gate.weight"))?,
166 });
167 }
168 let norm_f_w = load_v(&weights, "norm_f.weight")?;
169 let out_proj = load_residual_block(&weights, "out_proj")?;
170
171 let half_rope = config.rope_dim / 2;
172 let max_pos = config.max_ctx_tokens() + 256;
173 let inv_freq: Vec<f32> = (0..half_rope)
174 .map(|i| 1.0_f32 / 10000_f32.powf(2.0 * i as f32 / config.rope_dim as f32))
175 .collect();
176 let mut rope_cos = vec![0.0f32; max_pos * half_rope];
177 let mut rope_sin = vec![0.0f32; max_pos * half_rope];
178 for pos in 0..max_pos {
179 for i in 0..half_rope {
180 let theta = pos as f32 * inv_freq[i];
181 rope_cos[pos * half_rope + i] = theta.cos();
182 rope_sin[pos * half_rope + i] = theta.sin();
183 }
184 }
185 Ok(Self {
186 config,
187 in_proj,
188 blocks,
189 norm_f_w,
190 out_proj,
191 rope_cos,
192 rope_sin,
193 })
194 }
195
196 fn apply_partial_rope(
197 &self,
198 x: Tensor<3>,
199 time_ids: &[usize],
200 n_heads: usize,
201 head_dim: usize,
202 ) -> Result<Tensor<3>> {
203 let dims = x.dims();
204 let mut data = to_host(&x)?;
205 apply_interleaved_rope(
206 &mut data,
207 time_ids,
208 n_heads,
209 dims[1],
210 head_dim,
211 self.config.rope_dim,
212 &self.rope_cos,
213 &self.rope_sin,
214 )?;
215 Ok(Tensor::<3>::from_data(
216 TensorData::new(data, dims),
217 &device(),
218 ))
219 }
220
221 pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
222 let cfg = &self.config;
223 let ps = cfg.patch_size;
224 let max_ctx_len = cfg.max_ctx_tokens() * ps;
225 let ctx: &[f32] = if context.len() > max_ctx_len {
226 &context[context.len() - max_ctx_len..]
227 } else {
228 context
229 };
230 let n = ctx.len() as f64;
231 let loc = ctx.iter().map(|&v| v as f64).sum::<f64>() / n;
232 let var = ctx.iter().map(|&v| (v as f64 - loc).powi(2)).sum::<f64>() / (n - 1.0).max(1.0);
233 let scale = ((var + 1e-5_f64).sqrt()) as f32;
234
235 let ctx_norm: Vec<f32> = ctx.iter().map(|&v| (v - loc as f32) / scale).collect();
236 let rem = ctx_norm.len() % ps;
237 let ctx_padded: Vec<f32> = if rem != 0 {
238 let mut padded = vec![0.0f32; ps - rem];
239 padded.extend_from_slice(&ctx_norm);
240 padded
241 } else {
242 ctx_norm
243 };
244 let n_ctx = ctx_padded.len() / ps;
245 let ctx_flat: Vec<f32> = (0..n_ctx)
246 .flat_map(|i| {
247 let mut tok = ctx_padded[i * ps..(i + 1) * ps].to_vec();
248 tok.extend(vec![1.0f32; ps]);
249 tok
250 })
251 .collect();
252 let ctx_time_ids: Vec<usize> = (0..n_ctx).collect();
253 let num_pt = cfg.num_predict_token;
254 let num_q = cfg.num_quantiles;
255 let mq = cfg.median_quantile;
256 let n_future_patches = (horizon + ps - 1) / ps;
257
258 let dev = device();
259 let input_t = Tensor::<2>::from_data(TensorData::new(ctx_flat, [n_ctx, ps * 2]), &dev);
260 let h_ctx = residual_block_fwd(input_t, &self.in_proj);
261 let (h_enc, mut kv_cache) = self.prefill_encoder(h_ctx, &ctx_time_ids, n_ctx)?;
262 let last: Tensor<2> = h_enc.narrow(0, n_ctx - 1, 1);
263 let last_norm = burn_rms_norm_nd(last, self.norm_f_w.clone(), 1e-6);
264 let first_pred: Vec<f32> = to_host(&residual_block_fwd(last_norm, &self.out_proj))?;
265 let mut prev_patches: Vec<Vec<f32>> = (0..num_pt)
266 .map(|pt| {
267 let q_start = pt * num_q * ps + mq * ps;
268 first_pred[q_start..q_start + ps].to_vec()
269 })
270 .collect();
271 let mut collected_patches: Vec<Vec<f32>> = Vec::new();
272 let n_take = n_future_patches.min(num_pt);
273 for patch in prev_patches.iter().take(n_take) {
274 collected_patches.push(patch.clone());
275 }
276
277 let mut cached_len = n_ctx;
278 while collected_patches.len() < n_future_patches {
279 let new_time_ids: Vec<usize> = (cached_len..cached_len + num_pt).collect();
280 let new_flat: Vec<f32> = prev_patches
281 .iter()
282 .flat_map(|patch| {
283 let mut tok = patch.clone();
284 tok.extend(vec![0.0f32; ps]);
285 tok
286 })
287 .collect();
288 let new_t = Tensor::<2>::from_data(TensorData::new(new_flat, [num_pt, ps * 2]), &dev);
289 let h_in = residual_block_fwd(new_t, &self.in_proj);
290 let h_dec = self.decode_encoder(h_in, &new_time_ids, &mut kv_cache, cached_len)?;
291 let last2: Tensor<2> = h_dec.narrow(0, num_pt - 1, 1);
292 let last_norm2 = burn_rms_norm_nd(last2, self.norm_f_w.clone(), 1e-6);
293 let pred: Vec<f32> = to_host(&residual_block_fwd(last_norm2, &self.out_proj))?;
294 let new_patches: Vec<Vec<f32>> = (0..num_pt)
295 .map(|pt| {
296 let q_start = pt * num_q * ps + mq * ps;
297 pred[q_start..q_start + ps].to_vec()
298 })
299 .collect();
300 let need = n_future_patches - collected_patches.len();
301 for patch in new_patches.iter().take(need.min(num_pt)) {
302 collected_patches.push(patch.clone());
303 }
304 cached_len += num_pt;
305 prev_patches = new_patches;
306 }
307
308 Ok(collected_patches
309 .iter()
310 .flat_map(|p| p.iter().copied())
311 .take(horizon)
312 .map(|v| v * scale + loc as f32)
313 .collect())
314 }
315
316 fn prefill_encoder(
317 &self,
318 mut h: Tensor<2>,
319 time_ids: &[usize],
320 seq_len: usize,
321 ) -> Result<(Tensor<2>, Vec<(Tensor<3>, Tensor<3>)>)> {
322 let mut kv_cache = Vec::with_capacity(self.blocks.len());
323 for blk in &self.blocks {
324 let (h_out, k, v) = self.prefill_block(h, blk, time_ids, seq_len)?;
325 h = h_out;
326 kv_cache.push((k, v));
327 }
328 Ok((h, kv_cache))
329 }
330
331 fn prefill_block(
332 &self,
333 h: Tensor<2>,
334 blk: &BurnEncoderBlock,
335 time_ids: &[usize],
336 seq_len: usize,
337 ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
338 let h_norm = burn_rms_norm_nd(h.clone(), blk.norm1_w.clone(), 1e-6);
339 let (attn_out, k, v) = self.prefill_attn(&h_norm, blk, time_ids, seq_len)?;
340 let h = h + attn_out;
341 let h_norm2 = burn_rms_norm_nd(h.clone(), blk.norm2_w.clone(), 1e-6);
342 let ffn_out = swiglu_ffn(
343 h_norm2,
344 blk.ffn_fc1_w.clone(),
345 blk.ffn_fc2_w.clone(),
346 blk.ffn_gate_w.clone(),
347 );
348 Ok((h + ffn_out, k, v))
349 }
350
351 fn prefill_attn(
352 &self,
353 h: &Tensor<2>,
354 blk: &BurnEncoderBlock,
355 time_ids: &[usize],
356 seq_len: usize,
357 ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
358 let cfg = &self.config;
359 let qkv = burn_linear_nobias(h.clone(), blk.attn_qkv_w.clone());
360 let d = cfg.d_model;
361 let q: Tensor<2> = qkv.clone().narrow(1, 0, d);
362 let k: Tensor<2> = qkv.clone().narrow(1, d, d);
363 let v: Tensor<2> = qkv.narrow(1, 2 * d, d);
364 let q: Tensor<3> = q.reshape([seq_len, cfg.n_heads, cfg.head_dim]);
365 let k: Tensor<3> = k.reshape([seq_len, cfg.n_heads, cfg.head_dim]);
366 let q = qk_norm_heads(q, blk.attn_qn_w.clone(), seq_len, cfg.n_heads, cfg.head_dim);
367 let k = qk_norm_heads(k, blk.attn_kn_w.clone(), seq_len, cfg.n_heads, cfg.head_dim);
368 let q: Tensor<3> = q.permute([1, 0, 2]);
369 let k: Tensor<3> = k.permute([1, 0, 2]);
370 let v: Tensor<3> = v
371 .reshape([seq_len, cfg.n_heads, cfg.head_dim])
372 .permute([1, 0, 2]);
373 let q = self.apply_partial_rope(q, time_ids, cfg.n_heads, cfg.head_dim)?;
374 let k = self.apply_partial_rope(k.clone(), time_ids, cfg.n_heads, cfg.head_dim)?;
375 let scale = 1.0f32 / (cfg.head_dim as f32).sqrt();
376 let scores = q.matmul(k.clone().permute([0, 2, 1])).mul_scalar(scale);
377 let causal = make_causal_mask_t(seq_len);
378 let scores = scores + causal;
379 let scores = scores + blk.attn_vbias_t.clone();
380 let attn = activation::softmax(scores, 2);
381 let out = attn.matmul(v.clone());
382 let out: Tensor<2> = out.permute([1, 0, 2]).reshape([seq_len, d]);
383 Ok((burn_linear_nobias(out, blk.attn_o_w.clone()), k, v))
384 }
385
386 fn decode_encoder(
387 &self,
388 mut h: Tensor<2>,
389 new_time_ids: &[usize],
390 kv_cache: &mut Vec<(Tensor<3>, Tensor<3>)>,
391 cached_len: usize,
392 ) -> Result<Tensor<2>> {
393 for (li, blk) in self.blocks.iter().enumerate() {
394 let (h_out, k_new, v_new) =
395 self.decode_block_kv(h, blk, &kv_cache[li], new_time_ids, cached_len)?;
396 h = h_out;
397 kv_cache[li] = (k_new, v_new);
398 }
399 Ok(h)
400 }
401
402 fn decode_block_kv(
403 &self,
404 h: Tensor<2>,
405 blk: &BurnEncoderBlock,
406 cache: &(Tensor<3>, Tensor<3>),
407 new_time_ids: &[usize],
408 cached_len: usize,
409 ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
410 let h_norm = burn_rms_norm_nd(h.clone(), blk.norm1_w.clone(), 1e-6);
411 let (attn_out, k_new, v_new) =
412 self.decode_attn_kv(&h_norm, blk, cache, new_time_ids, cached_len)?;
413 let h = h + attn_out;
414 let h_norm2 = burn_rms_norm_nd(h.clone(), blk.norm2_w.clone(), 1e-6);
415 let ffn_out = swiglu_ffn(
416 h_norm2,
417 blk.ffn_fc1_w.clone(),
418 blk.ffn_fc2_w.clone(),
419 blk.ffn_gate_w.clone(),
420 );
421 Ok((h + ffn_out, k_new, v_new))
422 }
423
424 fn decode_attn_kv(
425 &self,
426 h: &Tensor<2>,
427 blk: &BurnEncoderBlock,
428 cache: &(Tensor<3>, Tensor<3>),
429 new_time_ids: &[usize],
430 cached_len: usize,
431 ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
432 let cfg = &self.config;
433 let new_len = new_time_ids.len();
434 let d = cfg.d_model;
435 let qkv = burn_linear_nobias(h.clone(), blk.attn_qkv_w.clone());
436 let q: Tensor<2> = qkv.clone().narrow(1, 0, d);
437 let k: Tensor<2> = qkv.clone().narrow(1, d, d);
438 let v: Tensor<2> = qkv.narrow(1, 2 * d, d);
439 let q: Tensor<3> = q.reshape([new_len, cfg.n_heads, cfg.head_dim]);
440 let k: Tensor<3> = k.reshape([new_len, cfg.n_heads, cfg.head_dim]);
441 let q = qk_norm_heads(q, blk.attn_qn_w.clone(), new_len, cfg.n_heads, cfg.head_dim);
442 let k = qk_norm_heads(k, blk.attn_kn_w.clone(), new_len, cfg.n_heads, cfg.head_dim);
443 let q: Tensor<3> = q.permute([1, 0, 2]);
444 let k: Tensor<3> = k.permute([1, 0, 2]);
445 let v: Tensor<3> = v
446 .reshape([new_len, cfg.n_heads, cfg.head_dim])
447 .permute([1, 0, 2]);
448 let q = self.apply_partial_rope(q, new_time_ids, cfg.n_heads, cfg.head_dim)?;
449 let k = self.apply_partial_rope(k, new_time_ids, cfg.n_heads, cfg.head_dim)?;
450 let k_full = Tensor::cat(vec![cache.0.clone(), k], 1);
451 let v_full = Tensor::cat(vec![cache.1.clone(), v], 1);
452 let scale = 1.0f32 / (cfg.head_dim as f32).sqrt();
453 let scores = q
454 .matmul(k_full.clone().permute([0, 2, 1]))
455 .mul_scalar(scale);
456 let scores = add_decode_causal_mask(scores, new_len, cached_len);
457 let scores = scores + blk.attn_vbias_t.clone();
458 let attn = activation::softmax(scores, 2);
459 let out = attn.matmul(v_full.clone());
460 let out: Tensor<2> = out.permute([1, 0, 2]).reshape([new_len, d]);
461 Ok((
462 burn_linear_nobias(out, blk.attn_o_w.clone()),
463 k_full,
464 v_full,
465 ))
466 }
467}
468
469#[cfg(test)]
470mod tests {
471 use super::*;
472 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
473
474 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
475 let mut x = seed;
476 (0..n)
477 .map(|_| {
478 x = x
479 .wrapping_mul(6364136223846793005)
480 .wrapping_add(1442695040888963407);
481 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
482 })
483 .collect()
484 }
485
486 fn tiny_config() -> Moirai2Config {
487 Moirai2Config {
488 d_model: 16,
489 n_layers: 1,
490 n_heads: 2,
491 head_dim: 8,
492 d_ff: 32,
493 patch_size: 2,
494 num_predict_token: 2,
495 num_quantiles: 3,
496 max_seq_len: 16,
497 rope_dim: 4,
498 median_quantile: 1,
499 }
500 }
501
502 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
503 *seed += 1;
504 let data: Vec<f32> = pseudo(rows * cols, *seed)
505 .iter()
506 .map(|v| v * scale)
507 .collect();
508 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
509 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
510 }
511
512 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
513 *seed += 1;
514 let bytes: Vec<u8> = pseudo(n, *seed)
515 .iter()
516 .flat_map(|v| v.to_le_bytes())
517 .collect();
518 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
519 }
520
521 fn put_res(
522 w: &mut GGUFWriter,
523 prefix: &str,
524 h: usize,
525 d_in: usize,
526 o: usize,
527 seed: &mut u64,
528 s: f32,
529 ) {
530 put2(w, &format!("{prefix}.hidden.weight"), h, d_in, seed, s);
531 put1(w, &format!("{prefix}.hidden.bias"), h, seed);
532 put2(w, &format!("{prefix}.output.weight"), o, h, seed, s);
533 put1(w, &format!("{prefix}.output.bias"), o, seed);
534 put2(w, &format!("{prefix}.residual.weight"), o, d_in, seed, s);
535 put1(w, &format!("{prefix}.residual.bias"), o, seed);
536 }
537
538 #[test]
539 fn burn_forecast_golden() {
540 let dir = std::env::temp_dir().join(format!("moirai2-burn-{}", std::process::id()));
542 let _ = std::fs::remove_dir_all(&dir);
543 std::fs::create_dir_all(&dir).unwrap();
544 let path = dir.join("model.gguf");
545 let mut w = GGUFWriter::new();
546 w.add_metadata(
547 "general.architecture",
548 GGUFMetaValue::String("moirai2".into()),
549 );
550 let mut seed = 1500u64;
551 let s = 0.1;
552 put_res(&mut w, "in_proj", 8, 4, 16, &mut seed, s);
553 put2(&mut w, "blk.0.attn_q.weight", 16, 16, &mut seed, s);
554 put2(&mut w, "blk.0.attn_k.weight", 16, 16, &mut seed, s);
555 put2(&mut w, "blk.0.attn_v.weight", 16, 16, &mut seed, s);
556 put1(&mut w, "blk.0.norm1.weight", 16, &mut seed);
557 put1(&mut w, "blk.0.norm2.weight", 16, &mut seed);
558 put2(&mut w, "blk.0.attn_o.weight", 16, 16, &mut seed, s);
559 put1(&mut w, "blk.0.attn_qn.weight", 8, &mut seed);
560 put1(&mut w, "blk.0.attn_kn.weight", 8, &mut seed);
561 put1(&mut w, "blk.0.attn_vbias.weight", 4, &mut seed);
562 put2(&mut w, "blk.0.ffn_fc1.weight", 32, 16, &mut seed, s);
563 put2(&mut w, "blk.0.ffn_fc2.weight", 16, 32, &mut seed, s);
564 put2(&mut w, "blk.0.ffn_gate.weight", 32, 16, &mut seed, s);
565 put1(&mut w, "norm_f.weight", 16, &mut seed);
566 put_res(&mut w, "out_proj", 8, 16, 12, &mut seed, s);
567 let f = std::fs::File::create(&path).unwrap();
568 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
569
570 let cfg = tiny_config();
571 let burn = BurnMoirai2Model::load(&path, cfg).unwrap();
572 let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
574 let t0 = std::time::Instant::now();
575 let a = burn.forecast(&ctx, 5).unwrap();
576 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
577 let b = burn.forecast(&ctx, 5).unwrap();
578 assert_eq!(a.len(), 5);
579 assert!(a.iter().all(|v| v.is_finite()));
580 assert!(a.iter().any(|&v| v != 0.0));
581 assert_eq!(a, b, "forecast must be deterministic");
582 println!("moirai2 synthetic: burn {burn_ms:.1}ms out={a:?}");
583 let golden: [f32; 5] = [9.501931, 10.167264, 10.172378, 9.644132, 9.523021];
585 assert_eq!(a.len(), golden.len());
586 for (x, y) in a.iter().zip(golden.iter()) {
587 assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
588 }
589 let _ = std::fs::remove_dir_all(&dir);
590 }
591}