zsfm_toto/infer/
mod.rs

1//! Toto-2 inference on the Burn backend (Flex CPU).
2//!
3//! Patched input scaler, transformer with GQA time/variate layers, xPos RoPE,
4//! PerDimScale, unit-scaled linears, SwiGLU FFN, and quantile output head,
5//! with weights loaded straight from GGUF via [`zsfm_burn::GgufWeights`].
6//!
7//! The `compute_f64` path is not ported: pass a config with
8//! `compute_f64 = false` (the default). The Burn port bails otherwise.
9
10use std::path::Path;
11
12use anyhow::{bail, Context, Result};
13use burn::tensor::{activation, Device, Tensor, TensorData};
14use zsfm_burn::linear::burn_linear_nd;
15use zsfm_burn::{open_gguf_weights, GgufWeights};
16use zsfm_core::{json_bool, json_f64, json_usize};
17
18const SILU_SCALE: f32 = 1.766782948312328;
19
20// ---------------------------------------------------------------------------
21// Config
22// ---------------------------------------------------------------------------
23
24#[derive(Clone, Debug)]
25pub struct InferConfig {
26    d_model: usize,
27    num_layers: usize,
28    num_heads: usize,
29    num_groups: usize,
30    qk_dim: usize,
31    v_dim: usize,
32    patch_size: usize,
33    norm_eps: f64,
34    layer_group_size: usize,
35    num_variate_layers_per_group: usize,
36    variate_layer_first: bool,
37    use_xpos: bool,
38    residual_mult: f64,
39    residual_attn_ratio: f64,
40    /// Run forward pass in F64 to match PyTorch numerical precision.
41    /// Weights are cast to F64 at load time (~2× memory). Default: false.
42    compute_f64: bool,
43}
44
45impl Default for InferConfig {
46    /// Matches the defaults of the upstream Datadog/Toto-2.0 `config.json`.
47    fn default() -> Self {
48        InferConfig {
49            d_model: 2048,
50            num_layers: 48,
51            num_heads: 32,
52            num_groups: 32,
53            qk_dim: 64,
54            v_dim: 64,
55            patch_size: 32,
56            norm_eps: 5e-4,
57            layer_group_size: 48,
58            num_variate_layers_per_group: 1,
59            variate_layer_first: false,
60            use_xpos: true,
61            residual_mult: 0.75,
62            residual_attn_ratio: 5.136215466577748,
63            compute_f64: false,
64        }
65    }
66}
67
68impl InferConfig {
69    fn is_variate_layer(&self, idx: usize) -> bool {
70        if self.variate_layer_first {
71            idx % self.layer_group_size < self.num_variate_layers_per_group
72        } else {
73            idx % self.layer_group_size >= self.layer_group_size - self.num_variate_layers_per_group
74        }
75    }
76    fn q_size(&self) -> usize {
77        self.qk_dim * self.num_heads
78    }
79    fn k_size(&self) -> usize {
80        self.qk_dim * self.num_groups
81    }
82    fn v_size(&self) -> usize {
83        self.v_dim * self.num_groups
84    }
85
86    // -- builder-style setters — chain from `InferConfig::default()` ---------
87
88    pub fn with_d_model(mut self, v: usize) -> Self {
89        self.d_model = v;
90        self
91    }
92    pub fn with_num_layers(mut self, v: usize) -> Self {
93        self.num_layers = v;
94        self
95    }
96    pub fn with_num_heads(mut self, v: usize) -> Self {
97        self.num_heads = v;
98        self
99    }
100    pub fn with_num_groups(mut self, v: usize) -> Self {
101        self.num_groups = v;
102        self
103    }
104    pub fn with_qk_dim(mut self, v: usize) -> Self {
105        self.qk_dim = v;
106        self
107    }
108    pub fn with_v_dim(mut self, v: usize) -> Self {
109        self.v_dim = v;
110        self
111    }
112    pub fn with_patch_size(mut self, v: usize) -> Self {
113        self.patch_size = v;
114        self
115    }
116    pub fn with_norm_eps(mut self, v: f64) -> Self {
117        self.norm_eps = v;
118        self
119    }
120    pub fn with_layer_group_size(mut self, v: usize) -> Self {
121        self.layer_group_size = v;
122        self
123    }
124    pub fn with_num_variate_layers_per_group(mut self, v: usize) -> Self {
125        self.num_variate_layers_per_group = v;
126        self
127    }
128    pub fn with_variate_layer_first(mut self, v: bool) -> Self {
129        self.variate_layer_first = v;
130        self
131    }
132    pub fn with_use_xpos(mut self, v: bool) -> Self {
133        self.use_xpos = v;
134        self
135    }
136    pub fn with_residual_mult(mut self, v: f64) -> Self {
137        self.residual_mult = v;
138        self
139    }
140    pub fn with_residual_attn_ratio(mut self, v: f64) -> Self {
141        self.residual_attn_ratio = v;
142        self
143    }
144
145    /// Run the forward pass in F64 (double precision) to match PyTorch numerical
146    /// accuracy. Uses ~2× memory. Default: F32.
147    pub fn with_compute_f64(mut self, v: bool) -> Self {
148        self.compute_f64 = v;
149        self
150    }
151
152    // -- read-only accessors needed outside this crate ------------------------
153
154    pub fn patch_size(&self) -> usize {
155        self.patch_size
156    }
157    pub fn compute_f64(&self) -> bool {
158        self.compute_f64
159    }
160
161    /// Overlay every recognised field from a parsed `config.json` onto
162    /// `InferConfig::default()`, using the same field names as the upstream
163    /// Datadog/Toto-2.0 config. Missing/unrecognised fields keep their default.
164    /// Shared by the CLI and (later) the Python bindings so the field-by-field
165    /// `.as_u64().unwrap_or(...)` boilerplate lives in exactly one place.
166    pub fn from_json_value(v: &serde_json::Value) -> Self {
167        Self::default()
168            .with_d_model(json_usize(v, "d_model", 2048))
169            .with_num_layers(json_usize(v, "num_layers", 48))
170            .with_num_heads(json_usize(v, "num_heads", 32))
171            .with_num_groups(json_usize(v, "num_groups", 32))
172            .with_qk_dim(json_usize(v, "qk_dim", 64))
173            .with_v_dim(json_usize(v, "v_dim", 64))
174            .with_patch_size(json_usize(v, "patch_size", 32))
175            .with_norm_eps(json_f64(v, "norm_eps", 5e-4))
176            .with_layer_group_size(json_usize(v, "layer_group_size", 48))
177            .with_num_variate_layers_per_group(json_usize(v, "num_variate_layers_per_group", 1))
178            .with_variate_layer_first(json_bool(v, "variate_layer_first", false))
179            .with_use_xpos(json_bool(v, "use_xpos", true))
180            .with_residual_mult(json_f64(v, "residual_mult", 0.75))
181            .with_residual_attn_ratio(json_f64(v, "residual_attn_ratio", 5.136215466577748))
182    }
183}
184
185fn device() -> Device {
186    Device::flex()
187}
188
189fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
190    weights.load_2d(name)
191}
192
193fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
194    weights.load_1d(name)
195}
196
197fn try_load_v(weights: &GgufWeights, name: &str) -> Result<Option<Tensor<1>>> {
198    match weights.load_1d(name) {
199        Ok(t) => Ok(Some(t)),
200        Err(_) => Ok(None),
201    }
202}
203
204fn load_w(weights: &GgufWeights, name: &str, expected_d_out: usize) -> Result<Tensor<2>> {
205    weights.load_weight_2d(name, expected_d_out)
206}
207
208/// Unit-scaled linear: `scale * (x @ w^T + b)` with `scale = 1/sqrt(d_in)`.
209fn uu_linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
210    let d_in = w.dims()[1] as f32;
211    burn_linear_nd(x, w, b).mul_scalar(1.0 / d_in.sqrt())
212}
213
214fn uu_linear_readout<const D: usize>(
215    x: Tensor<D>,
216    w: Tensor<2>,
217    b: Option<Tensor<1>>,
218) -> Tensor<D> {
219    let d_in = w.dims()[1] as f32;
220    burn_linear_nd(x, w, b).mul_scalar(1.0 / d_in)
221}
222
223fn uu_silu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
224    activation::silu(x).mul_scalar(SILU_SCALE)
225}
226
227fn rms_norm_weightless<const D: usize>(x: Tensor<D>, eps: f64) -> Tensor<D> {
228    let dims = x.dims();
229    let last = dims[D - 1];
230    let lead: usize = dims[..D - 1].iter().product();
231    let flat: Tensor<2> = x.reshape([lead, last]);
232    let mean_sq = flat.clone().square().mean_dim(1);
233    let rms = (mean_sq.add_scalar(eps as f32)).sqrt();
234    (flat / rms).reshape(dims)
235}
236
237fn residual_add<const D: usize>(h: Tensor<D>, skip: Tensor<D>, tau: f64) -> Tensor<D> {
238    let denom = (1.0 + tau * tau).sqrt() as f32;
239    h.mul_scalar(tau as f32 / denom) + skip.mul_scalar(1.0 / denom)
240}
241
242fn apply_per_dim_scale<const D: usize>(q: Tensor<D>, pds_w: Tensor<1>) -> Tensor<D> {
243    // softplus(pds) / ln2, broadcast over the last dim.
244    let sp = pds_w.exp() + 1.0;
245    let r = sp.log() / std::f32::consts::LN_2;
246    let dims = q.dims();
247    let last = dims[D - 1];
248    let lead: usize = dims[..D - 1].iter().product();
249    (q.reshape([lead, last]) * r.unsqueeze_dim::<2>(0)).reshape(dims)
250}
251
252struct BurnXposRope {
253    proj_width: usize,
254    cos: Vec<Vec<f32>>,
255    sin: Vec<Vec<f32>>,
256    xpos_base_scale: Vec<f32>,
257}
258
259impl BurnXposRope {
260    fn build(qk_dim: usize, max_len: usize) -> Self {
261        let proj_width = qk_dim / 2;
262        let half = proj_width / 2;
263        let base = 10000_f32;
264        let theta: Vec<f32> = (0..half)
265            .map(|i| 1.0 / base.powf(2.0 * i as f32 / proj_width as f32))
266            .collect();
267        let mut cos = vec![vec![0.0f32; proj_width]; max_len];
268        let mut sin = vec![vec![0.0f32; proj_width]; max_len];
269        for m in 0..max_len {
270            for i in 0..half {
271                let angle = m as f32 * theta[i];
272                cos[m][2 * i] = angle.cos();
273                cos[m][2 * i + 1] = angle.cos();
274                sin[m][2 * i] = angle.sin();
275                sin[m][2 * i + 1] = angle.sin();
276            }
277        }
278        let xpos_base_scale: Vec<f32> = (0..half)
279            .map(|i| ((2 * i) as f32 + 0.4 * proj_width as f32) / (1.4 * proj_width as f32))
280            .collect();
281        Self {
282            proj_width,
283            cos,
284            sin,
285            xpos_base_scale,
286        }
287    }
288
289    fn apply(&self, x: Tensor<4>, seq_ids: &[u32], xpos_exponent: f32) -> Tensor<4> {
290        let dims = x.dims();
291        let qk_dim = dims[3];
292        let proj_width = self.proj_width;
293        let half = proj_width / 2;
294        let seq_len = seq_ids.len();
295        let max_pos = seq_ids.iter().copied().max().unwrap_or(0) as f32;
296        let center = ((max_pos as u32 + 1) / 2) as f32;
297        let mut cos_data = vec![0.0f32; seq_len * proj_width];
298        let mut sin_data = vec![0.0f32; seq_len * proj_width];
299        for (si, &pos) in seq_ids.iter().enumerate() {
300            let power = (pos as f32 - center) / 256.0;
301            for i in 0..half {
302                let xpos_s = self.xpos_base_scale[i].powf(power).powf(xpos_exponent);
303                let c = self.cos[pos as usize][2 * i] * xpos_s;
304                let s = self.sin[pos as usize][2 * i] * xpos_s;
305                cos_data[si * proj_width + 2 * i] = c;
306                cos_data[si * proj_width + 2 * i + 1] = c;
307                sin_data[si * proj_width + 2 * i] = s;
308                sin_data[si * proj_width + 2 * i + 1] = s;
309            }
310        }
311        let dev = device();
312        let pos_cos =
313            Tensor::<4>::from_data(TensorData::new(cos_data, [1, 1, seq_len, proj_width]), &dev);
314        let pos_sin =
315            Tensor::<4>::from_data(TensorData::new(sin_data, [1, 1, seq_len, proj_width]), &dev);
316        // Interleaved rotate_half: [a0,b0,a1,b1..] -> [-b0,a0,-b1,a1..]
317        let x_rot = x.clone().narrow(3, 0, proj_width);
318        let shp = x_rot.dims();
319        let pairs = x_rot.clone().reshape([shp[0], shp[1], shp[2], half, 2]);
320        let x1 = pairs.clone().narrow(4, 0, 1);
321        let x2 = pairs.narrow(4, 1, 1);
322        let rot_x = Tensor::cat(vec![x2.mul_scalar(-1.0), x1], 4)
323            .reshape([shp[0], shp[1], shp[2], proj_width]);
324        let rotated = x_rot * pos_cos + rot_x * pos_sin;
325        if qk_dim > proj_width {
326            let x_pass = x.narrow(3, proj_width, qk_dim - proj_width);
327            Tensor::cat(vec![rotated, x_pass], 3)
328        } else {
329            rotated
330        }
331    }
332}
333
334struct BurnResidualMlpW {
335    l1_w: Tensor<2>,
336    l1_b: Tensor<1>,
337    l2_w: Tensor<2>,
338    l2_b: Tensor<1>,
339    skip_w: Tensor<2>,
340    skip_b: Tensor<1>,
341    tau: f64,
342    is_output: bool,
343}
344
345struct BurnBlockW {
346    attn_qkv_w: Tensor<2>,
347    attn_qkv_b: Option<Tensor<1>>,
348    attn_out_w: Tensor<2>,
349    attn_out_b: Option<Tensor<1>>,
350    attn_pds: Option<Tensor<1>>,
351    attn_tau: f64,
352    ffn_up_w: Tensor<2>,
353    ffn_down_w: Tensor<2>,
354    mlp_tau: f64,
355}
356
357pub struct BurnTotoModel {
358    config: InferConfig,
359    rope: BurnXposRope,
360    patch_proj: BurnResidualMlpW,
361    blocks: Vec<BurnBlockW>,
362    output_head: BurnResidualMlpW,
363}
364
365fn load_residual_mlp(
366    weights: &GgufWeights,
367    prefix: &str,
368    tau: f64,
369    is_output: bool,
370) -> Result<BurnResidualMlpW> {
371    Ok(BurnResidualMlpW {
372        l1_w: load_t(weights, &format!("{prefix}.linear1.weight"))?,
373        l1_b: load_v(weights, &format!("{prefix}.linear1.bias"))?,
374        l2_w: load_t(weights, &format!("{prefix}.linear2.weight"))?,
375        l2_b: load_v(weights, &format!("{prefix}.linear2.bias"))?,
376        skip_w: load_t(weights, &format!("{prefix}.skip_proj.weight"))?,
377        skip_b: load_v(weights, &format!("{prefix}.skip_proj.bias"))?,
378        tau,
379        is_output,
380    })
381}
382
383fn compute_taus(
384    num_layers: usize,
385    residual_mult: f64,
386    residual_attn_ratio: f64,
387) -> (Vec<f64>, Vec<f64>) {
388    let total_depth = 2 * num_layers;
389    let alpha_mlp = residual_mult * (2.0 / (1.0 + residual_attn_ratio.powi(2))).sqrt();
390    let alpha_attn = residual_attn_ratio * alpha_mlp;
391    let tau = |index: usize| -> f64 {
392        let n_attn = (index + 1) / 2;
393        let n_mlp = index / 2;
394        let num = if index % 2 == 0 {
395            alpha_attn
396        } else {
397            alpha_mlp
398        };
399        let den = (total_depth as f64 / 2.0
400            + n_attn as f64 * alpha_attn.powi(2)
401            + n_mlp as f64 * alpha_mlp.powi(2))
402        .sqrt();
403        num / den
404    };
405    (0..num_layers)
406        .map(|i| (tau(2 * i), tau(2 * i + 1)))
407        .unzip()
408}
409
410impl BurnTotoModel {
411    pub fn patch_size(&self) -> usize {
412        self.config.patch_size()
413    }
414
415    pub fn load(gguf_path: &Path, config: InferConfig) -> Result<Self> {
416        if config.compute_f64 {
417            bail!("BurnTotoModel supports compute_f64=false only");
418        }
419        let weights = open_gguf_weights(gguf_path)
420            .with_context(|| format!("open {}", gguf_path.display()))?;
421
422        let patch_proj = load_residual_mlp(&weights, "patch_proj", 1.0, false)?;
423
424        let (attn_taus, mlp_taus) = compute_taus(
425            config.num_layers,
426            config.residual_mult,
427            config.residual_attn_ratio,
428        );
429        let mut blocks = Vec::with_capacity(config.num_layers);
430        for n in 0..config.num_layers {
431            blocks.push(BurnBlockW {
432                attn_qkv_w: load_t(&weights, &format!("blk.{n}.attn_qkv.weight"))?,
433                attn_qkv_b: try_load_v(&weights, &format!("blk.{n}.attn_qkv.bias"))?,
434                attn_out_w: load_t(&weights, &format!("blk.{n}.attn_output.weight"))?,
435                attn_out_b: try_load_v(&weights, &format!("blk.{n}.attn_output.bias"))?,
436                attn_pds: try_load_v(&weights, &format!("blk.{n}.attn_pds.weight"))?,
437                attn_tau: attn_taus[n],
438                ffn_up_w: load_t(&weights, &format!("blk.{n}.ffn_up.weight"))?,
439                // Mirror the Q8 realignment the loader applies to ffn_down.
440                ffn_down_w: load_w(
441                    &weights,
442                    &format!("blk.{n}.ffn_down.weight"),
443                    config.d_model,
444                )?,
445                mlp_tau: mlp_taus[n],
446            });
447        }
448        let output_head = load_residual_mlp(&weights, "output_head", 1.0, true)?;
449        let rope = BurnXposRope::build(config.qk_dim, 8192);
450        Ok(Self {
451            config,
452            rope,
453            patch_proj,
454            blocks,
455            output_head,
456        })
457    }
458
459    pub fn forecast(
460        &self,
461        target: &[Vec<f32>],
462        mask: &[Vec<bool>],
463        prediction_length: usize,
464    ) -> Result<Vec<Vec<Vec<f32>>>> {
465        let n_var = target.len();
466        let ctx_len = target[0].len();
467        let patch_size = self.config.patch_size;
468        if ctx_len % patch_size != 0 {
469            bail!("ctx_len ({ctx_len}) must be divisible by patch_size ({patch_size})");
470        }
471        let ctx_patches = ctx_len / patch_size;
472        let fcst_patches = (prediction_length + patch_size - 1) / patch_size;
473        let total_patches = ctx_patches + fcst_patches;
474
475        let mut locs = Vec::with_capacity(n_var);
476        let mut scales = Vec::with_capacity(n_var);
477        for v in 0..n_var {
478            let (loc, scale) = causal_patched_std_scaler(&target[v], &mask[v], patch_size);
479            locs.push(loc);
480            scales.push(scale);
481        }
482
483        let mut patch_data = vec![0.0f32; n_var * total_patches * 2 * patch_size];
484        for v in 0..n_var {
485            for p in 0..ctx_patches {
486                let base = v * total_patches * 2 * patch_size + p * 2 * patch_size;
487                for i in 0..patch_size {
488                    let t = p * patch_size + i;
489                    let obs = mask[v][t];
490                    let val = if obs {
491                        (((target[v][t] - locs[v][t]) / scales[v][t]) as f64).asinh() as f32
492                    } else {
493                        0.0
494                    };
495                    patch_data[base + i] = val;
496                    patch_data[base + patch_size + i] = if obs { 0.0 } else { 1.0 };
497                }
498            }
499            for p in ctx_patches..total_patches {
500                let base = v * total_patches * 2 * patch_size + p * 2 * patch_size;
501                for i in 0..patch_size {
502                    patch_data[base + i] = 0.0;
503                    patch_data[base + patch_size + i] = 1.0;
504                }
505            }
506        }
507
508        let dev = device();
509        let x = Tensor::<4>::from_data(
510            TensorData::new(patch_data, [1, n_var, total_patches, 2 * patch_size]),
511            &dev,
512        );
513        let x = self.forward_residual_mlp(x, &self.patch_proj)?;
514        let x = self.forward_transformer(x, n_var, total_patches)?;
515        let x = self.forward_residual_mlp(x, &self.output_head)?;
516        let x: Tensor<4> = x.narrow(2, ctx_patches - 1, fcst_patches);
517        let fcst_steps = fcst_patches * patch_size;
518        let out: Tensor<5> = x.reshape([1, n_var, fcst_patches, patch_size, 9]);
519        let out: Tensor<5> = out.permute([4, 0, 1, 2, 3]);
520        let out: Tensor<3> = out.reshape([9, n_var, fcst_steps]);
521        let h = out.dims()[2];
522        let out: Tensor<3> = if fcst_steps > prediction_length {
523            out.narrow(2, 0, prediction_length)
524        } else {
525            out
526        };
527        let out_data: Vec<f32> = out
528            .to_data()
529            .try_to_vec()
530            .map_err(|e| anyhow::anyhow!("{e:?}"))?;
531
532        let loc_final: Vec<f32> = (0..n_var).map(|v| locs[v][ctx_len - 1]).collect();
533        let scale_final: Vec<f32> = (0..n_var).map(|v| scales[v][ctx_len - 1]).collect();
534        let mut result = vec![vec![vec![0.0f32; prediction_length]; n_var]; 9];
535        for q in 0..9 {
536            for v in 0..n_var {
537                for t in 0..prediction_length {
538                    let raw = out_data[(q * n_var + v) * h + t] as f64;
539                    result[q][v][t] = (raw.sinh() as f32) * scale_final[v] + loc_final[v];
540                }
541            }
542        }
543        Ok(result)
544    }
545
546    fn forward_residual_mlp<const D: usize>(
547        &self,
548        x: Tensor<D>,
549        w: &BurnResidualMlpW,
550    ) -> Result<Tensor<D>> {
551        let h = uu_silu(uu_linear(x.clone(), w.l1_w.clone(), Some(w.l1_b.clone())));
552        let h = if w.is_output {
553            uu_linear_readout(h, w.l2_w.clone(), Some(w.l2_b.clone()))
554        } else {
555            uu_linear(h, w.l2_w.clone(), Some(w.l2_b.clone()))
556        };
557        let skip = if w.is_output {
558            uu_linear_readout(x, w.skip_w.clone(), Some(w.skip_b.clone()))
559        } else {
560            uu_linear(x, w.skip_w.clone(), Some(w.skip_b.clone()))
561        };
562        Ok(residual_add(h, skip, w.tau))
563    }
564
565    fn forward_transformer(
566        &self,
567        mut x: Tensor<4>,
568        n_var: usize,
569        num_patches: usize,
570    ) -> Result<Tensor<4>> {
571        for (idx, blk) in self.blocks.iter().enumerate() {
572            x = if self.config.is_variate_layer(idx) {
573                self.forward_variate_layer(x, blk, n_var, num_patches)?
574            } else {
575                self.forward_time_layer(x, blk, n_var, num_patches)?
576            };
577        }
578        Ok(rms_norm_weightless(x, self.config.norm_eps))
579    }
580
581    fn forward_time_layer(
582        &self,
583        x: Tensor<4>,
584        blk: &BurnBlockW,
585        n_var: usize,
586        num_patches: usize,
587    ) -> Result<Tensor<4>> {
588        let cfg = &self.config;
589        let state: Tensor<3> = x.reshape([n_var, num_patches, cfg.d_model]);
590        let normed = rms_norm_weightless(state.clone(), cfg.norm_eps);
591        let seq_ids: Vec<u32> = (0..num_patches as u32).collect();
592        let attn_out = self.forward_attention(&normed, blk, &seq_ids, false)?;
593        let state = residual_add(attn_out, state, blk.attn_tau);
594        let normed = rms_norm_weightless(state.clone(), cfg.norm_eps);
595        let ffn_out = self.forward_ffn(&normed, blk)?;
596        let state = residual_add(ffn_out, state, blk.mlp_tau);
597        Ok(state.reshape([1, n_var, num_patches, cfg.d_model]))
598    }
599
600    fn forward_variate_layer(
601        &self,
602        x: Tensor<4>,
603        blk: &BurnBlockW,
604        n_var: usize,
605        num_patches: usize,
606    ) -> Result<Tensor<4>> {
607        let cfg = &self.config;
608        let state: Tensor<3> = x
609            .permute([0, 2, 1, 3])
610            .reshape([num_patches, n_var, cfg.d_model]);
611        let normed = rms_norm_weightless(state.clone(), cfg.norm_eps);
612        let attn_out = self.forward_attention(&normed, blk, &[], true)?;
613        let state = residual_add(attn_out, state, blk.attn_tau);
614        let normed = rms_norm_weightless(state.clone(), cfg.norm_eps);
615        let ffn_out = self.forward_ffn(&normed, blk)?;
616        let state = residual_add(ffn_out, state, blk.mlp_tau);
617        Ok(state
618            .reshape([1, num_patches, n_var, cfg.d_model])
619            .permute([0, 2, 1, 3]))
620    }
621
622    fn forward_attention(
623        &self,
624        state: &Tensor<3>,
625        blk: &BurnBlockW,
626        seq_ids: &[u32],
627        is_variate: bool,
628    ) -> Result<Tensor<3>> {
629        let cfg = &self.config;
630        let dims = state.dims();
631        let (batch, seq) = (dims[0], dims[1]);
632        let qkv = uu_linear(
633            state.clone(),
634            blk.attn_qkv_w.clone(),
635            blk.attn_qkv_b.clone(),
636        );
637        let q: Tensor<3> = qkv.clone().narrow(2, 0, cfg.q_size());
638        let k: Tensor<3> = qkv.clone().narrow(2, cfg.q_size(), cfg.k_size());
639        let v: Tensor<3> = qkv.narrow(2, cfg.q_size() + cfg.k_size(), cfg.v_size());
640        let pack = |t: Tensor<3>, heads: usize, hd: usize| {
641            t.reshape([batch, seq, heads, hd]).permute([0, 2, 1, 3])
642        };
643        let (q, k, v) = (
644            pack(q, cfg.num_heads, cfg.qk_dim),
645            pack(k, cfg.num_groups, cfg.qk_dim),
646            pack(v, cfg.num_groups, cfg.v_dim),
647        );
648        let q = match blk.attn_pds.as_ref() {
649            Some(pds) => apply_per_dim_scale(q, pds.clone()),
650            None => q,
651        };
652        let (q, k) = if !is_variate && !seq_ids.is_empty() && cfg.use_xpos {
653            (
654                self.rope.apply(q, seq_ids, 1.0),
655                self.rope.apply(k, seq_ids, -1.0),
656            )
657        } else {
658            (q, k)
659        };
660        let scale = 1.0f32 / cfg.qk_dim as f32;
661        let scores = q.matmul(k.transpose()).mul_scalar(scale);
662        let scores = if !is_variate {
663            scores + burn_causal_mask(seq)
664        } else {
665            scores
666        };
667        let attn = activation::softmax(scores, 3);
668        let out = attn.matmul(v);
669        let out: Tensor<3> =
670            out.permute([0, 2, 1, 3])
671                .reshape([batch, seq, cfg.num_heads * cfg.v_dim]);
672        Ok(uu_linear(
673            out,
674            blk.attn_out_w.clone(),
675            blk.attn_out_b.clone(),
676        ))
677    }
678
679    fn forward_ffn(&self, x: &Tensor<3>, blk: &BurnBlockW) -> Result<Tensor<3>> {
680        let fc1_out = uu_linear(x.clone(), blk.ffn_up_w.clone(), None);
681        let half = fc1_out.dims()[2] / 2;
682        let gate = fc1_out.clone().narrow(2, 0, half);
683        let val = fc1_out.narrow(2, half, half);
684        let activated = gate * activation::silu(val);
685        Ok(uu_linear(activated, blk.ffn_down_w.clone(), None))
686    }
687}
688
689fn burn_causal_mask(seq: usize) -> Tensor<4> {
690    let dev = device();
691    let mut data = vec![0.0f32; seq * seq];
692    for i in 0..seq {
693        for j in (i + 1)..seq {
694            data[i * seq + j] = f32::NEG_INFINITY;
695        }
696    }
697    Tensor::<2>::from_data(TensorData::new(data, [seq, seq]), &dev)
698        .unsqueeze_dim::<3>(0)
699        .unsqueeze_dim::<4>(0)
700}
701
702fn causal_patched_std_scaler(
703    data: &[f32],
704    mask: &[bool],
705    patch_size: usize,
706) -> (Vec<f32>, Vec<f32>) {
707    let n = data.len();
708    let num_patches = n / patch_size;
709    let mut cum_count = 0.0f64;
710    let mut m1 = 0.0f64;
711    let mut m2 = 0.0f64;
712    let correction = 1.0f64;
713    let minimum_scale = 1e-6f64;
714    let mut patch_loc = vec![0.0f32; num_patches];
715    let mut patch_scale = vec![1e-6f32; num_patches];
716    for p in 0..num_patches {
717        for i in 0..patch_size {
718            let t = p * patch_size + i;
719            if mask[t] {
720                let x = data[t] as f64;
721                cum_count += 1.0;
722                let prev_m1 = m1;
723                m1 += (x - m1) / cum_count;
724                m2 += (x - prev_m1) * (x - m1);
725            }
726        }
727        patch_loc[p] = m1 as f32;
728        let denom = (cum_count - correction).max(1.0);
729        patch_scale[p] = (m2 / denom).sqrt().max(minimum_scale) as f32;
730    }
731    let mut loc = vec![0.0f32; n];
732    let mut scale = vec![1e-6f32; n];
733    for p in 0..num_patches {
734        for i in 0..patch_size {
735            let t = p * patch_size + i;
736            loc[t] = patch_loc[p];
737            scale[t] = patch_scale[p];
738        }
739    }
740    (loc, scale)
741}
742
743#[cfg(test)]
744mod tests {
745    use super::*;
746    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
747
748    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
749        let mut x = seed;
750        (0..n)
751            .map(|_| {
752                x = x
753                    .wrapping_mul(6364136223846793005)
754                    .wrapping_add(1442695040888963407);
755                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
756            })
757            .collect()
758    }
759
760    fn tiny_config() -> InferConfig {
761        InferConfig::default()
762            .with_d_model(16)
763            .with_num_layers(1)
764            .with_num_heads(4)
765            .with_num_groups(4)
766            .with_qk_dim(8)
767            .with_v_dim(8)
768            .with_patch_size(4)
769    }
770
771    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
772        *seed += 1;
773        let data: Vec<f32> = pseudo(rows * cols, *seed)
774            .iter()
775            .map(|v| v * scale)
776            .collect();
777        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
778        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
779    }
780
781    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
782        *seed += 1;
783        let bytes: Vec<u8> = pseudo(n, *seed)
784            .iter()
785            .flat_map(|v| v.to_le_bytes())
786            .collect();
787        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
788    }
789
790    fn put_mlp(
791        w: &mut GGUFWriter,
792        prefix: &str,
793        h1: usize,
794        d_in: usize,
795        h2: usize,
796        seed: &mut u64,
797        s: f32,
798    ) {
799        put2(w, &format!("{prefix}.linear1.weight"), h1, d_in, seed, s);
800        put1(w, &format!("{prefix}.linear1.bias"), h1, seed);
801        put2(w, &format!("{prefix}.linear2.weight"), h2, h1, seed, s);
802        put1(w, &format!("{prefix}.linear2.bias"), h2, seed);
803        put2(w, &format!("{prefix}.skip_proj.weight"), h2, d_in, seed, s);
804        put1(w, &format!("{prefix}.skip_proj.bias"), h2, seed);
805    }
806
807    #[test]
808    fn burn_forecast_golden() {
809        // d=16, heads=groups=4, qk=v=8, patch=4, 1 layer, all time layers.
810        let dir = std::env::temp_dir().join(format!("toto-burn-{}", std::process::id()));
811        let _ = std::fs::remove_dir_all(&dir);
812        std::fs::create_dir_all(&dir).unwrap();
813        let path = dir.join("model.gguf");
814        let mut w = GGUFWriter::new();
815        w.add_metadata("general.architecture", GGUFMetaValue::String("toto".into()));
816        let mut seed = 1100u64;
817        let s = 0.05;
818        put_mlp(&mut w, "patch_proj", 32, 8, 16, &mut seed, s);
819        put2(&mut w, "blk.0.attn_qkv.weight", 96, 16, &mut seed, s);
820        put1(&mut w, "blk.0.attn_qkv.bias", 96, &mut seed);
821        put2(&mut w, "blk.0.attn_output.weight", 16, 32, &mut seed, s);
822        put1(&mut w, "blk.0.attn_output.bias", 16, &mut seed);
823        put1(&mut w, "blk.0.attn_pds.weight", 8, &mut seed);
824        put2(&mut w, "blk.0.ffn_up.weight", 32, 16, &mut seed, s);
825        put2(&mut w, "blk.0.ffn_down.weight", 16, 16, &mut seed, s);
826        put_mlp(&mut w, "output_head", 32, 16, 36, &mut seed, s);
827        let f = std::fs::File::create(&path).unwrap();
828        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
829
830        let cfg = tiny_config();
831        assert!(!cfg.compute_f64);
832        let burn = BurnTotoModel::load(&path, cfg).unwrap();
833        let ctx = vec![(0..8).map(|i| 10.0 + 0.5 * i as f32).collect::<Vec<_>>()];
834        let mask = vec![vec![true; 8]];
835        let t0 = std::time::Instant::now();
836        let a = burn.forecast(&ctx, &mask, 4).unwrap();
837        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
838        let b = burn.forecast(&ctx, &mask, 4).unwrap();
839        assert_eq!(a.len(), 9);
840        for q in &a {
841            assert_eq!(q.len(), 1);
842            for v in q {
843                assert_eq!(v.len(), 4);
844            }
845        }
846        assert!(a.iter().flatten().flatten().all(|v| v.is_finite()));
847        assert!(a.iter().flatten().flatten().any(|&v| v != 0.0));
848        assert_eq!(a, b, "forecast must be deterministic");
849        println!("toto synthetic: burn {burn_ms:.1}ms out={a:?}");
850        // Golden values recorded from the verified Burn port (ex-candle parity 1e-3).
851        let flat: Vec<f32> = a.iter().flatten().flatten().copied().collect();
852        let golden: [f32; 36] = [
853            11.704744, 11.739721, 11.7525215, 11.725984, 11.722208, 11.724002, 11.743584,
854            11.738862, 11.74826, 11.730129, 11.741938, 11.736687, 11.761163, 11.724331, 11.729776,
855            11.752375, 11.711724, 11.757467, 11.735854, 11.761217, 11.6961, 11.734921, 11.741099,
856            11.715156, 11.7233715, 11.694718, 11.721759, 11.7119465, 11.716567, 11.721282,
857            11.730694, 11.710778, 11.740484, 11.723714, 11.738366, 11.721424,
858        ];
859        assert_eq!(flat.len(), golden.len());
860        for (x, y) in flat.iter().zip(golden.iter()) {
861            assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
862        }
863        let _ = std::fs::remove_dir_all(&dir);
864    }
865}
zsfm_toto/infer/mod.rs

zsfm_toto/infer/
mod.rs
