zsfm_lag_llama/infer/
mod.rs

1//! Lag-Llama inference on the Burn backend (Flex CPU).
2//!
3//! Two-phase inference:
4//! 1. Prefill: full forward pass over context tokens on Burn tensors,
5//!    collects K/V cache per layer.
6//! 2. Decode: per-step single-token forward pass with pure-Rust host
7//!    kernels (`rms_norm_raw`, `raw_gemv`, `rope_single_inplace`,
8//!    `mha_decode_raw`, `silu_mlp_raw`).
9//!
10//! Weights load straight from GGUF via [`zsfm_burn::GgufWeights`].
11
12use std::path::Path;
13
14use anyhow::{Context, Result};
15use burn::tensor::{activation, Device, Tensor, TensorData};
16use simdeez::prelude::*;
17use zsfm_burn::linear::{burn_linear_nd, burn_linear_nobias};
18use zsfm_burn::norm::burn_rms_norm_nd;
19use zsfm_burn::{open_gguf_weights, GgufWeights};
20
21use crate::config::LagLlamaConfig;
22
23fn device() -> Device {
24    Device::flex()
25}
26
27fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
28    weights.load_2d(name)
29}
30
31fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
32    weights.load_1d(name)
33}
34
35fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
36    t.to_data()
37        .try_to_vec()
38        .map_err(|e| anyhow::anyhow!("{e:?}"))
39}
40
41struct BurnTransformerBlock {
42    rms1_w: Tensor<1>,
43    rms2_w: Tensor<1>,
44    qkv_w: Tensor<2>,
45    c_w: Tensor<2>,
46    fc1_w: Tensor<2>,
47    fc2_w: Tensor<2>,
48    proj_w: Tensor<2>,
49}
50
51struct RawBlockWeights {
52    rms1: Vec<f32>,
53    rms2: Vec<f32>,
54    qkv: Vec<f32>,
55    c: Vec<f32>,
56    fc1: Vec<f32>,
57    fc2: Vec<f32>,
58    proj: Vec<f32>,
59}
60
61pub struct BurnLagLlamaModel {
62    config: LagLlamaConfig,
63    rope_cos: Tensor<2>,
64    rope_sin: Tensor<2>,
65    wte_w: Tensor<2>,
66    wte_b: Tensor<1>,
67    blocks: Vec<BurnTransformerBlock>,
68    raw_blocks: Vec<RawBlockWeights>,
69    rope_cos_raw: Vec<f32>,
70    rope_sin_raw: Vec<f32>,
71    norm_f_raw: Vec<f32>,
72    wte_w_raw: Vec<f32>,
73    wte_b_raw: Vec<f32>,
74    mu_w_raw: Vec<f32>,
75    mu_b_raw: Vec<f32>,
76}
77
78impl BurnLagLlamaModel {
79    pub fn load(gguf_path: &Path, config: LagLlamaConfig) -> Result<Self> {
80        let weights = open_gguf_weights(gguf_path)
81            .with_context(|| format!("open {}", gguf_path.display()))?;
82
83        let wte_w = load_t(&weights, "enc.wte.weight")?;
84        let wte_b = load_v(&weights, "enc.wte.bias")?;
85        let mut blocks = Vec::with_capacity(config.n_layer);
86        for n in 0..config.n_layer {
87            let p = |s: &str| format!("blk.{n}.{s}");
88            let qw = load_t(&weights, &p("attn_q.weight"))?;
89            let kvw = load_t(&weights, &p("attn_kv.weight"))?;
90            // kv holds K and V stacked: check total size.
91            let qkv_w = Tensor::cat(vec![qw, kvw], 0);
92            blocks.push(BurnTransformerBlock {
93                rms1_w: load_v(&weights, &p("rms1.weight"))?,
94                rms2_w: load_v(&weights, &p("rms2.weight"))?,
95                qkv_w,
96                c_w: load_t(&weights, &p("attn_c.weight"))?,
97                fc1_w: load_t(&weights, &p("mlp_fc1.weight"))?,
98                fc2_w: load_t(&weights, &p("mlp_fc2.weight"))?,
99                proj_w: load_t(&weights, &p("mlp_proj.weight"))?,
100            });
101        }
102
103        // norm_f is 1D in the checkpoint; load as vector.
104        let norm_f_raw: Vec<f32> = weights.load_flat("norm_f.weight")?;
105        let mu_w_t = load_t(&weights, "head.mu.weight")?;
106        let mu_b_raw: Vec<f32> = weights.load_flat("head.mu.bias")?;
107        let mu_w_raw: Vec<f32> = to_host(&mu_w_t)?;
108        let wte_w_raw: Vec<f32> = to_host(&wte_w)?;
109        let wte_b_raw: Vec<f32> = to_host(&wte_b)?;
110
111        let head_dim = config.n_embd_per_head;
112        let half = head_dim / 2;
113        let inv_freq: Vec<f32> = (0..half)
114            .map(|i| 1.0_f32 / 10000_f32.powf(2.0 * i as f32 / head_dim as f32))
115            .collect();
116        let max_pos = config.max_context_length + 4096;
117        let mut cos_vals = vec![0.0f32; max_pos * half];
118        let mut sin_vals = vec![0.0f32; max_pos * half];
119        for p in 0..max_pos {
120            let pos = p as f32;
121            for i in 0..half {
122                let theta = pos * inv_freq[i];
123                cos_vals[p * half + i] = theta.cos();
124                sin_vals[p * half + i] = theta.sin();
125            }
126        }
127        let dev = device();
128        let rope_cos =
129            Tensor::<2>::from_data(TensorData::new(cos_vals.clone(), [max_pos, half]), &dev);
130        let rope_sin =
131            Tensor::<2>::from_data(TensorData::new(sin_vals.clone(), [max_pos, half]), &dev);
132
133        let mut raw_blocks = Vec::with_capacity(config.n_layer);
134        for blk in &blocks {
135            raw_blocks.push(RawBlockWeights {
136                rms1: to_host(&blk.rms1_w)?,
137                rms2: to_host(&blk.rms2_w)?,
138                qkv: to_host(&blk.qkv_w)?,
139                c: to_host(&blk.c_w)?,
140                fc1: to_host(&blk.fc1_w)?,
141                fc2: to_host(&blk.fc2_w)?,
142                proj: to_host(&blk.proj_w)?,
143            });
144        }
145
146        Ok(Self {
147            config,
148            rope_cos,
149            rope_sin,
150            wte_w,
151            wte_b: wte_b.clone(),
152            blocks,
153            raw_blocks,
154            rope_cos_raw: cos_vals,
155            rope_sin_raw: sin_vals,
156            norm_f_raw,
157            wte_w_raw,
158            wte_b_raw,
159            mu_w_raw,
160            mu_b_raw,
161        })
162    }
163
164    pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
165        let cfg = &self.config;
166        let max_lag = *cfg.lags_seq.iter().max().unwrap_or(&0);
167        let (loc, scale) = robust_stats(context);
168        let scale = scale.max(1e-8);
169        let mut hist: Vec<f32> = vec![0.0; max_lag + 1];
170        for &v in context {
171            hist.push((v - loc) / scale);
172        }
173        let ctx_buf_end = hist.len();
174        let ctx_buf_start = ctx_buf_end.saturating_sub(cfg.max_context_length);
175        let seq_len = ctx_buf_end - ctx_buf_start;
176
177        // --- Phase 1: Burn prefill (full sequence, once) ---
178        let feat_ctx = build_feature_matrix(&hist, ctx_buf_start, seq_len, cfg);
179        let dev = device();
180        let x =
181            Tensor::<2>::from_data(TensorData::new(feat_ctx, [seq_len, cfg.feature_size]), &dev);
182        let mut h = burn_linear_nd(x, self.wte_w.clone(), Some(self.wte_b.clone()));
183        let mut kv_caches: Vec<(Tensor<3>, Tensor<3>)> = Vec::with_capacity(cfg.n_layer);
184        for blk in &self.blocks {
185            let (h_out, k, v) = self.prefill_block(&h, blk, seq_len, ctx_buf_start)?;
186            h = h_out;
187            kv_caches.push((k, v));
188        }
189        let h_flat: Vec<f32> = to_host(&h)?;
190        let n_embd = cfg.n_embd;
191        let mut h_raw: Vec<f32> = h_flat[(seq_len - 1) * n_embd..seq_len * n_embd].to_vec();
192
193        let n_head = cfg.n_head;
194        let head_dim = cfg.n_embd_per_head;
195        let mlp_hidden = cfg.mlp_hidden;
196        let feat_size = cfg.feature_size;
197        let mut kv_raw = extract_kv_caches_raw(&kv_caches, n_head, head_dim, horizon)?;
198        let mut kv_len = seq_len;
199
200        // --- Phase 2: verbatim pure-Rust decode loop ---
201        let max_kv_len = seq_len + horizon;
202        let mut h_tmp = vec![0.0f32; n_embd];
203        let mut qkv_buf = vec![0.0f32; 3 * n_embd];
204        let mut proj_buf = vec![0.0f32; n_embd];
205        let mut mlp_gate = vec![0.0f32; mlp_hidden];
206        let mut mlp_up = vec![0.0f32; mlp_hidden];
207        let mut mlp_out = vec![0.0f32; n_embd];
208        let mut attn_out = vec![0.0f32; n_embd];
209        let mut scores_scratch = vec![0.0f32; n_head * max_kv_len];
210        let mut preds = Vec::with_capacity(horizon);
211        let mut rope_offset = ctx_buf_start + seq_len;
212
213        for step in 0..horizon {
214            h_tmp.copy_from_slice(&h_raw);
215            rms_norm_raw(&mut h_tmp, &self.norm_f_raw, 1e-5);
216            let pred_scaled = raw_dot(&h_tmp, &self.mu_w_raw) + self.mu_b_raw[0];
217            preds.push(pred_scaled * scale + loc);
218            if step == horizon - 1 {
219                break;
220            }
221            hist.push(pred_scaled);
222            let abs_t = hist.len() - 1;
223            let feat_one = build_one_feature(&hist, abs_t, cfg);
224            raw_gemv_bias(
225                &feat_one,
226                &self.wte_w_raw,
227                &self.wte_b_raw,
228                n_embd,
229                feat_size,
230                &mut h_raw,
231            );
232            let new_kv_len = kv_len + 1;
233            for li in 0..cfg.n_layer {
234                let blk = &self.raw_blocks[li];
235                let (ref mut k_heads, ref mut v_heads) = kv_raw[li];
236                h_tmp.copy_from_slice(&h_raw);
237                rms_norm_raw(&mut h_tmp, &blk.rms1, 1e-5);
238                raw_gemv(&h_tmp, &blk.qkv, 3 * n_embd, n_embd, &mut qkv_buf);
239                rope_single_inplace(
240                    &mut qkv_buf[..n_embd],
241                    rope_offset,
242                    &self.rope_cos_raw,
243                    &self.rope_sin_raw,
244                    n_head,
245                    head_dim,
246                );
247                rope_single_inplace(
248                    &mut qkv_buf[n_embd..2 * n_embd],
249                    rope_offset,
250                    &self.rope_cos_raw,
251                    &self.rope_sin_raw,
252                    n_head,
253                    head_dim,
254                );
255                for hi in 0..n_head {
256                    k_heads[hi].extend_from_slice(
257                        &qkv_buf[n_embd + hi * head_dim..n_embd + (hi + 1) * head_dim],
258                    );
259                    v_heads[hi].extend_from_slice(
260                        &qkv_buf[2 * n_embd + hi * head_dim..2 * n_embd + (hi + 1) * head_dim],
261                    );
262                }
263                mha_decode_raw(
264                    &qkv_buf[..n_embd],
265                    k_heads,
266                    v_heads,
267                    n_head,
268                    head_dim,
269                    new_kv_len,
270                    &mut scores_scratch,
271                    &mut attn_out,
272                );
273                raw_gemv(&attn_out, &blk.c, n_embd, n_embd, &mut proj_buf);
274                for i in 0..n_embd {
275                    h_raw[i] += proj_buf[i];
276                }
277                h_tmp.copy_from_slice(&h_raw);
278                rms_norm_raw(&mut h_tmp, &blk.rms2, 1e-5);
279                silu_mlp_raw(
280                    &h_tmp,
281                    &blk.fc1,
282                    &blk.fc2,
283                    &blk.proj,
284                    n_embd,
285                    mlp_hidden,
286                    &mut mlp_gate,
287                    &mut mlp_up,
288                    &mut mlp_out,
289                );
290                for i in 0..n_embd {
291                    h_raw[i] += mlp_out[i];
292                }
293            }
294            kv_len = new_kv_len;
295            rope_offset += 1;
296        }
297        Ok(preds)
298    }
299
300    fn prefill_block(
301        &self,
302        hidden: &Tensor<2>,
303        blk: &BurnTransformerBlock,
304        seq_len: usize,
305        rope_start: usize,
306    ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
307        let h = burn_rms_norm_nd(hidden.clone(), blk.rms1_w.clone(), 1e-5);
308        let (attn_out, k, v) = self.prefill_attn(&h, blk, seq_len, rope_start)?;
309        let h = hidden.clone() + attn_out;
310        let h2 = burn_rms_norm_nd(h.clone(), blk.rms2_w.clone(), 1e-5);
311        let h2 = silu_mlp_burn(&h2, &blk.fc1_w, &blk.fc2_w, &blk.proj_w);
312        Ok((h2 + h, k, v))
313    }
314
315    fn prefill_attn(
316        &self,
317        hidden: &Tensor<2>,
318        blk: &BurnTransformerBlock,
319        seq_len: usize,
320        rope_start: usize,
321    ) -> Result<(Tensor<2>, Tensor<3>, Tensor<3>)> {
322        let cfg = &self.config;
323        let n_head = cfg.n_head;
324        let head_dim = cfg.n_embd_per_head;
325        let n_embd = cfg.n_embd;
326        let qkv = burn_linear_nobias(hidden.clone(), blk.qkv_w.clone());
327        let q: Tensor<2> = qkv.clone().narrow(1, 0, n_embd);
328        let k: Tensor<2> = qkv.clone().narrow(1, n_embd, n_embd);
329        let v: Tensor<2> = qkv.narrow(1, 2 * n_embd, n_embd);
330        let q: Tensor<3> = q.reshape([seq_len, n_head, head_dim]).permute([1, 0, 2]);
331        let k: Tensor<3> = k.reshape([seq_len, n_head, head_dim]).permute([1, 0, 2]);
332        let v: Tensor<3> = v.reshape([seq_len, n_head, head_dim]).permute([1, 0, 2]);
333        let q = apply_rope_3d(&q, rope_start, seq_len, &self.rope_cos, &self.rope_sin);
334        let k = apply_rope_3d(&k, rope_start, seq_len, &self.rope_cos, &self.rope_sin);
335        let scale = 1.0f32 / (head_dim as f32).sqrt();
336        let scores = q.matmul(k.clone().permute([0, 2, 1])).mul_scalar(scale);
337        let mask = burn_causal_mask(seq_len);
338        let scores = scores + mask;
339        let attn = activation::softmax(scores, 2);
340        let out = attn.matmul(v.clone());
341        let out: Tensor<2> = out.permute([1, 0, 2]).reshape([seq_len, n_embd]);
342        let out = burn_linear_nobias(out, blk.c_w.clone());
343        Ok((out, k, v))
344    }
345}
346
347fn apply_rope_3d(
348    x: &Tensor<3>,
349    offset: usize,
350    seq_len: usize,
351    cos_table: &Tensor<2>,
352    sin_table: &Tensor<2>,
353) -> Tensor<3> {
354    let half = cos_table.dims()[1];
355    let cos_t = cos_table
356        .clone()
357        .narrow(0, offset, seq_len)
358        .unsqueeze_dim::<3>(0);
359    let sin_t = sin_table
360        .clone()
361        .narrow(0, offset, seq_len)
362        .unsqueeze_dim::<3>(0);
363    let x1 = x.clone().narrow(2, 0, half);
364    let x2 = x.clone().narrow(2, half, half);
365    let rot1 = x1.clone() * cos_t.clone() - x2.clone() * sin_t.clone();
366    let rot2 = x1 * sin_t + x2 * cos_t;
367    Tensor::cat(vec![rot1, rot2], 2)
368}
369
370fn burn_causal_mask(seq_len: usize) -> Tensor<3> {
371    let dev = device();
372    let mut mask_data = vec![0.0f32; seq_len * seq_len];
373    for i in 0..seq_len {
374        for j in (i + 1)..seq_len {
375            mask_data[i * seq_len + j] = f32::NEG_INFINITY;
376        }
377    }
378    Tensor::<2>::from_data(TensorData::new(mask_data, [seq_len, seq_len]), &dev)
379        .unsqueeze_dim::<3>(0)
380}
381
382fn silu_mlp_burn(
383    x: &Tensor<2>,
384    fc1_w: &Tensor<2>,
385    fc2_w: &Tensor<2>,
386    proj_w: &Tensor<2>,
387) -> Tensor<2> {
388    let h = activation::silu(burn_linear_nobias(x.clone(), fc1_w.clone()));
389    let h = h * burn_linear_nobias(x.clone(), fc2_w.clone());
390    burn_linear_nobias(h, proj_w.clone())
391}
392
393fn extract_kv_caches_raw(
394    kv_caches: &[(Tensor<3>, Tensor<3>)],
395    n_head: usize,
396    head_dim: usize,
397    horizon: usize,
398) -> Result<Vec<(Vec<Vec<f32>>, Vec<Vec<f32>>)>> {
399    let mut result = Vec::with_capacity(kv_caches.len());
400    for (k_t, v_t) in kv_caches {
401        let k_flat: Vec<f32> = to_host(k_t)?;
402        let v_flat: Vec<f32> = to_host(v_t)?;
403        let seq = k_flat.len() / (n_head * head_dim);
404        let cap = (seq + horizon) * head_dim;
405        let mut k_heads: Vec<Vec<f32>> = Vec::with_capacity(n_head);
406        let mut v_heads: Vec<Vec<f32>> = Vec::with_capacity(n_head);
407        for h in 0..n_head {
408            // Layout is [H, S, D] row-major: head h occupies a contiguous chunk.
409            let start = h * seq * head_dim;
410            let end = start + seq * head_dim;
411            let mut kh = Vec::with_capacity(cap);
412            kh.extend_from_slice(&k_flat[start..end]);
413            let mut vh = Vec::with_capacity(cap);
414            vh.extend_from_slice(&v_flat[start..end]);
415            k_heads.push(kh);
416            v_heads.push(vh);
417        }
418        result.push((k_heads, v_heads));
419    }
420    Ok(result)
421}
422
423fn build_feature_matrix(
424    hist: &[f32],
425    start: usize,
426    seq_len: usize,
427    cfg: &LagLlamaConfig,
428) -> Vec<f32> {
429    let mut feat = vec![0.0f32; seq_len * cfg.feature_size];
430    for t in 0..seq_len {
431        let abs_t = start + t;
432        for (li, &lag) in cfg.lags_seq.iter().enumerate() {
433            let src = abs_t as isize - lag as isize;
434            if src >= 0 && (src as usize) < hist.len() {
435                feat[t * cfg.feature_size + li] = hist[src as usize];
436            }
437        }
438    }
439    feat
440}
441
442fn build_one_feature(hist: &[f32], abs_t: usize, cfg: &LagLlamaConfig) -> Vec<f32> {
443    let mut feat = vec![0.0f32; cfg.feature_size];
444    for (li, &lag) in cfg.lags_seq.iter().enumerate() {
445        let src = abs_t as isize - lag as isize;
446        if src >= 0 && (src as usize) < hist.len() {
447            feat[li] = hist[src as usize];
448        }
449    }
450    feat
451}
452
453fn robust_stats(x: &[f32]) -> (f32, f32) {
454    let n = x.len();
455    if n == 0 {
456        return (0.0, 1.0);
457    }
458    let mut sorted = x.to_vec();
459    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
460    let median = if n % 2 == 0 {
461        (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
462    } else {
463        sorted[n / 2]
464    };
465    let mut devs: Vec<f32> = sorted.iter().map(|&v| (v - median).abs()).collect();
466    devs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
467    let mad = if n % 2 == 0 {
468        (devs[n / 2 - 1] + devs[n / 2]) / 2.0
469    } else {
470        devs[n / 2]
471    };
472    (median, mad)
473}
474
475simd_runtime_generate!(
476    fn simd_sq_sum(row: &[f32]) -> f32 {
477        let mut r = &row[..];
478        let mut acc = S::Vf32::zeroes();
479        while r.len() >= S::Vf32::WIDTH {
480            let v = S::Vf32::load_from_slice(r);
481            acc = v.mul_add(v, acc);
482            r = &r[S::Vf32::WIDTH..];
483        }
484        let mut sum = acc.horizontal_add();
485        for &x in r {
486            sum += x * x;
487        }
488        sum
489    }
490);
491
492simd_runtime_generate!(
493    fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
494        let mut aa = &a[..];
495        let mut bb = &b[..];
496        let mut acc = S::Vf32::zeroes();
497        while aa.len() >= S::Vf32::WIDTH {
498            let va = S::Vf32::load_from_slice(aa);
499            let vb = S::Vf32::load_from_slice(bb);
500            acc = va.mul_add(vb, acc);
501            aa = &aa[S::Vf32::WIDTH..];
502            bb = &bb[S::Vf32::WIDTH..];
503        }
504        let mut sum = acc.horizontal_add();
505        for (&x, &y) in aa.iter().zip(bb.iter()) {
506            sum += x * y;
507        }
508        sum
509    }
510);
511
512#[inline(always)]
513fn fast_exp_f32(x: f32) -> f32 {
514    let x = x.max(-87.3365_f32);
515    f32::from_bits(((x * 12102203.0_f32) as i32 + 1064866805_i32) as u32)
516}
517
518#[inline(always)]
519fn raw_dot(a: &[f32], b: &[f32]) -> f32 {
520    simd_dot(a, b)
521}
522
523fn rms_norm_raw(x: &mut [f32], w: &[f32], eps: f32) {
524    let n = x.len() as f32;
525    let rms = (simd_sq_sum(x) / n + eps).sqrt();
526    for i in 0..x.len() {
527        x[i] = x[i] / rms * w[i];
528    }
529}
530
531fn raw_gemv(x: &[f32], w: &[f32], n_out: usize, n_in: usize, out: &mut [f32]) {
532    for i in 0..n_out {
533        out[i] = raw_dot(x, &w[i * n_in..(i + 1) * n_in]);
534    }
535}
536
537fn raw_gemv_bias(x: &[f32], w: &[f32], b: &[f32], n_out: usize, n_in: usize, out: &mut [f32]) {
538    for i in 0..n_out {
539        out[i] = raw_dot(x, &w[i * n_in..(i + 1) * n_in]) + b[i];
540    }
541}
542
543fn rope_single_inplace(
544    qk: &mut [f32],
545    pos: usize,
546    cos: &[f32],
547    sin: &[f32],
548    n_head: usize,
549    head_dim: usize,
550) {
551    let half = head_dim / 2;
552    let cos_row = &cos[pos * half..(pos + 1) * half];
553    let sin_row = &sin[pos * half..(pos + 1) * half];
554    for h in 0..n_head {
555        let base = h * head_dim;
556        for i in 0..half {
557            let x1 = qk[base + i];
558            let x2 = qk[base + half + i];
559            qk[base + i] = x1 * cos_row[i] - x2 * sin_row[i];
560            qk[base + half + i] = x1 * sin_row[i] + x2 * cos_row[i];
561        }
562    }
563}
564
565fn mha_decode_raw(
566    q: &[f32],
567    k_heads: &[Vec<f32>],
568    v_heads: &[Vec<f32>],
569    n_head: usize,
570    head_dim: usize,
571    kv_len: usize,
572    scores_scratch: &mut [f32],
573    out: &mut [f32],
574) {
575    let scale_inv = 1.0 / (head_dim as f32).sqrt();
576    for h in 0..n_head {
577        let q_h = &q[h * head_dim..(h + 1) * head_dim];
578        let k_h = &k_heads[h];
579        let v_h = &v_heads[h];
580        let sc = &mut scores_scratch[h * kv_len..(h + 1) * kv_len];
581        let k_ptr = k_h.as_ptr();
582        for j in 0..kv_len {
583            let kj = unsafe { std::slice::from_raw_parts(k_ptr.add(j * head_dim), head_dim) };
584            let mut dot = 0.0f32;
585            for d in 0..head_dim {
586                dot += q_h[d] * kj[d];
587            }
588            sc[j] = dot * scale_inv;
589        }
590        let max_s = sc.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
591        let mut sum = 0.0f32;
592        for s in sc.iter_mut() {
593            *s = fast_exp_f32(*s - max_s);
594            sum += *s;
595        }
596        let inv_sum = 1.0 / sum;
597        for s in sc.iter_mut() {
598            *s *= inv_sum;
599        }
600        let out_h = &mut out[h * head_dim..(h + 1) * head_dim];
601        out_h.iter_mut().for_each(|v| *v = 0.0);
602        let v_ptr = v_h.as_ptr();
603        for j in 0..kv_len {
604            let sc_j = sc[j];
605            let v_j = unsafe { std::slice::from_raw_parts(v_ptr.add(j * head_dim), head_dim) };
606            for d in 0..head_dim {
607                out_h[d] += sc_j * v_j[d];
608            }
609        }
610    }
611}
612
613fn silu_mlp_raw(
614    x: &[f32],
615    fc1: &[f32],
616    fc2: &[f32],
617    proj: &[f32],
618    n_embd: usize,
619    mlp_hidden: usize,
620    gate: &mut [f32],
621    up: &mut [f32],
622    out: &mut [f32],
623) {
624    for i in 0..mlp_hidden {
625        let v = raw_dot(x, &fc1[i * n_embd..(i + 1) * n_embd]);
626        gate[i] = v / (1.0 + fast_exp_f32(-v));
627    }
628    for i in 0..mlp_hidden {
629        up[i] = raw_dot(x, &fc2[i * n_embd..(i + 1) * n_embd]);
630    }
631    for j in 0..mlp_hidden {
632        gate[j] *= up[j];
633    }
634    raw_gemv(gate, proj, n_embd, mlp_hidden, out);
635}
636
637#[cfg(test)]
638mod tests {
639    use super::*;
640    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
641
642    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
643        let mut x = seed;
644        (0..n)
645            .map(|_| {
646                x = x
647                    .wrapping_mul(6364136223846793005)
648                    .wrapping_add(1442695040888963407);
649                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
650            })
651            .collect()
652    }
653
654    fn tiny_config() -> LagLlamaConfig {
655        LagLlamaConfig {
656            n_layer: 1,
657            n_head: 2,
658            n_embd_per_head: 8,
659            n_embd: 16,
660            mlp_hidden: 32,
661            feature_size: 4,
662            n_lags: 2,
663            n_time_feat: 2,
664            max_context_length: 16,
665            lags_seq: vec![1, 2],
666        }
667    }
668
669    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
670        *seed += 1;
671        let data: Vec<f32> = pseudo(rows * cols, *seed)
672            .iter()
673            .map(|v| v * scale)
674            .collect();
675        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
676        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
677    }
678
679    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
680        *seed += 1;
681        let bytes: Vec<u8> = pseudo(n, *seed)
682            .iter()
683            .flat_map(|v| v.to_le_bytes())
684            .collect();
685        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
686    }
687
688    #[test]
689    fn burn_forecast_golden() {
690        let dir = std::env::temp_dir().join(format!("lagllama-burn-{}", std::process::id()));
691        let _ = std::fs::remove_dir_all(&dir);
692        std::fs::create_dir_all(&dir).unwrap();
693        let path = dir.join("model.gguf");
694        let mut w = GGUFWriter::new();
695        w.add_metadata(
696            "general.architecture",
697            GGUFMetaValue::String("lag-llama".into()),
698        );
699        let mut seed = 2700u64;
700        let s = 0.1;
701        put2(&mut w, "enc.wte.weight", 16, 4, &mut seed, s);
702        put1(&mut w, "enc.wte.bias", 16, &mut seed);
703        put2(&mut w, "blk.0.attn_q.weight", 16, 16, &mut seed, s);
704        put2(&mut w, "blk.0.attn_kv.weight", 32, 16, &mut seed, s);
705        put1(&mut w, "blk.0.rms1.weight", 16, &mut seed);
706        put1(&mut w, "blk.0.rms2.weight", 16, &mut seed);
707        put2(&mut w, "blk.0.attn_c.weight", 16, 16, &mut seed, s);
708        put2(&mut w, "blk.0.mlp_fc1.weight", 32, 16, &mut seed, s);
709        put2(&mut w, "blk.0.mlp_fc2.weight", 32, 16, &mut seed, s);
710        put2(&mut w, "blk.0.mlp_proj.weight", 16, 32, &mut seed, s);
711        put1(&mut w, "norm_f.weight", 16, &mut seed);
712        put2(&mut w, "head.mu.weight", 1, 16, &mut seed, s);
713        put1(&mut w, "head.mu.bias", 1, &mut seed);
714        let f = std::fs::File::create(&path).unwrap();
715        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
716
717        let cfg = tiny_config();
718        let burn = BurnLagLlamaModel::load(&path, cfg).unwrap();
719        let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
720        let t0 = std::time::Instant::now();
721        let a = burn.forecast(&ctx, 2).unwrap();
722        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
723        let b = burn.forecast(&ctx, 2).unwrap();
724        assert_eq!(a.len(), 2);
725        assert!(a.iter().all(|v| v.is_finite()));
726        assert!(a.iter().any(|&v| v != 0.0));
727        assert_eq!(a, b, "forecast must be deterministic");
728        println!("lag-llama synthetic: burn {burn_ms:.1}ms out={a:?}");
729        // Golden recorded from the verified Burn port (ex-candle parity 0.00e0).
730        let golden: [f32; 2] = [10.582177, 10.570611];
731        assert_eq!(a.len(), golden.len());
732        for (x, y) in a.iter().zip(golden.iter()) {
733            assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
734        }
735        let _ = std::fs::remove_dir_all(&dir);
736    }
737}
zsfm_lag_llama/infer/mod.rs

zsfm_lag_llama/infer/
mod.rs
