zsfm_tirex/infer/
mod.rs

1//! TiRex inference on the Burn backend (Flex CPU) for the two tensor
2//! stages, with all host kernels reused verbatim from `super`.
3//!
4//! TiRex is already ~90% engine-agnostic host code (sLSTM recurrence,
5//! patching, scaling, SIMD kernels). Only `residual_block_forward` and
6//! `ffn_forward` touch tensors, so this port re-implements just those two
7//! on Burn and keeps every host helper in this module.
8//! Weight loading reads GGUF straight via [`zsfm_burn::GgufWeights`].
9
10use std::path::Path;
11
12use anyhow::{Context, Result};
13use burn::tensor::{activation, Device, Tensor, TensorData};
14use simdeez::math::{SimdMathF32Core, SimdMathF32Hyperbolic};
15use simdeez::prelude::*;
16use zsfm_burn::linear::burn_linear_nd;
17use zsfm_burn::{open_gguf_weights, GgufWeights};
18
19use crate::config::TiRexConfig;
20
21// ---------------------------------------------------------------------------
22// SIMD kernels (cross-platform: SSE2 / AVX2 / AVX-512 / NEON / WASM / scalar)
23// ---------------------------------------------------------------------------
24
25simd_runtime_generate!(
26    fn simd_sq_sum(row: &[f32]) -> f32 {
27        let mut r = &row[..];
28        let mut acc = S::Vf32::zeroes();
29        while r.len() >= S::Vf32::WIDTH {
30            let v = S::Vf32::load_from_slice(r);
31            acc = v.mul_add(v, acc);
32            r = &r[S::Vf32::WIDTH..];
33        }
34        let mut sum = acc.horizontal_add();
35        for &x in r {
36            sum += x * x;
37        }
38        sum
39    }
40);
41
42simd_runtime_generate!(
43    fn simd_scale_weight(row: &mut [f32], w: &[f32], scale: f32) {
44        let sc = S::Vf32::set1(scale);
45        let mut r = &mut row[..];
46        let mut wv = &w[..];
47        while r.len() >= S::Vf32::WIDTH {
48            let v = S::Vf32::load_from_slice(r);
49            let wi = S::Vf32::load_from_slice(wv);
50            (v * sc * wi).copy_to_slice(r);
51            r = &mut r[S::Vf32::WIDTH..];
52            wv = &wv[S::Vf32::WIDTH..];
53        }
54        for i in 0..r.len() {
55            r[i] *= scale * wv[i];
56        }
57    }
58);
59
60simd_runtime_generate!(
61    fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
62        let mut aa = &a[..];
63        let mut bb = &b[..];
64        let mut acc = S::Vf32::zeroes();
65        while aa.len() >= S::Vf32::WIDTH {
66            let va = S::Vf32::load_from_slice(aa);
67            let vb = S::Vf32::load_from_slice(bb);
68            acc = va.mul_add(vb, acc);
69            aa = &aa[S::Vf32::WIDTH..];
70            bb = &bb[S::Vf32::WIDTH..];
71        }
72        let mut sum = acc.horizontal_add();
73        for (&x, &y) in aa.iter().zip(bb.iter()) {
74            sum += x * y;
75        }
76        sum
77    }
78);
79
80simd_runtime_generate!(
81    fn simd_gate_update(
82        iraw: &[f32],
83        fraw: &[f32],
84        zraw: &[f32],
85        oraw: &[f32],
86        c: &[f32],
87        n: &[f32],
88        m: &[f32],
89        cnew: &mut [f32],
90        nnew: &mut [f32],
91        hnew: &mut [f32],
92        mnew: &mut [f32],
93        is_first: bool,
94    ) {
95        let clamp = S::Vf32::set1(15.0f32);
96        let one = S::Vf32::set1(1.0f32);
97        let eps = S::Vf32::set1(1e-8f32);
98        let zero = S::Vf32::zeroes();
99
100        let mut ir = &iraw[..];
101        let mut fr = &fraw[..];
102        let mut zr = &zraw[..];
103        let mut or_ = &oraw[..];
104        let mut cv = &c[..];
105        let mut nv = &n[..];
106        let mut mv = &m[..];
107        let mut cnw = &mut cnew[..];
108        let mut nnw = &mut nnew[..];
109        let mut hnw = &mut hnew[..];
110        let mut mnw = &mut mnew[..];
111
112        while ir.len() >= S::Vf32::WIDTH {
113            let iv = S::Vf32::load_from_slice(ir);
114            let fv = S::Vf32::load_from_slice(fr).min(clamp);
115            let zv = S::Vf32::load_from_slice(zr);
116            let ov = S::Vf32::load_from_slice(or_);
117            let cv_ = S::Vf32::load_from_slice(cv);
118            let nv_ = S::Vf32::load_from_slice(nv);
119            let mv_ = S::Vf32::load_from_slice(mv);
120
121            // log_sigmoid(fv): stable two-branch form, select by sign
122            let ls_pos = -(one + (-fv).exp_u35()).ln_u35(); // fv >= 0: -ln(1+exp(-fv))
123            let ls_neg = fv - (one + fv.exp_u35()).ln_u35(); // fv < 0:  fv - ln(1+exp(fv))
124            let log_sig_f = fv.cmp_lt(zero).blendv(ls_pos, ls_neg);
125
126            let logfplusm = mv_ + log_sig_f;
127            let mnew_v = if is_first { iv } else { iv.max(logfplusm) };
128
129            let ogate = one / (one + (-ov).exp_u35()); // sigmoid(ov)
130            let igate = (iv - mnew_v).exp_u35().min(one);
131            let fgate = (logfplusm - mnew_v).exp_u35().min(one);
132            let zgate = zv.tanh_u35();
133
134            let cnew_v = fgate.mul_add(cv_, igate * zgate);
135            let nnew_v = fgate.mul_add(nv_, igate);
136
137            let mask = nnew_v.abs().cmp_gt(eps);
138            let hnew_v = mask.blendv(zero, ogate * cnew_v / nnew_v);
139
140            cnew_v.copy_to_slice(cnw);
141            nnew_v.copy_to_slice(nnw);
142            hnew_v.copy_to_slice(hnw);
143            mnew_v.copy_to_slice(mnw);
144
145            ir = &ir[S::Vf32::WIDTH..];
146            fr = &fr[S::Vf32::WIDTH..];
147            zr = &zr[S::Vf32::WIDTH..];
148            or_ = &or_[S::Vf32::WIDTH..];
149            cv = &cv[S::Vf32::WIDTH..];
150            nv = &nv[S::Vf32::WIDTH..];
151            mv = &mv[S::Vf32::WIDTH..];
152            cnw = &mut cnw[S::Vf32::WIDTH..];
153            nnw = &mut nnw[S::Vf32::WIDTH..];
154            hnw = &mut hnw[S::Vf32::WIDTH..];
155            mnw = &mut mnw[S::Vf32::WIDTH..];
156        }
157
158        for i in 0..ir.len() {
159            let iv_s = ir[i];
160            let fv_s = fr[i].min(15.0);
161            let zv_s = zr[i];
162            let ov_s = or_[i];
163            let c_s = cv[i];
164            let n_s = nv[i];
165            let m_s = mv[i];
166            let ls = if fv_s >= 0.0 {
167                -(1.0 + (-fv_s).exp()).ln()
168            } else {
169                fv_s - (1.0 + fv_s.exp()).ln()
170            };
171            let lfpm = m_s + ls;
172            let mn = if is_first { iv_s } else { iv_s.max(lfpm) };
173            let og = 1.0 / (1.0 + (-ov_s).exp());
174            let ig = (iv_s - mn).exp().min(1.0);
175            let fg = (lfpm - mn).exp().min(1.0);
176            let zg = zv_s.tanh();
177            cnw[i] = fg * c_s + ig * zg;
178            nnw[i] = fg * n_s + ig;
179            hnw[i] = if nnw[i].abs() > 1e-8 {
180                og * cnw[i] / nnw[i]
181            } else {
182                0.0
183            };
184            mnw[i] = mn;
185        }
186    }
187);
188
189/// StandardScaler: compute (loc, scale) ignoring NaN values.
190fn standard_scaler(x: &[f32]) -> (f32, f32) {
191    let eps = 1e-5f32;
192    let valid: Vec<f32> = x.iter().copied().filter(|v| !v.is_nan()).collect();
193    if valid.is_empty() {
194        return (0.0, 1.0);
195    }
196    let loc = valid.iter().sum::<f32>() / valid.len() as f32;
197    let variance = valid.iter().map(|&v| (v - loc) * (v - loc)).sum::<f32>() / valid.len() as f32;
198    let scale = variance.sqrt();
199    let scale = if scale == 0.0 { loc.abs() + eps } else { scale };
200    (loc, scale)
201}
202
203/// Pad (left with NaN) or truncate (from front) context to target_len.
204fn adjust_context(x: &[f32], target_len: usize) -> Vec<f32> {
205    if x.len() >= target_len {
206        x[x.len() - target_len..].to_vec()
207    } else {
208        let pad = target_len - x.len();
209        let mut out = vec![f32::NAN; target_len];
210        out[pad..].copy_from_slice(x);
211        out
212    }
213}
214
215/// RMSNorm in-place: x = x / rms(x) * w
216fn rms_norm_inplace(x: &mut [f32], w: &[f32], d: usize, eps: f32) {
217    let s = x.len() / d;
218    for t in 0..s {
219        let row = &mut x[t * d..(t + 1) * d];
220        let ss = simd_sq_sum(row);
221        let scale = 1.0 / (ss / d as f32 + eps).sqrt();
222        simd_scale_weight(row, w, scale);
223    }
224}
225
226/// Fused headwise linear for ng gate groups: [S, D] → [S, ng*D] written into `out`.
227///
228/// w layout: [NH, ng*DH, DH] — for head h, gate g: rows g*DH..(g+1)*DH hold w[h,g,:].
229/// Output layout: out[t, g, h, o] at flat index t*ng*d + g*d + h*dh + o.
230/// This matches the x_g layout expected by the sLSTM recurrence, so no interleaving
231/// copy is needed after calling this function.
232fn headwise_linear_batch_ng(
233    x: &[f32],
234    w: &[f32],
235    s: usize,
236    nh: usize,
237    dh: usize,
238    ng: usize,
239    out: &mut [f32],
240) {
241    let d = nh * dh;
242    let ngdh = ng * dh;
243    for t in 0..s {
244        for h in 0..nh {
245            let x_h = &x[t * d + h * dh..t * d + (h + 1) * dh];
246            let w_h_off = h * ngdh * dh;
247            for g in 0..ng {
248                let out_base = t * ng * d + g * d + h * dh;
249                let w_gate_off = w_h_off + g * dh * dh;
250                for o in 0..dh {
251                    let w_row = &w[w_gate_off + o * dh..w_gate_off + (o + 1) * dh];
252                    out[out_base + o] = simd_dot(x_h, w_row);
253                }
254            }
255        }
256    }
257}
258
259/// Compute recurrent contribution Ry from h_prev and the sLSTM kernel.
260///
261/// kernel_t layout: [NH, NG*DH, DH] (transposed at load time from [NH, DH, NG*DH])
262/// so the inner dot-product dimension (di) is contiguous, enabling simd_dot.
263/// Output written into `out`: [NG*NH*DH = 2048] in [NG, NH, DH] order.
264fn compute_ry(
265    h: &[f32],
266    kernel_t: &[f32],
267    nh: usize,
268    dh: usize,
269    ng: usize,
270    ry_raw: &mut [f32],
271    out: &mut [f32],
272) {
273    // ry_raw[NH, NG*DH]: for each head, dot h_head with each row of kernel_t[head]
274    for head in 0..nh {
275        let h_h = &h[head * dh..(head + 1) * dh];
276        let k_head = &kernel_t[head * ng * dh * dh..];
277        for gate_d in 0..ng * dh {
278            let k_row = &k_head[gate_d * dh..(gate_d + 1) * dh]; // contiguous — simd_dot applies
279            ry_raw[head * ng * dh + gate_d] = simd_dot(h_h, k_row);
280        }
281    }
282    // Permute [NH, NG, DH] → [NG, NH, DH]
283    for head in 0..nh {
284        for g in 0..ng {
285            for d in 0..dh {
286                out[g * nh * dh + head * dh + d] = ry_raw[head * ng * dh + g * dh + d];
287            }
288        }
289    }
290}
291
292fn device() -> Device {
293    Device::flex()
294}
295
296fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
297    weights.load_2d(name)
298}
299
300fn load_vec(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
301    weights.load_flat(name)
302}
303
304fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
305    t.to_data()
306        .try_to_vec()
307        .map_err(|e| anyhow::anyhow!("{e:?}"))
308}
309
310struct BurnEmbedBlock {
311    hidden_w: Tensor<2>,
312    hidden_b: Vec<f32>,
313    output_w: Tensor<2>,
314    output_b: Vec<f32>,
315    residual_w: Tensor<2>,
316    residual_b: Vec<f32>,
317}
318
319struct BurnBlock {
320    norm_slstm: Vec<f32>,
321    fizo_w: Vec<f32>,
322    slstm_kernel_t: Vec<f32>,
323    slstm_bias: Vec<f32>,
324    group_norm_w: Vec<f32>,
325    norm_ffn: Vec<f32>,
326    ffn_gate_w: Tensor<2>,
327    ffn_up_w: Tensor<2>,
328    ffn_down_w: Tensor<2>,
329}
330
331pub struct BurnTiRexModel {
332    config: TiRexConfig,
333    in_emb: BurnEmbedBlock,
334    blocks: Vec<BurnBlock>,
335    out_norm: Vec<f32>,
336    out_emb: BurnEmbedBlock,
337}
338
339fn load_embed_block(weights: &GgufWeights, prefix: &str) -> Result<BurnEmbedBlock> {
340    let p = |s: &str| format!("{prefix}.{s}");
341    Ok(BurnEmbedBlock {
342        hidden_w: load_t(weights, &p("hidden.weight"))?,
343        hidden_b: load_vec(weights, &p("hidden.bias"))?,
344        output_w: load_t(weights, &p("output.weight"))?,
345        output_b: load_vec(weights, &p("output.bias"))?,
346        residual_w: load_t(weights, &p("residual.weight"))?,
347        residual_b: load_vec(weights, &p("residual.bias"))?,
348    })
349}
350
351impl BurnTiRexModel {
352    pub fn load(gguf_path: &Path, config: TiRexConfig) -> Result<Self> {
353        let weights = open_gguf_weights(gguf_path)
354            .with_context(|| format!("open {}", gguf_path.display()))?;
355
356        let in_emb = load_embed_block(&weights, "in_emb")?;
357        let nh = config.num_heads;
358        let dh = config.head_dim();
359        let wpp = dh * dh;
360        let mut blocks = Vec::with_capacity(config.num_blocks);
361        for n in 0..config.num_blocks {
362            let p = |s: &str| format!("blk.{n}.{s}");
363            let fgate_w = load_vec(&weights, &p("fgate.weight"))?;
364            let igate_w = load_vec(&weights, &p("igate.weight"))?;
365            let zgate_w = load_vec(&weights, &p("zgate.weight"))?;
366            let ogate_w = load_vec(&weights, &p("ogate.weight"))?;
367            let mut fizo_w = vec![0.0f32; nh * 4 * wpp];
368            for h in 0..nh {
369                fizo_w[h * 4 * wpp..h * 4 * wpp + wpp]
370                    .copy_from_slice(&fgate_w[h * wpp..(h + 1) * wpp]);
371                fizo_w[h * 4 * wpp + wpp..h * 4 * wpp + 2 * wpp]
372                    .copy_from_slice(&igate_w[h * wpp..(h + 1) * wpp]);
373                fizo_w[h * 4 * wpp + 2 * wpp..h * 4 * wpp + 3 * wpp]
374                    .copy_from_slice(&zgate_w[h * wpp..(h + 1) * wpp]);
375                fizo_w[h * 4 * wpp + 3 * wpp..h * 4 * wpp + 4 * wpp]
376                    .copy_from_slice(&ogate_w[h * wpp..(h + 1) * wpp]);
377            }
378            let raw_kernel = load_vec(&weights, &p("slstm_kernel"))?;
379            let ng = 4usize;
380            let mut slstm_kernel_t = vec![0.0f32; nh * ng * dh * dh];
381            for head in 0..nh {
382                for gate_d in 0..(ng * dh) {
383                    for di in 0..dh {
384                        slstm_kernel_t[head * ng * dh * dh + gate_d * dh + di] =
385                            raw_kernel[head * dh * ng * dh + di * ng * dh + gate_d];
386                    }
387                }
388            }
389            blocks.push(BurnBlock {
390                norm_slstm: load_vec(&weights, &p("norm_slstm"))?,
391                fizo_w,
392                slstm_kernel_t,
393                slstm_bias: load_vec(&weights, &p("slstm_bias"))?,
394                group_norm_w: load_vec(&weights, &p("group_norm"))?,
395                norm_ffn: load_vec(&weights, &p("norm_ffn"))?,
396                ffn_gate_w: load_t(&weights, &p("ffn_gate.weight"))?,
397                ffn_up_w: load_t(&weights, &p("ffn_up.weight"))?,
398                ffn_down_w: load_t(&weights, &p("ffn_down.weight"))?,
399            });
400        }
401        let out_norm = load_vec(&weights, "out_norm")?;
402        let out_emb = load_embed_block(&weights, "out_emb")?;
403        Ok(Self {
404            config,
405            in_emb,
406            blocks,
407            out_norm,
408            out_emb,
409        })
410    }
411}
412
413/// Burn residual MLP block over host row-major data.
414fn residual_block_forward_burn(
415    x: &[f32],
416    s: usize,
417    in_dim: usize,
418    h_dim: usize,
419    out_dim: usize,
420    emb: &BurnEmbedBlock,
421) -> Vec<f32> {
422    let dev = device();
423    let xt = Tensor::<2>::from_data(TensorData::new(x.to_vec(), [s, in_dim]), &dev);
424    let bh = Tensor::<1>::from_data(TensorData::new(emb.hidden_b.clone(), [h_dim]), &dev);
425    let h = activation::relu(burn_linear_nd(xt.clone(), emb.hidden_w.clone(), Some(bh)));
426    let bo = Tensor::<1>::from_data(TensorData::new(emb.output_b.clone(), [out_dim]), &dev);
427    let out = burn_linear_nd(h, emb.output_w.clone(), Some(bo));
428    let br = Tensor::<1>::from_data(TensorData::new(emb.residual_b.clone(), [out_dim]), &dev);
429    let res = burn_linear_nd(xt, emb.residual_w.clone(), Some(br));
430    let y: Vec<f32> = (out + res).to_data().try_to_vec().unwrap();
431    y
432}
433
434/// Burn SiLU-gated FFN over host row-major data.
435fn ffn_forward_burn(
436    x: &[f32],
437    s: usize,
438    in_dim: usize,
439    up_dim: usize,
440    blk: &BurnBlock,
441) -> Vec<f32> {
442    let dev = device();
443    let xt = Tensor::<2>::from_data(TensorData::new(x.to_vec(), [s, in_dim]), &dev);
444    let gate = burn_linear_nd(xt.clone(), blk.ffn_gate_w.clone(), None);
445    let up = burn_linear_nd(xt, blk.ffn_up_w.clone(), None);
446    let h = activation::silu(gate) * up;
447    let out = burn_linear_nd(h, blk.ffn_down_w.clone(), None);
448    out.to_data().try_to_vec().unwrap()
449}
450
451impl BurnTiRexModel {
452    pub fn forecast(
453        &self,
454        context: &[f32],
455        prediction_length: usize,
456    ) -> Result<(Vec<Vec<f32>>, Vec<f32>)> {
457        let cfg = &self.config;
458        let patch_size = cfg.patch_size;
459        let n_quantiles = cfg.num_quantiles;
460        let median_idx = cfg
461            .quantiles
462            .iter()
463            .position(|&q| (q - 0.5).abs() < 1e-6)
464            .unwrap_or(4);
465        let n_steps = prediction_length.div_ceil(patch_size);
466        let mut all_q: Vec<f32> = Vec::with_capacity(n_steps * patch_size * n_quantiles);
467        let mut ctx: Vec<f32> = context.to_vec();
468        for _ in 0..n_steps {
469            let full_ctx = adjust_context(&ctx, cfg.train_ctx_len);
470            let (loc, scale) = standard_scaler(&full_ctx);
471            let s_vals: Vec<f32> = full_ctx
472                .iter()
473                .map(|&x| {
474                    if x.is_nan() {
475                        0.0f32
476                    } else {
477                        (x - loc) / scale
478                    }
479                })
480                .collect();
481            let s_mask: Vec<f32> = full_ctx
482                .iter()
483                .map(|&x| if x.is_nan() { 0.0f32 } else { 1.0f32 })
484                .collect();
485            let num_patches = cfg.train_ctx_len / patch_size;
486            let mut patched_vals = vec![0.0f32; num_patches * patch_size];
487            let mut patched_mask = vec![0.0f32; num_patches * patch_size];
488            for p in 0..num_patches {
489                let start = p * patch_size;
490                patched_vals[p * patch_size..(p + 1) * patch_size]
491                    .copy_from_slice(&s_vals[start..start + patch_size]);
492                patched_mask[p * patch_size..(p + 1) * patch_size]
493                    .copy_from_slice(&s_mask[start..start + patch_size]);
494            }
495            let in_dim = patch_size * 2;
496            let mut x_in = vec![0.0f32; num_patches * in_dim];
497            for p in 0..num_patches {
498                x_in[p * in_dim..p * in_dim + patch_size]
499                    .copy_from_slice(&patched_vals[p * patch_size..(p + 1) * patch_size]);
500                x_in[p * in_dim + patch_size..p * in_dim + in_dim]
501                    .copy_from_slice(&patched_mask[p * patch_size..(p + 1) * patch_size]);
502            }
503            let mut hidden = residual_block_forward_burn(
504                &x_in,
505                num_patches,
506                in_dim,
507                cfg.input_ff_dim,
508                cfg.embedding_dim,
509                &self.in_emb,
510            );
511            let ng = 4usize;
512            let d = cfg.embedding_dim;
513            let mut sc_xg = vec![0.0f32; num_patches * ng * d];
514            let mut sc_hout = vec![0.0f32; num_patches * d];
515            let mut sc_y = vec![0.0f32; num_patches * d];
516            let mut sc_xn = vec![0.0f32; num_patches * d];
517            let mut sc_raw = vec![0.0f32; ng * d];
518            let mut sc_ry_raw = vec![0.0f32; ng * d];
519            let mut sc_ry_out = vec![0.0f32; ng * d];
520            let mut sc_hnew = vec![0.0f32; d];
521            let mut sc_cnew = vec![0.0f32; d];
522            let mut sc_nnew = vec![0.0f32; d];
523            let mut sc_mnew = vec![0.0f32; d];
524            for block in &self.blocks {
525                hidden = self.forward_slstm_block(
526                    &hidden,
527                    num_patches,
528                    block,
529                    &mut sc_xg,
530                    &mut sc_hout,
531                    &mut sc_y,
532                    &mut sc_xn,
533                    &mut sc_raw,
534                    &mut sc_ry_raw,
535                    &mut sc_ry_out,
536                    &mut sc_hnew,
537                    &mut sc_cnew,
538                    &mut sc_nnew,
539                    &mut sc_mnew,
540                )?;
541            }
542            rms_norm_inplace(&mut hidden, &self.out_norm, cfg.embedding_dim, 1e-6);
543            let out_dim = cfg.output_dim();
544            let preds = residual_block_forward_burn(
545                &hidden,
546                num_patches,
547                cfg.embedding_dim,
548                cfg.input_ff_dim,
549                out_dim,
550                &self.out_emb,
551            );
552            let last = &preds[(num_patches - 1) * out_dim..num_patches * out_dim];
553            for q in 0..n_quantiles {
554                for d in 0..patch_size {
555                    let v = last[q * patch_size + d] * scale + loc;
556                    all_q.push(v);
557                }
558            }
559            ctx.extend(std::iter::repeat(f32::NAN).take(patch_size));
560        }
561        let total = n_steps * patch_size;
562        let mut quantiles: Vec<Vec<f32>> = vec![Vec::with_capacity(prediction_length); n_quantiles];
563        for t in 0..prediction_length.min(total) {
564            let step = t / patch_size;
565            let d = t % patch_size;
566            for q in 0..n_quantiles {
567                quantiles[q].push(all_q[step * n_quantiles * patch_size + q * patch_size + d]);
568            }
569        }
570        let median: Vec<f32> = (0..prediction_length)
571            .map(|t| quantiles[median_idx][t])
572            .collect();
573        Ok((quantiles, median))
574    }
575
576    #[allow(clippy::too_many_arguments)]
577    fn forward_slstm_block(
578        &self,
579        x: &[f32],
580        s: usize,
581        blk: &BurnBlock,
582        sc_xg: &mut [f32],
583        sc_hout: &mut [f32],
584        sc_y: &mut [f32],
585        sc_xn: &mut [f32],
586        sc_raw: &mut [f32],
587        sc_ry_raw: &mut [f32],
588        sc_ry_out: &mut [f32],
589        sc_hnew: &mut Vec<f32>,
590        sc_cnew: &mut Vec<f32>,
591        sc_nnew: &mut Vec<f32>,
592        sc_mnew: &mut Vec<f32>,
593    ) -> Result<Vec<f32>> {
594        let cfg = &self.config;
595        let d = cfg.embedding_dim;
596        let nh = cfg.num_heads;
597        let dh = cfg.head_dim();
598        let ng = 4usize;
599        sc_xn[..s * d].copy_from_slice(&x[..s * d]);
600        rms_norm_inplace(&mut sc_xn[..s * d], &blk.norm_slstm, d, 1e-6);
601        headwise_linear_batch_ng(&sc_xn[..s * d], &blk.fizo_w, s, nh, dh, ng, sc_xg);
602        let mut h = vec![0.0f32; d];
603        let mut c = vec![0.0f32; d];
604        let mut n = vec![0.0f32; d];
605        let mut m = vec![f32::NEG_INFINITY; d];
606        for t in 0..s {
607            let wx = &sc_xg[t * ng * d..(t + 1) * ng * d];
608            compute_ry(&h, &blk.slstm_kernel_t, nh, dh, ng, sc_ry_raw, sc_ry_out);
609            for i in 0..ng * d {
610                sc_raw[i] = wx[i] + sc_ry_out[i] + blk.slstm_bias[i];
611            }
612            let is_first = t == 0;
613            simd_gate_update(
614                &sc_raw[..d],
615                &sc_raw[d..2 * d],
616                &sc_raw[2 * d..3 * d],
617                &sc_raw[3 * d..],
618                &c,
619                &n,
620                &m,
621                sc_cnew,
622                sc_nnew,
623                sc_hnew,
624                sc_mnew,
625                is_first,
626            );
627            std::mem::swap(&mut h, sc_hnew);
628            std::mem::swap(&mut c, sc_cnew);
629            std::mem::swap(&mut n, sc_nnew);
630            std::mem::swap(&mut m, sc_mnew);
631            sc_hout[t * d..(t + 1) * d].copy_from_slice(&h);
632        }
633        for t in 0..s {
634            let h_t = &sc_hout[t * d..(t + 1) * d];
635            for head in 0..nh {
636                let start = head * dh;
637                let h_slice = &h_t[start..start + dh];
638                let mean = h_slice.iter().sum::<f32>() / dh as f32;
639                let var = h_slice
640                    .iter()
641                    .map(|&v| (v - mean) * (v - mean))
642                    .sum::<f32>()
643                    / dh as f32;
644                let inv_std = 1.0 / (var + 1e-5f32).sqrt();
645                let w_slice = &blk.group_norm_w[start..start + dh];
646                for d_i in 0..dh {
647                    sc_y[t * d + start + d_i] =
648                        (h_t[start + d_i] - mean) * inv_std * (1.0 + w_slice[d_i]);
649                }
650            }
651        }
652        let mut x_out = x.to_vec();
653        for i in 0..s * d {
654            x_out[i] += sc_y[i];
655        }
656        sc_xn[..s * d].copy_from_slice(&x_out[..s * d]);
657        rms_norm_inplace(&mut sc_xn[..s * d], &blk.norm_ffn, d, 1e-6);
658        let ffn_out = ffn_forward_burn(&sc_xn[..s * d], s, d, cfg.ffn_up_dim, blk);
659        for i in 0..s * d {
660            x_out[i] += ffn_out[i];
661        }
662        Ok(x_out)
663    }
664}
665
666#[cfg(test)]
667mod tests {
668    use super::*;
669    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
670
671    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
672        let mut x = seed;
673        (0..n)
674            .map(|_| {
675                x = x
676                    .wrapping_mul(6364136223846793005)
677                    .wrapping_add(1442695040888963407);
678                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
679            })
680            .collect()
681    }
682
683    fn tiny_config() -> TiRexConfig {
684        TiRexConfig {
685            patch_size: 4,
686            num_blocks: 1,
687            embedding_dim: 16,
688            num_heads: 2,
689            input_ff_dim: 32,
690            ffn_up_dim: 24,
691            train_ctx_len: 16,
692            quantiles: vec![0.25, 0.5, 0.75],
693            num_quantiles: 3,
694        }
695    }
696
697    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
698        *seed += 1;
699        let data: Vec<f32> = pseudo(rows * cols, *seed)
700            .iter()
701            .map(|v| v * scale)
702            .collect();
703        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
704        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
705    }
706
707    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
708        *seed += 1;
709        let bytes: Vec<u8> = pseudo(n, *seed)
710            .iter()
711            .flat_map(|v| v.to_le_bytes())
712            .collect();
713        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
714    }
715
716    fn put_emb(
717        w: &mut GGUFWriter,
718        prefix: &str,
719        h: usize,
720        d_in: usize,
721        o: usize,
722        seed: &mut u64,
723        s: f32,
724    ) {
725        put2(w, &format!("{prefix}.hidden.weight"), h, d_in, seed, s);
726        put1(w, &format!("{prefix}.hidden.bias"), h, seed);
727        put2(w, &format!("{prefix}.output.weight"), o, h, seed, s);
728        put1(w, &format!("{prefix}.output.bias"), o, seed);
729        put2(w, &format!("{prefix}.residual.weight"), o, d_in, seed, s);
730        put1(w, &format!("{prefix}.residual.bias"), o, seed);
731    }
732
733    #[test]
734    fn burn_forecast_golden() {
735        let dir = std::env::temp_dir().join(format!("tirex-burn-{}", std::process::id()));
736        let _ = std::fs::remove_dir_all(&dir);
737        std::fs::create_dir_all(&dir).unwrap();
738        let path = dir.join("model.gguf");
739        let mut w = GGUFWriter::new();
740        w.add_metadata(
741            "general.architecture",
742            GGUFMetaValue::String("tirex".into()),
743        );
744        let mut seed = 2900u64;
745        let s = 0.05;
746        put_emb(&mut w, "in_emb", 32, 8, 16, &mut seed, s);
747        for t in ["fgate", "igate", "zgate", "ogate"] {
748            put2(&mut w, &format!("blk.0.{t}.weight"), 16, 16, &mut seed, s);
749        }
750        // slstm_kernel logical [NH, DH, NG*DH] = [2,8,32].
751        put1(&mut w, "blk.0.slstm_kernel", 2 * 8 * 32, &mut seed);
752        // NOTE: raw kernel values unscaled (structure, not numerics, matters here).
753        put1(&mut w, "blk.0.slstm_bias", 64, &mut seed);
754        put1(&mut w, "blk.0.norm_slstm", 16, &mut seed);
755        put1(&mut w, "blk.0.group_norm", 16, &mut seed);
756        put1(&mut w, "blk.0.norm_ffn", 16, &mut seed);
757        put2(&mut w, "blk.0.ffn_gate.weight", 24, 16, &mut seed, s);
758        put2(&mut w, "blk.0.ffn_up.weight", 24, 16, &mut seed, s);
759        put2(&mut w, "blk.0.ffn_down.weight", 16, 24, &mut seed, s);
760        put1(&mut w, "out_norm", 16, &mut seed);
761        put_emb(&mut w, "out_emb", 32, 16, 12, &mut seed, s);
762        let f = std::fs::File::create(&path).unwrap();
763        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
764
765        let cfg = tiny_config();
766        let burn = BurnTiRexModel::load(&path, cfg).unwrap();
767        let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
768        let t0 = std::time::Instant::now();
769        let (aq, am) = burn.forecast(&ctx, 2).unwrap();
770        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
771        let (bq, bm) = burn.forecast(&ctx, 2).unwrap();
772        assert_eq!(aq, bq, "quantiles must be deterministic");
773        assert_eq!(am, bm, "median must be deterministic");
774        assert!(aq.iter().all(|r| r.iter().all(|v| v.is_finite())));
775        assert!(am.iter().all(|v| v.is_finite()));
776        println!("tirex synthetic: burn {burn_ms:.1}ms q={aq:?} m={am:?}");
777        // Golden recorded from the verified Burn port (ex-candle parity 0.00e0).
778        let golden_q: Vec<Vec<f32>> = vec![
779            vec![10.033624, 10.955177],
780            vec![10.487212, 9.479774],
781            vec![10.795641, 10.196359],
782        ];
783        let golden_m: Vec<f32> = vec![10.487212, 9.479774];
784        assert_eq!(aq.len(), golden_q.len());
785        assert_eq!(am.len(), golden_m.len());
786        for (ra, re) in aq.iter().zip(golden_q.iter()) {
787            for (x, y) in ra.iter().zip(re.iter()) {
788                assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
789            }
790        }
791        for (x, y) in am.iter().zip(golden_m.iter()) {
792            assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
793        }
794        let _ = std::fs::remove_dir_all(&dir);
795    }
796}
zsfm_tirex/infer/mod.rs

zsfm_tirex/infer/
mod.rs
