zsfm_timesfm/infer/
mod.rs

1//! TimesFM 2.5 Burn-only inference engine (Burn backend, Flex CPU).
2//!
3//! Same RevIN running stats, tokenizer, prefill + KV-cached decode, RoPE,
4//! and denormalization, with weights loaded straight from GGUF via
5//! [`zsfm_burn::GgufWeights`]. Architecture constants mirror the model.
6
7use std::path::Path;
8
9use anyhow::{Context, Result};
10use burn::tensor::{activation, Device, Tensor, TensorData};
11use zsfm_burn::{open_gguf_weights, GgufWeights};
12
13const D_MODEL: usize = 1280;
14const N_HEADS: usize = 16;
15const HEAD_DIM: usize = 80;
16const N_LAYERS: usize = 20;
17const INPUT_PATCH: usize = 32;
18const OUTPUT_PATCH: usize = 128;
19const N_OUTPUTS: usize = 10;
20const DECODE_IDX: usize = 5;
21const M_PATCHES: usize = OUTPUT_PATCH / INPUT_PATCH;
22const ROPE_THETA: f64 = 10000.0;
23const MAX_SEQ: usize = 16384 / INPUT_PATCH + 256;
24const NORM_EPS: f32 = 1e-6;
25const REVIN_TOL: f32 = 1e-6;
26
27fn device() -> Device {
28    Device::flex()
29}
30
31fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
32    weights.load_2d(name)
33}
34
35fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
36    weights.load_1d(name)
37}
38
39fn load_w(weights: &GgufWeights, name: &str, expected_d_out: usize) -> Result<Tensor<2>> {
40    weights.load_weight_2d(name, expected_d_out)
41}
42
43fn load_flat(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
44    weights.load_flat(name)
45}
46
47struct BurnResidualBlockW {
48    hidden_w: Tensor<2>,
49    hidden_b: Option<Tensor<1>>,
50    output_w: Tensor<2>,
51    output_b: Option<Tensor<1>>,
52    skip_w: Tensor<2>,
53    skip_b: Option<Tensor<1>>,
54}
55
56struct BurnAttnW {
57    qkv_w: Tensor<2>,
58    out_w: Tensor<2>,
59    q_norm_w: Tensor<1>,
60    k_norm_w: Tensor<1>,
61    q_scale: Tensor<1>,
62    pre_norm_w: Tensor<1>,
63    post_norm_w: Tensor<1>,
64}
65
66struct BurnFfnW {
67    up_w: Tensor<2>,
68    down_w: Tensor<2>,
69    pre_norm_w: Tensor<1>,
70    post_norm_w: Tensor<1>,
71}
72
73struct BurnBlockW {
74    attn: BurnAttnW,
75    ffn: BurnFfnW,
76}
77
78pub struct BurnTimesFMModel {
79    tokenizer: BurnResidualBlockW,
80    blocks: Vec<BurnBlockW>,
81    out_point: BurnResidualBlockW,
82    rope: BurnRope,
83}
84
85struct BurnRope {
86    cos: Tensor<2>,
87    sin: Tensor<2>,
88    head_dim: usize,
89}
90
91impl BurnRope {
92    fn build(head_dim: usize, max_seq: usize, theta: f64) -> Self {
93        let dev = device();
94        let half = head_dim / 2;
95        let inv_freq: Vec<f64> = (0..half)
96            .map(|i| 1.0 / theta.powf(2.0 * i as f64 / head_dim as f64))
97            .collect();
98        let mut cos = vec![0.0f32; max_seq * head_dim];
99        let mut sin = vec![0.0f32; max_seq * head_dim];
100        for p in 0..max_seq {
101            for i in 0..half {
102                let angle = (p as f64 * inv_freq[i]) as f32;
103                let (s, c) = angle.sin_cos();
104                cos[p * head_dim + i] = c;
105                cos[p * head_dim + half + i] = c;
106                sin[p * head_dim + i] = s;
107                sin[p * head_dim + half + i] = s;
108            }
109        }
110        Self {
111            cos: Tensor::<2>::from_data(TensorData::new(cos, [max_seq, head_dim]), &dev),
112            sin: Tensor::<2>::from_data(TensorData::new(sin, [max_seq, head_dim]), &dev),
113            head_dim,
114        }
115    }
116
117    /// Apply RoPE to x with shape [b, seq, n_heads, head_dim].
118    fn apply(&self, x: Tensor<4>, start_pos: usize) -> Tensor<4> {
119        let [b, seq, _h, hd] = x.dims();
120        let half = hd / 2;
121        let cos = self
122            .cos
123            .clone()
124            .narrow(0, start_pos, seq)
125            .unsqueeze_dim::<3>(0)
126            .unsqueeze_dim::<4>(2);
127        let sin = self
128            .sin
129            .clone()
130            .narrow(0, start_pos, seq)
131            .unsqueeze_dim::<3>(0)
132            .unsqueeze_dim::<4>(2);
133        let x1 = x.clone().narrow(3, 0, half);
134        let x2 = x.clone().narrow(3, half, half);
135        let cos1 = cos.clone().narrow(3, 0, half);
136        let sin1 = sin.clone().narrow(3, 0, half);
137        let _ = (b, self.head_dim);
138        let first = x1.clone() * cos1.clone() - x2.clone() * sin1.clone();
139        let second = x2 * cos1 + x1 * sin1;
140        Tensor::cat(vec![first, second], 3)
141    }
142}
143
144fn linear(x: Tensor<2>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<2> {
145    match b {
146        Some(bias) => zsfm_burn::linear::burn_linear_bias(x, w, bias),
147        None => zsfm_burn::linear::burn_linear_nobias(x, w),
148    }
149}
150
151fn linear_nd<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
152    let dims = x.dims();
153    let last = dims[D - 1];
154    let lead: usize = dims[..D - 1].iter().product();
155    let out_d = w.dims()[0];
156    let flat = x.reshape([lead, last]);
157    let y = linear(flat, w, b);
158    let mut out_dims = dims;
159    out_dims[D - 1] = out_d;
160    y.reshape(out_dims)
161}
162
163fn rms_norm_2d(x: Tensor<2>, w: Tensor<1>) -> Tensor<2> {
164    zsfm_burn::norm::burn_rms_norm(x, w, NORM_EPS)
165}
166
167fn rms_norm_nd<const D: usize>(x: Tensor<D>, w: Tensor<1>) -> Tensor<D> {
168    let dims = x.dims();
169    let last = dims[D - 1];
170    let lead: usize = dims[..D - 1].iter().product();
171    let y = zsfm_burn::norm::burn_rms_norm(x.reshape([lead, last]), w, NORM_EPS);
172    y.reshape(dims)
173}
174
175fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
176    activation::softmax(x, D - 1)
177}
178
179fn silu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
180    activation::silu(x)
181}
182
183fn softplus(x: f32) -> f32 {
184    if x > 20.0 {
185        x
186    } else {
187        (1.0f32 + x.exp()).ln()
188    }
189}
190
191fn compute_q_scale(raw: Vec<f32>) -> Vec<f32> {
192    let factor = std::f32::consts::LOG2_E / (HEAD_DIM as f32).sqrt();
193    raw.into_iter().map(|x| factor * softplus(x)).collect()
194}
195
196fn load_residual_block(
197    weights: &GgufWeights,
198    prefix: &str,
199    hidden_d_out: usize,
200    out_d_out: usize,
201    with_bias: bool,
202) -> Result<BurnResidualBlockW> {
203    let hidden_w = load_w(&weights, &format!("{prefix}.hidden.weight"), hidden_d_out)?;
204    let hidden_b = with_bias
205        .then(|| load_v(&weights, &format!("{prefix}.hidden.bias")))
206        .transpose()?;
207    let output_w = load_w(&weights, &format!("{prefix}.output.weight"), out_d_out)?;
208    let output_b = with_bias
209        .then(|| load_v(&weights, &format!("{prefix}.output.bias")))
210        .transpose()?;
211    let skip_w = load_w(&weights, &format!("{prefix}.skip.weight"), out_d_out)?;
212    let skip_b = with_bias
213        .then(|| load_v(&weights, &format!("{prefix}.skip.bias")))
214        .transpose()?;
215    Ok(BurnResidualBlockW {
216        hidden_w,
217        hidden_b,
218        output_w,
219        output_b,
220        skip_w,
221        skip_b,
222    })
223}
224
225fn forward_residual_block(x: Tensor<2>, w: &BurnResidualBlockW) -> Tensor<2> {
226    let h = silu(linear(x.clone(), w.hidden_w.clone(), w.hidden_b.clone()));
227    let out = linear(h, w.output_w.clone(), w.output_b.clone());
228    let skip = linear(x, w.skip_w.clone(), w.skip_b.clone());
229    out + skip
230}
231
232fn burn_causal_mask(n: usize) -> Tensor<4> {
233    let dev = device();
234    let data: Vec<f32> = (0..n)
235        .flat_map(|q| (0..n).map(move |k| if k <= q { 0.0f32 } else { f32::NEG_INFINITY }))
236        .collect();
237    Tensor::<2>::from_data(TensorData::new(data, [n, n]), &dev)
238        .unsqueeze_dim::<3>(0)
239        .unsqueeze_dim::<4>(0)
240}
241
242fn burn_decode_mask(new_len: usize, cache_len: usize) -> Tensor<4> {
243    let dev = device();
244    let total = cache_len + new_len;
245    let data: Vec<f32> = (0..new_len)
246        .flat_map(|q_rel| {
247            (0..total).map(move |k_abs| {
248                if k_abs <= cache_len + q_rel {
249                    0.0f32
250                } else {
251                    f32::NEG_INFINITY
252                }
253            })
254        })
255        .collect();
256    Tensor::<2>::from_data(TensorData::new(data, [new_len, total]), &dev)
257        .unsqueeze_dim::<3>(0)
258        .unsqueeze_dim::<4>(0)
259}
260
261impl BurnTimesFMModel {
262    pub fn load(gguf_path: &Path) -> Result<Self> {
263        let weights = open_gguf_weights(gguf_path)
264            .with_context(|| format!("open {}", gguf_path.display()))?;
265
266        let tokenizer = load_residual_block(&weights, "tokenizer", D_MODEL, D_MODEL, true)?;
267
268        let mut blocks = Vec::with_capacity(N_LAYERS);
269        for n in 0..N_LAYERS {
270            let b = format!("blk.{n}");
271            let qkv_w = load_w(&weights, &format!("{b}.attn_qkv.weight"), 3 * D_MODEL)?;
272            let out_w = load_w(&weights, &format!("{b}.attn_out.weight"), D_MODEL)?;
273            let q_norm_w = load_v(&weights, &format!("{b}.attn_q_norm.weight"))?;
274            let k_norm_w = load_v(&weights, &format!("{b}.attn_k_norm.weight"))?;
275            let q_scale_raw: Vec<f32> = load_flat(&weights, &format!("{b}.attn_q_scale.weight"))?;
276            let q_scale = Tensor::<1>::from_data(
277                TensorData::new(compute_q_scale(q_scale_raw), [HEAD_DIM]),
278                &device(),
279            );
280            let pre_norm_w = load_v(&weights, &format!("{b}.pre_attn_norm.weight"))?;
281            let post_norm_w = load_v(&weights, &format!("{b}.post_attn_norm.weight"))?;
282            let up_w = load_w(&weights, &format!("{b}.ffn_up.weight"), D_MODEL)?;
283            let down_w = load_w(&weights, &format!("{b}.ffn_down.weight"), D_MODEL)?;
284            let pre_ff = load_v(&weights, &format!("{b}.pre_ff_norm.weight"))?;
285            let post_ff = load_v(&weights, &format!("{b}.post_ff_norm.weight"))?;
286            blocks.push(BurnBlockW {
287                attn: BurnAttnW {
288                    qkv_w,
289                    out_w,
290                    q_norm_w,
291                    k_norm_w,
292                    q_scale,
293                    pre_norm_w,
294                    post_norm_w,
295                },
296                ffn: BurnFfnW {
297                    up_w,
298                    down_w,
299                    pre_norm_w: pre_ff,
300                    post_norm_w: post_ff,
301                },
302            });
303        }
304
305        let out_point = load_residual_block(&weights, "out_point", D_MODEL, D_MODEL, false)?;
306        let rope = BurnRope::build(HEAD_DIM, MAX_SEQ, ROPE_THETA);
307        Ok(Self {
308            tokenizer,
309            blocks,
310            out_point,
311            rope,
312        })
313    }
314
315    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> Result<Vec<Vec<f32>>> {
316        let p = INPUT_PATCH;
317        let o = OUTPUT_PATCH;
318        let q = N_OUTPUTS;
319        let m = M_PATCHES;
320
321        let len_front = if context.len() % p == 0 {
322            0
323        } else {
324            p - context.len() % p
325        };
326        let mut vals = vec![0.0f32; len_front + context.len()];
327        vals[len_front..].copy_from_slice(context);
328        let mut mask = vec![true; len_front];
329        mask.extend(vec![false; context.len()]);
330        let n_ctx = (len_front + context.len()) / p;
331
332        let mut patch_mus = Vec::with_capacity(n_ctx);
333        let mut patch_sigmas = Vec::with_capacity(n_ctx);
334        let (mut rs_n, mut rs_mu, mut rs_sigma) = (0.0f32, 0.0f32, 0.0f32);
335        for i in 0..n_ctx {
336            let pv = &vals[i * p..(i + 1) * p];
337            let pm = &mask[i * p..(i + 1) * p];
338            (rs_n, rs_mu, rs_sigma) = update_running_stats(rs_n, rs_mu, rs_sigma, pv, pm);
339            patch_mus.push(rs_mu);
340            patch_sigmas.push(rs_sigma);
341        }
342        let (last_n, last_mu, last_sigma) = (rs_n, rs_mu, rs_sigma);
343        let ctx_input = build_tokenizer_input(&vals, &mask, n_ctx, p, &patch_mus, &patch_sigmas);
344
345        let num_decode_steps = if prediction_length <= o {
346            0
347        } else {
348            (prediction_length - 1) / o
349        };
350        let total_steps = 1 + num_decode_steps;
351        let mut all_outputs: Vec<Vec<[f32; N_OUTPUTS]>> = Vec::with_capacity(total_steps);
352
353        let (ctx_out, mut kv_cache) = self.prefill(ctx_input, n_ctx);
354        let ctx_flat: Vec<f32> = ctx_out
355            .to_data()
356            .try_to_vec()
357            .map_err(|e| anyhow::anyhow!("{e:?}"))?;
358        let ctx_denorm = denorm_flat(&ctx_flat, &patch_mus, &patch_sigmas, o, q);
359        all_outputs.push(ctx_denorm[n_ctx - 1].clone());
360
361        let (mut ar_n, mut ar_mu, mut ar_sigma) = (last_n, last_mu, last_sigma);
362        let mut last_ar: Vec<f32> = ctx_denorm[n_ctx - 1]
363            .iter()
364            .map(|row| row[DECODE_IDX])
365            .collect();
366
367        for step in 0..num_decode_steps {
368            let new_vals_flat = last_ar.clone();
369            let new_mask_flat = vec![false; o];
370            let mut new_mus = Vec::with_capacity(m);
371            let mut new_sigmas = Vec::with_capacity(m);
372            for mi in 0..m {
373                let pv = &new_vals_flat[mi * p..(mi + 1) * p];
374                let pm = &new_mask_flat[mi * p..(mi + 1) * p];
375                (ar_n, ar_mu, ar_sigma) = update_running_stats(ar_n, ar_mu, ar_sigma, pv, pm);
376                new_mus.push(ar_mu);
377                new_sigmas.push(ar_sigma);
378            }
379            let new_input =
380                build_tokenizer_input(&new_vals_flat, &new_mask_flat, m, p, &new_mus, &new_sigmas);
381            let rope_offset = n_ctx + m * step;
382            let new_out = self.decode_chunk(new_input, &mut kv_cache, rope_offset);
383            let new_flat: Vec<f32> = new_out
384                .to_data()
385                .try_to_vec()
386                .map_err(|e| anyhow::anyhow!("{e:?}"))?;
387            let last_m_denorm = denorm_flat(&new_flat, &new_mus, &new_sigmas, o, q);
388            let step_out = &last_m_denorm[m - 1];
389            last_ar = step_out.iter().map(|row| row[DECODE_IDX]).collect();
390            all_outputs.push(step_out.clone());
391        }
392
393        let total_available = all_outputs.len() * o;
394        let n_take = prediction_length.min(total_available);
395        let mut result: Vec<Vec<f32>> = vec![Vec::with_capacity(n_take); q];
396        'outer: for step in &all_outputs {
397            for timestep in step {
398                for qi in 0..q {
399                    if result[qi].len() >= prediction_length {
400                        break 'outer;
401                    }
402                    result[qi].push(timestep[qi]);
403                }
404            }
405        }
406        Ok(result)
407    }
408
409    fn prefill(
410        &self,
411        tokenizer_input: Tensor<2>,
412        n_patches: usize,
413    ) -> (Tensor<2>, Vec<(Tensor<4>, Tensor<4>)>) {
414        let x = forward_residual_block(tokenizer_input, &self.tokenizer);
415        let mut hidden = x.unsqueeze_dim::<3>(0);
416        let causal = burn_causal_mask(n_patches);
417        let mut kv_cache = Vec::with_capacity(N_LAYERS);
418        for block in &self.blocks {
419            let (h_out, k, v) = self.prefill_block(hidden, block, n_patches, causal.clone());
420            hidden = h_out;
421            kv_cache.push((k, v));
422        }
423        let out_seq: Tensor<2> = hidden.squeeze_dim(0);
424        let out = forward_residual_block(out_seq, &self.out_point);
425        (out, kv_cache)
426    }
427
428    fn prefill_block(
429        &self,
430        x: Tensor<3>,
431        w: &BurnBlockW,
432        n_patches: usize,
433        causal: Tensor<4>,
434    ) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
435        let normed = rms_norm_nd(x.clone(), w.attn.pre_norm_w.clone());
436        let (attn_out, k, v) = self.prefill_attn(normed, &w.attn, n_patches, causal);
437        let attn_out = rms_norm_nd(attn_out, w.attn.post_norm_w.clone()) + x;
438        let normed_ff = rms_norm_nd(attn_out.clone(), w.ffn.pre_norm_w.clone());
439        let ff_out = self.forward_ffn(normed_ff, &w.ffn);
440        let out = rms_norm_nd(ff_out, w.ffn.post_norm_w.clone()) + attn_out;
441        (out, k, v)
442    }
443
444    fn prefill_attn(
445        &self,
446        x: Tensor<3>,
447        w: &BurnAttnW,
448        n_patches: usize,
449        causal: Tensor<4>,
450    ) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
451        burn_prefill_attn(x, w, &self.rope, n_patches, causal)
452    }
453
454    fn decode_chunk(
455        &self,
456        new_input: Tensor<2>,
457        kv_cache: &mut Vec<(Tensor<4>, Tensor<4>)>,
458        rope_offset: usize,
459    ) -> Tensor<2> {
460        let x = forward_residual_block(new_input, &self.tokenizer);
461        let mut hidden = x.unsqueeze_dim::<3>(0);
462        let cached_len = rope_offset;
463        let decode_mask = burn_decode_mask(M_PATCHES, cached_len);
464        for (li, block) in self.blocks.iter().enumerate() {
465            let (h_out, k_new, v_new) = self.decode_block_kv(
466                hidden,
467                block,
468                &kv_cache[li],
469                rope_offset,
470                decode_mask.clone(),
471            );
472            hidden = h_out;
473            kv_cache[li] = (k_new, v_new);
474        }
475        let out_seq: Tensor<2> = hidden.squeeze_dim(0);
476        forward_residual_block(out_seq, &self.out_point)
477    }
478
479    fn decode_block_kv(
480        &self,
481        x: Tensor<3>,
482        w: &BurnBlockW,
483        cache: &(Tensor<4>, Tensor<4>),
484        rope_offset: usize,
485        decode_mask: Tensor<4>,
486    ) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
487        let normed = rms_norm_nd(x.clone(), w.attn.pre_norm_w.clone());
488        let (attn_out, k_new, v_new) =
489            self.decode_attn_kv(normed, &w.attn, cache, rope_offset, decode_mask);
490        let attn_out = rms_norm_nd(attn_out, w.attn.post_norm_w.clone()) + x;
491        let normed_ff = rms_norm_nd(attn_out.clone(), w.ffn.pre_norm_w.clone());
492        let ff_out = self.forward_ffn(normed_ff, &w.ffn);
493        let out = rms_norm_nd(ff_out, w.ffn.post_norm_w.clone()) + attn_out;
494        (out, k_new, v_new)
495    }
496
497    fn decode_attn_kv(
498        &self,
499        x: Tensor<3>,
500        w: &BurnAttnW,
501        cache: &(Tensor<4>, Tensor<4>),
502        rope_offset: usize,
503        mask: Tensor<4>,
504    ) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
505        burn_decode_attn(x, w, &self.rope, cache, rope_offset, mask)
506    }
507
508    fn forward_ffn(&self, x: Tensor<3>, w: &BurnFfnW) -> Tensor<3> {
509        let h = linear_nd(x, w.up_w.clone(), None);
510        let h = silu(h);
511        linear_nd(h, w.down_w.clone(), None)
512    }
513}
514
515fn burn_prefill_attn(
516    x: Tensor<3>,
517    w: &BurnAttnW,
518    rope: &BurnRope,
519    n_patches: usize,
520    causal: Tensor<4>,
521) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
522    let qkv = linear_nd(x, w.qkv_w.clone(), None);
523    let q: Tensor<3> = qkv.clone().narrow(2, 0, D_MODEL);
524    let k: Tensor<3> = qkv.clone().narrow(2, D_MODEL, D_MODEL);
525    let v: Tensor<3> = qkv.narrow(2, 2 * D_MODEL, D_MODEL);
526    let q: Tensor<4> = q.reshape([1, n_patches, N_HEADS, HEAD_DIM]);
527    let k: Tensor<4> = k.reshape([1, n_patches, N_HEADS, HEAD_DIM]);
528    let v: Tensor<4> = v.reshape([1, n_patches, N_HEADS, HEAD_DIM]);
529    let q = rope.apply(q, 0);
530    let k = rope.apply(k, 0);
531    let q = rms_norm_nd(q, w.q_norm_w.clone());
532    let k = rms_norm_nd(k, w.k_norm_w.clone());
533    let scale = w
534        .q_scale
535        .clone()
536        .unsqueeze_dim::<2>(0)
537        .unsqueeze_dim::<3>(0)
538        .unsqueeze_dim::<4>(0);
539    let q = q * scale;
540    let q: Tensor<4> = q.permute([0, 2, 1, 3]);
541    let k: Tensor<4> = k.permute([0, 2, 1, 3]);
542    let v: Tensor<4> = v.permute([0, 2, 1, 3]);
543    let scores = q.clone().matmul(k.clone().transpose()) + causal;
544    let attn_w = softmax_last(scores);
545    let ctx = attn_w.matmul(v.clone());
546    let ctx: Tensor<4> = ctx.permute([0, 2, 1, 3]);
547    let ctx: Tensor<3> = ctx.reshape([1, n_patches, D_MODEL]);
548    (linear_nd(ctx, w.out_w.clone(), None), k, v)
549}
550
551fn burn_decode_attn(
552    x: Tensor<3>,
553    w: &BurnAttnW,
554    rope: &BurnRope,
555    cache: &(Tensor<4>, Tensor<4>),
556    rope_offset: usize,
557    mask: Tensor<4>,
558) -> (Tensor<3>, Tensor<4>, Tensor<4>) {
559    let qkv = linear_nd(x, w.qkv_w.clone(), None);
560    let q: Tensor<3> = qkv.clone().narrow(2, 0, D_MODEL);
561    let k: Tensor<3> = qkv.clone().narrow(2, D_MODEL, D_MODEL);
562    let v: Tensor<3> = qkv.narrow(2, 2 * D_MODEL, D_MODEL);
563    let q: Tensor<4> = q.reshape([1, M_PATCHES, N_HEADS, HEAD_DIM]);
564    let k: Tensor<4> = k.reshape([1, M_PATCHES, N_HEADS, HEAD_DIM]);
565    let v: Tensor<4> = v.reshape([1, M_PATCHES, N_HEADS, HEAD_DIM]);
566    let q = rope.apply(q, rope_offset);
567    let k = rope.apply(k, rope_offset);
568    let q = rms_norm_nd(q, w.q_norm_w.clone());
569    let k = rms_norm_nd(k, w.k_norm_w.clone());
570    let scale = w
571        .q_scale
572        .clone()
573        .unsqueeze_dim::<2>(0)
574        .unsqueeze_dim::<3>(0)
575        .unsqueeze_dim::<4>(0);
576    let q = q * scale;
577    let q: Tensor<4> = q.permute([0, 2, 1, 3]);
578    let k: Tensor<4> = k.permute([0, 2, 1, 3]);
579    let v: Tensor<4> = v.permute([0, 2, 1, 3]);
580    let k_full = Tensor::cat(vec![cache.0.clone(), k], 2);
581    let v_full = Tensor::cat(vec![cache.1.clone(), v], 2);
582    let scores = q.matmul(k_full.clone().transpose()) + mask;
583    let attn_w = softmax_last(scores);
584    let ctx = attn_w.matmul(v_full.clone());
585    let ctx: Tensor<4> = ctx.permute([0, 2, 1, 3]);
586    let ctx: Tensor<3> = ctx.reshape([1, M_PATCHES, D_MODEL]);
587    (linear_nd(ctx, w.out_w.clone(), None), k_full, v_full)
588}
589
590fn update_running_stats(
591    n: f32,
592    mu: f32,
593    sigma: f32,
594    vals: &[f32],
595    mask: &[bool],
596) -> (f32, f32, f32) {
597    let mut inc_n = 0.0f32;
598    let mut sum_x = 0.0f32;
599    for (i, &m) in mask.iter().enumerate() {
600        if !m {
601            inc_n += 1.0;
602            sum_x += vals[i];
603        }
604    }
605    let inc_mu = if inc_n == 0.0 { 0.0 } else { sum_x / inc_n };
606    let inc_var = if inc_n == 0.0 {
607        0.0
608    } else {
609        mask.iter()
610            .enumerate()
611            .filter(|(_, &m)| !m)
612            .map(|(i, _)| (vals[i] - inc_mu).powi(2))
613            .sum::<f32>()
614            / inc_n
615    };
616    let inc_sigma = inc_var.sqrt();
617    let new_n = n + inc_n;
618    let safe_new_n = if new_n == 0.0 { 1.0 } else { new_n };
619    let new_mu = if new_n == 0.0 {
620        0.0
621    } else {
622        (n * mu + inc_mu * inc_n) / safe_new_n
623    };
624    let t1 = n * sigma.powi(2);
625    let t2 = inc_n * inc_sigma.powi(2);
626    let t3 = n * (mu - new_mu).powi(2);
627    let t4 = inc_n * (inc_mu - new_mu).powi(2);
628    let new_var = if new_n == 0.0 {
629        0.0
630    } else {
631        (t1 + t2 + t3 + t4) / safe_new_n
632    };
633    (new_n, new_mu, new_var.max(0.0).sqrt())
634}
635
636fn build_tokenizer_input(
637    vals: &[f32],
638    mask: &[bool],
639    n_patches: usize,
640    p: usize,
641    mus: &[f32],
642    sigmas: &[f32],
643) -> Tensor<2> {
644    let mut data = vec![0.0f32; n_patches * 2 * p];
645    for pi in 0..n_patches {
646        let mu = mus[pi];
647        let sigma = sigmas[pi];
648        let sigma_safe = if sigma < REVIN_TOL { 1.0f32 } else { sigma };
649        for i in 0..p {
650            let idx = pi * p + i;
651            let is_masked = mask[idx];
652            let normed = if is_masked {
653                0.0
654            } else {
655                (vals[idx] - mu) / sigma_safe
656            };
657            let base = pi * 2 * p;
658            data[base + i] = normed;
659            data[base + p + i] = if is_masked { 1.0 } else { 0.0 };
660        }
661    }
662    Tensor::<2>::from_data(TensorData::new(data, [n_patches, 2 * p]), &device())
663}
664
665fn denorm_flat(
666    flat: &[f32],
667    mus: &[f32],
668    sigmas: &[f32],
669    o: usize,
670    q: usize,
671) -> Vec<Vec<[f32; N_OUTPUTS]>> {
672    let n = mus.len();
673    let mut result: Vec<Vec<[f32; N_OUTPUTS]>> = vec![vec![[0.0f32; N_OUTPUTS]; o]; n];
674    for pi in 0..n {
675        let mu = mus[pi];
676        let sigma = sigmas[pi];
677        for ti in 0..o {
678            for qi in 0..q {
679                let raw = flat[pi * (o * q) + ti * q + qi];
680                result[pi][ti][qi] = raw * sigma + mu;
681            }
682        }
683    }
684    result
685}
686
687#[cfg(test)]
688mod tests {
689    use super::*;
690
691    pub(crate) fn pseudo(n: usize, seed: u64) -> Vec<f32> {
692        let mut x = seed;
693        (0..n)
694            .map(|_| {
695                x = x
696                    .wrapping_mul(6364136223846793005)
697                    .wrapping_add(1442695040888963407);
698                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
699            })
700            .collect()
701    }
702
703    pub(crate) fn max_err(a: &[f32], b: &[f32]) -> f32 {
704        a.iter()
705            .zip(b.iter())
706            .map(|(x, y)| (x - y).abs())
707            .fold(0.0f32, f32::max)
708    }
709
710    /// Host reference for the GPT-NeoX-style RoPE in [`BurnRope`].
711    fn host_rope(x: &[f32], b: usize, seq: usize, h: usize, hd: usize, start: usize) -> Vec<f32> {
712        let half = hd / 2;
713        let inv_freq: Vec<f64> = (0..half)
714            .map(|i| 1.0 / ROPE_THETA.powf(2.0 * i as f64 / hd as f64))
715            .collect();
716        let mut out = vec![0.0f32; x.len()];
717        for bi in 0..b {
718            for s in 0..seq {
719                let pos = start + s;
720                for hi in 0..h {
721                    let base = (bi * seq * h + s * h + hi) * hd;
722                    for i in 0..half {
723                        let angle = (pos as f64 * inv_freq[i]) as f32;
724                        let (sn, cs) = angle.sin_cos();
725                        let x1 = x[base + i];
726                        let x2 = x[base + half + i];
727                        out[base + i] = x1 * cs - x2 * sn;
728                        out[base + half + i] = x2 * cs + x1 * sn;
729                    }
730                }
731            }
732        }
733        out
734    }
735
736    #[test]
737    fn rope_matches_host() {
738        let burn_rope = BurnRope::build(HEAD_DIM, 16, ROPE_THETA);
739        let xv = pseudo(1 * 4 * 2 * HEAD_DIM, 7);
740        let bx =
741            Tensor::<4>::from_data(TensorData::new(xv.clone(), [1, 4, 2, HEAD_DIM]), &device());
742        let got: Vec<f32> = burn_rope.apply(bx, 2).to_data().try_to_vec().unwrap();
743        let expected = host_rope(&xv, 1, 4, 2, HEAD_DIM, 2);
744        let err = max_err(&expected, &got);
745        assert!(err < 1e-5, "rope mismatch: {err:.2e}");
746    }
747
748    fn host_linear(
749        x: &[f32],
750        n: usize,
751        w: &[f32],
752        r: usize,
753        c: usize,
754        b: Option<&[f32]>,
755    ) -> Vec<f32> {
756        assert_eq!(x.len(), n * c);
757        assert_eq!(w.len(), r * c);
758        let mut y = vec![0.0f32; n * r];
759        for i in 0..n {
760            for o in 0..r {
761                let mut acc = 0.0f32;
762                for k in 0..c {
763                    acc += w[o * c + k] * x[i * c + k];
764                }
765                if let Some(bv) = b {
766                    acc += bv[o];
767                }
768                y[i * r + o] = acc;
769            }
770        }
771        y
772    }
773
774    fn host_silu(x: &[f32]) -> Vec<f32> {
775        x.iter().map(|&v| v / (1.0 + (-v).exp())).collect()
776    }
777
778    #[test]
779    fn residual_block_matches_host() {
780        // Small dims: in=8, hidden=16, out=16, with bias.
781        let (din, h, dout, n) = (8usize, 16usize, 16usize, 4usize);
782        let w = |r: usize, c: usize, s: u64| pseudo(r * c, s);
783        let hw = w(h, din, 21);
784        let hb = pseudo(h, 22);
785        let ow = w(dout, h, 23);
786        let ob = pseudo(dout, 24);
787        let sw = w(dout, din, 25);
788        let sb = pseudo(dout, 26);
789        let xv = pseudo(n * din, 27);
790
791        let hh = host_silu(&host_linear(&xv, n, &hw, h, din, Some(&hb)));
792        let out = host_linear(&hh, n, &ow, dout, h, Some(&ob));
793        let skip = host_linear(&xv, n, &sw, dout, din, Some(&sb));
794        let expected: Vec<f32> = out.iter().zip(skip.iter()).map(|(a, b)| a + b).collect();
795
796        let bdev = device();
797        let bxb = Tensor::<2>::from_data(TensorData::new(xv, [n, din]), &bdev);
798        let bw = BurnResidualBlockW {
799            hidden_w: Tensor::<2>::from_data(TensorData::new(hw, [h, din]), &bdev),
800            hidden_b: Some(Tensor::<1>::from_data(TensorData::new(hb, [h]), &bdev)),
801            output_w: Tensor::<2>::from_data(TensorData::new(ow, [dout, h]), &bdev),
802            output_b: Some(Tensor::<1>::from_data(TensorData::new(ob, [dout]), &bdev)),
803            skip_w: Tensor::<2>::from_data(TensorData::new(sw, [dout, din]), &bdev),
804            skip_b: Some(Tensor::<1>::from_data(TensorData::new(sb, [dout]), &bdev)),
805        };
806        let got: Vec<f32> = forward_residual_block(bxb, &bw)
807            .to_data()
808            .try_to_vec()
809            .unwrap();
810        let err = max_err(&expected, &got);
811        assert!(err < 1e-4, "residual mismatch: {err:.2e}");
812    }
813
814    #[test]
815    fn prefill_attention_burn_golden() {
816        // Real widths (D_MODEL=1280, 16 heads x 80), 2 patches, random weights.
817        // Exercises split/rope/norm/scale/permute/scores/softmax; golden pins output.
818        let n = 2usize;
819        let mk = |r: usize, c: usize, s: u64| {
820            pseudo(r * c, s)
821                .into_iter()
822                .map(|v| v * 0.02)
823                .collect::<Vec<_>>()
824        };
825
826        let qkv_v = mk(3 * D_MODEL, D_MODEL, 31);
827        let out_v = mk(D_MODEL, D_MODEL, 32);
828        let qn_v = pseudo(HEAD_DIM, 33);
829        let kn_v = pseudo(HEAD_DIM, 34);
830        let qs_v = compute_q_scale(pseudo(HEAD_DIM, 35));
831        let xv = pseudo(n * D_MODEL, 36);
832
833        let bdev = device();
834        let bx = Tensor::<3>::from_data(TensorData::new(xv, [1, n, D_MODEL]), &bdev);
835        let bw = BurnAttnW {
836            qkv_w: Tensor::<2>::from_data(TensorData::new(qkv_v, [3 * D_MODEL, D_MODEL]), &bdev),
837            out_w: Tensor::<2>::from_data(
838                TensorData::new(out_v.clone(), [D_MODEL, D_MODEL]),
839                &bdev,
840            ),
841            q_norm_w: Tensor::<1>::from_data(TensorData::new(qn_v, [HEAD_DIM]), &bdev),
842            k_norm_w: Tensor::<1>::from_data(TensorData::new(kn_v, [HEAD_DIM]), &bdev),
843            q_scale: Tensor::<1>::from_data(TensorData::new(qs_v, [HEAD_DIM]), &bdev),
844            pre_norm_w: Tensor::<1>::from_data(
845                TensorData::new(vec![1.0f32; D_MODEL], [D_MODEL]),
846                &bdev,
847            ),
848            post_norm_w: Tensor::<1>::from_data(
849                TensorData::new(vec![1.0f32; D_MODEL], [D_MODEL]),
850                &bdev,
851            ),
852        };
853        let brope = BurnRope::build(HEAD_DIM, 32, ROPE_THETA);
854        let t0 = std::time::Instant::now();
855        let (got_t, _, _) = burn_prefill_attn(bx, &bw, &brope, n, burn_causal_mask(n));
856        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
857        let got: Vec<f32> = got_t.to_data().try_to_vec().unwrap();
858        assert_eq!(got.len(), n * D_MODEL);
859        assert!(got.iter().all(|v| v.is_finite()));
860        let sum: f32 = got.iter().sum();
861        let head = [got[0], got[1], got[511], got[1024], got[2559]];
862        println!("timesfm attn [n={n}]: burn {burn_ms:.1}ms sum={sum:.6} head={head:?}");
863        // Golden recorded from the verified Burn implementation.
864        let golden_sum = -210195.125f32;
865        let golden_head = [-82.57538, -79.27846, -84.41953, -81.87318, -82.0369];
866        assert!(
867            (sum - golden_sum).abs() < 2.0,
868            "attn sum drift: {sum} vs {golden_sum}"
869        );
870        for (g, e) in head.iter().zip(golden_head.iter()) {
871            assert!((g - e).abs() < 1e-3, "attn head drift: {g} vs {e}");
872        }
873    }
874}
zsfm_timesfm/infer/mod.rs

zsfm_timesfm/infer/
mod.rs
