zsfm_moirai/infer/
mod.rs

1//! Moirai-1.0-R-large inference engine.
2//!
3//! Architecture: masked bidirectional transformer encoder with multi-scale patches.
4//! - Mean-scale normalization → patch embedding (in_proj per patch size) →
5//!   mask tokens for future patches → 24 encoder blocks (RMSNorm + QK-norm
6//!   attention + SwiGLU FFN) → final RMSNorm → Student-t head →
7//!   inverse scale → point forecast (mean = loc)
8
9use std::io::{BufReader, Read, Seek};
10use std::path::Path;
11
12use anyhow::{Context, Result};
13use candle_core::quantized::gguf_file;
14use candle_core::{DType, Device, Tensor, D};
15
16use std::collections::HashMap;
17use std::sync::Mutex;
18
19use rayon::prelude::*;
20
21use crate::config::MoiraiConfig;
22
23// Chosen patch size for inference: 32 (index 2 in [8,16,32,64,128])
24const PATCH_SIZE: usize = 32;
25const PATCH_IDX: usize  = 2;
26
27// ---------------------------------------------------------------------------
28// Weight structs
29// ---------------------------------------------------------------------------
30
31struct EncoderBlock {
32    norm1_w:         Tensor,
33    norm2_w:         Tensor,
34    attn_qkv_w:      Tensor, // fused [3*d_model, d_model]
35    attn_o_w:        Tensor,
36    attn_qn_w:       Tensor, // Q per-head norm scale [head_dim]
37    attn_kn_w:       Tensor, // K per-head norm scale [head_dim]
38    vbias_obs: Tensor,   // [n_heads, 1, 1] — observed-key per-head bias
39    vbias_mask: Tensor,  // [n_heads, 1, 1] — masked-key per-head bias
40    ffn_fc1_w:       Tensor,
41    ffn_fc2_w:       Tensor,
42    ffn_gate_w:      Tensor,
43}
44
45pub struct MoiraiModel {
46    device: Device,
47    config: MoiraiConfig,
48    in_proj_w:     Tensor, // [5, 1024, 128]
49    in_proj_b:     Tensor, // [5, 1024]
50    mask_embed:    Tensor, // [1, 1024]
51    blocks:        Vec<EncoderBlock>,
52    norm_f_w:      Tensor,
53    head_st_loc_w: Tensor, // [5, 128, 1024]
54    head_st_loc_b: Tensor, // [5, 128]
55    rope_inv_freq: Vec<f32>,
56    rope_cache:    Mutex<HashMap<usize, (Tensor, Tensor)>>,
57}
58
59// ---------------------------------------------------------------------------
60// GGUF loading
61// ---------------------------------------------------------------------------
62
63fn load_t(
64    content: &gguf_file::Content,
65    reader: &mut (impl Read + Seek),
66    name: &str,
67    device: &Device,
68) -> Result<Tensor> {
69    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
70}
71
72impl MoiraiModel {
73    pub fn load(gguf_path: &Path, config: MoiraiConfig) -> Result<Self> {
74        let device = Device::Cpu;
75        let file = std::fs::File::open(gguf_path)
76            .with_context(|| format!("open {}", gguf_path.display()))?;
77        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
78        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
79
80        let in_proj_w  = load_t(&content, &mut reader, "in_proj.weight", &device)?;
81        let in_proj_b  = load_t(&content, &mut reader, "in_proj.bias", &device)?;
82        let mask_embed = load_t(&content, &mut reader, "mask_embed.weight", &device)?;
83
84        let mut blocks = Vec::with_capacity(config.n_layers);
85        for n in 0..config.n_layers {
86            let p = |s: &str| format!("blk.{n}.{s}");
87            let q_w = load_t(&content, &mut reader, &p("attn_q.weight"), &device)?;
88            let k_w = load_t(&content, &mut reader, &p("attn_k.weight"), &device)?;
89            let v_w = load_t(&content, &mut reader, &p("attn_v.weight"), &device)?;
90            let attn_qkv_w = Tensor::cat(&[&q_w, &k_w, &v_w], 0)
91                .with_context(|| format!("qkv cat blk.{n}"))?;
92            let norm1_w    = load_t(&content, &mut reader, &p("norm1.weight"), &device)?;
93            let norm2_w    = load_t(&content, &mut reader, &p("norm2.weight"), &device)?;
94            let attn_o_w   = load_t(&content, &mut reader, &p("attn_o.weight"), &device)?;
95            let attn_qn_w  = load_t(&content, &mut reader, &p("attn_qn.weight"), &device)?;
96            let attn_kn_w  = load_t(&content, &mut reader, &p("attn_kn.weight"), &device)?;
97            let vbias_raw = load_t(&content, &mut reader, &p("attn_vbias.weight"), &device)?
98                .flatten_all()?.to_vec1::<f32>()?;
99            let n_heads = config.n_heads;
100            let vbias_obs  = Tensor::from_vec(vbias_raw[0..n_heads].to_vec(), (n_heads, 1, 1), &device)?;
101            let vbias_mask = Tensor::from_vec(vbias_raw[n_heads..2*n_heads].to_vec(), (n_heads, 1, 1), &device)?;
102            let ffn_fc1_w  = load_t(&content, &mut reader, &p("ffn_fc1.weight"), &device)?;
103            let ffn_fc2_w  = load_t(&content, &mut reader, &p("ffn_fc2.weight"), &device)?;
104            let ffn_gate_w = load_t(&content, &mut reader, &p("ffn_gate.weight"), &device)?;
105            blocks.push(EncoderBlock {
106                norm1_w, norm2_w, attn_qkv_w, attn_o_w, attn_qn_w, attn_kn_w,
107                vbias_obs, vbias_mask, ffn_fc1_w, ffn_fc2_w, ffn_gate_w,
108            });
109        }
110
111        let norm_f_w      = load_t(&content, &mut reader, "norm_f.weight", &device)?;
112        let head_st_loc_w = load_t(&content, &mut reader, "head.st_loc.weight", &device)?;
113        let head_st_loc_b = load_t(&content, &mut reader, "head.st_loc.bias", &device)?;
114
115        let head_dim = config.head_dim;
116        let half = head_dim / 2;
117        let rope_inv_freq: Vec<f32> = (0..half)
118            .map(|i| 1.0_f32 / 10000_f32.powf(2.0 * i as f32 / head_dim as f32))
119            .collect();
120
121        Ok(Self {
122            device, config,
123            in_proj_w, in_proj_b, mask_embed,
124            blocks, norm_f_w,
125            head_st_loc_w, head_st_loc_b,
126            rope_inv_freq,
127            rope_cache: Mutex::new(HashMap::new()),
128        })
129    }
130
131    // -----------------------------------------------------------------------
132    // Forecasting
133    // -----------------------------------------------------------------------
134
135    /// Zero-shot forecasting: embed context + masked future, single encoder pass.
136    pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
137        let cfg = &self.config;
138        let patch_size = PATCH_SIZE;
139        let patch_idx  = PATCH_IDX;
140
141        // Mean-scale normalization (Moirai's default scaling)
142        let loc   = context.iter().map(|&v| v as f64).sum::<f64>() / context.len() as f64;
143        let scale = context.iter().map(|&v| (v as f64 - loc).abs()).sum::<f64>()
144                    / context.len() as f64;
145        let scale = (scale.max(1e-8)) as f32;
146        let loc   = loc as f32;
147
148        // Scale context and trim/pad to max_seq_len
149        let max_ts = cfg.max_seq_len; // 512
150        let mut ctx_scaled: Vec<f32> = context.iter().map(|&v| (v - loc) / scale).collect();
151        if ctx_scaled.len() > max_ts {
152            let start = ctx_scaled.len() - max_ts;
153            ctx_scaled = ctx_scaled[start..].to_vec();
154        }
155        // Pad to next multiple of patch_size
156        let ctx_len = ctx_scaled.len();
157        let ctx_padded_len = ((ctx_len + patch_size - 1) / patch_size) * patch_size;
158        if ctx_padded_len > ctx_len {
159            let mut padded = vec![0.0f32; ctx_padded_len - ctx_len];
160            padded.extend_from_slice(&ctx_scaled);
161            ctx_scaled = padded;
162        }
163        let n_ctx_patches = ctx_scaled.len() / patch_size;
164
165        // Number of future patches
166        let n_fc_patches = (horizon + patch_size - 1) / patch_size;
167        let total_patches = n_ctx_patches + n_fc_patches;
168
169        // Embed context patches: in_proj_w[patch_idx] is [1024, 128], use [:, :patch_size]
170        // in_proj shape: [5, 1024, 128] stored as GGUF [128, 1024, 5] (reversed)
171        // After dequantize candle gives us the tensor in its stored order
172        // The Python shape [5, 1024, 128] → GGUF reversal → [128, 1024, 5]
173        // We need to work with this carefully.
174        let d_model      = cfg.d_model;       // 1024
175        let max_ps       = cfg.max_patch_size; // 128
176
177        // Python in_proj.weight[patch_idx]: [d_model=1024, max_ps=128]
178        // In GGUF (reversed dims): the tensor is stored as [max_ps=128, d_model=1024, 5_dim]?
179        // Actually, for a 3D tensor [5, 1024, 128], GGUF stores shape in Python order but
180        // Candle dequantizes preserving the original shape. Let's reshape to [5, 1024, 128].
181        let in_proj_w_3d = self.in_proj_w.reshape((5, d_model, max_ps))?;
182        // Extract patch_idx slice: [d_model, max_ps]
183        let proj_w_slice = in_proj_w_3d.get(patch_idx)?.contiguous()?; // [d_model, max_ps]
184        // Take only patch_size columns: [d_model, patch_size]
185        let proj_w = proj_w_slice.narrow(1, 0, patch_size)?.contiguous()?;
186
187        // Bias: in_proj_b[patch_idx] → [d_model=1024]
188        let in_proj_b_2d = self.in_proj_b.reshape((5, d_model))?;
189        let proj_b = in_proj_b_2d.get(patch_idx)?.contiguous()?; // [d_model]
190
191        // Build patch embeddings for context: [n_ctx_patches, d_model]
192        let mut patch_flat = vec![0.0f32; n_ctx_patches * patch_size];
193        for i in 0..n_ctx_patches {
194            let src = &ctx_scaled[i * patch_size..(i + 1) * patch_size];
195            patch_flat[i * patch_size..(i + 1) * patch_size].copy_from_slice(src);
196        }
197        let ctx_patches_t = Tensor::from_vec(
198            patch_flat, (n_ctx_patches, patch_size), &self.device,
199        )?;
200        // [n_ctx, patch_size] @ [patch_size, d_model] + [d_model] → [n_ctx, d_model]
201        let ctx_emb = ctx_patches_t
202            .matmul(&proj_w.t()?)?
203            .broadcast_add(&proj_b)?;
204
205        // Mask embeddings for future patches: tile mask_embed [n_fc, d_model]
206        // mask_embed shape in GGUF: [1024, 1] (Python [1, 1024] → reversed)
207        let mask_embed = self.mask_embed.reshape((1, d_model))?;
208        let fc_emb = mask_embed.expand((n_fc_patches, d_model))?;
209
210        // Concatenate: [total_patches, d_model]
211        let mut h = Tensor::cat(&[&ctx_emb, &fc_emb], 0)?;
212
213        // Build is_masked indicator for var_attn_bias: [total_patches]
214        // 0 = context, 1 = masked future
215        let mut is_masked = vec![0u8; total_patches];
216        for i in n_ctx_patches..total_patches {
217            is_masked[i] = 1;
218        }
219
220        // Encoder
221        for blk in &self.blocks {
222            h = self.forward_block(&h, blk, &is_masked, total_patches)?;
223        }
224        h = zsfm_nn::rms_norm(&h, Some(&self.norm_f_w), 1e-6)?;
225
226        // Apply Student-t loc head to future positions
227        // head_st_loc_w shape: Python [5, 128, 1024] → GGUF [1024, 128, 5]
228        // After reshape: [5, 128, 1024]
229        let loc_w_3d = self.head_st_loc_w.reshape((5, max_ps, d_model))?;
230        let loc_b_2d = self.head_st_loc_b.reshape((5, max_ps))?;
231        let loc_w = loc_w_3d.get(patch_idx)?.narrow(0, 0, patch_size)?.contiguous()?; // [patch_size, d_model]
232        let loc_b = loc_b_2d.get(patch_idx)?.narrow(0, 0, patch_size)?.contiguous()?; // [patch_size]
233
234        let future_h = h.narrow(0, n_ctx_patches, n_fc_patches)?; // [n_fc, d_model]
235        // [n_fc, d_model] @ [d_model, patch_size] + [patch_size] → [n_fc, patch_size]
236        let pred = future_h.matmul(&loc_w.t()?)?.broadcast_add(&loc_b)?;
237
238        let pred_flat: Vec<f32> = pred.flatten_all()?.to_vec1()?;
239        let result: Vec<f32> = pred_flat
240            .iter()
241            .take(horizon)
242            .map(|&v| v * scale + loc)
243            .collect();
244
245        Ok(result)
246    }
247
248    fn forward_block(
249        &self,
250        hidden: &Tensor,
251        blk: &EncoderBlock,
252        is_masked: &[u8],
253        seq_len: usize,
254    ) -> Result<Tensor> {
255        let res = hidden;
256        let h = zsfm_nn::rms_norm(hidden, Some(&blk.norm1_w), 1e-6)?;
257        let h = self.qk_attn(&h, blk, is_masked, seq_len)?;
258        let h = (h + res)?;
259
260        let res2 = h.clone();
261        let h2 = zsfm_nn::rms_norm(&h, Some(&blk.norm2_w), 1e-6)?;
262        let h2 = zsfm_nn::swiglu_ffn(&h2, &blk.ffn_fc1_w, &blk.ffn_fc2_w, &blk.ffn_gate_w)?;
263        Ok((h2 + res2)?)
264    }
265
266    fn qk_attn(
267        &self,
268        hidden: &Tensor,
269        blk: &EncoderBlock,
270        is_masked: &[u8],
271        seq_len: usize,
272    ) -> Result<Tensor> {
273        let cfg = &self.config;
274        let n_heads  = cfg.n_heads;
275        let head_dim = cfg.head_dim;
276        let d_model  = cfg.d_model;
277
278        let qkv = zsfm_nn::linear_nobias(hidden, &blk.attn_qkv_w)?;
279        let q = qkv.narrow(D::Minus1, 0, d_model)?;
280        let k = qkv.narrow(D::Minus1, d_model, d_model)?;
281        let v = qkv.narrow(D::Minus1, 2 * d_model, d_model)?;
282
283        // Reshape: [seq, n_heads, head_dim]
284        let q = q.reshape((seq_len, n_heads, head_dim))?;
285        let k = k.reshape((seq_len, n_heads, head_dim))?;
286
287        // QK per-head norm then RoPE
288        let q = qk_norm_heads(&q, &blk.attn_qn_w, seq_len, n_heads, head_dim)?;
289        let k = qk_norm_heads(&k, &blk.attn_kn_w, seq_len, n_heads, head_dim)?;
290
291        // [seq, n_heads, head_dim] → [n_heads, seq, head_dim]
292        let q = q.permute((1, 0, 2))?.contiguous()?;
293        let k = k.permute((1, 0, 2))?.contiguous()?;
294
295        // Apply RoPE for positional encoding (cached by seq_len)
296        let (cos_t, sin_t) = {
297            let mut cache = self.rope_cache.lock().unwrap();
298            if !cache.contains_key(&seq_len) {
299                let half = self.rope_inv_freq.len();
300                let mut cos_v = vec![0.0f32; seq_len * half];
301                let mut sin_v = vec![0.0f32; seq_len * half];
302                for pos in 0..seq_len {
303                    for i in 0..half {
304                        let theta = pos as f32 * self.rope_inv_freq[i];
305                        cos_v[pos * half + i] = theta.cos();
306                        sin_v[pos * half + i] = theta.sin();
307                    }
308                }
309                let half_dim = self.rope_inv_freq.len();
310                let cos_t = Tensor::from_vec(cos_v, (seq_len, half_dim), &self.device)?.unsqueeze(0)?;
311                let sin_t = Tensor::from_vec(sin_v, (seq_len, half_dim), &self.device)?.unsqueeze(0)?;
312                cache.insert(seq_len, (cos_t, sin_t));
313            }
314            let (c, s) = &cache[&seq_len];
315            (c.clone(), s.clone())
316        };
317
318        let q = apply_rope_with_tables(&q, &cos_t, &sin_t, head_dim)?;
319        let k = apply_rope_with_tables(&k, &cos_t, &sin_t, head_dim)?;
320        let v = v.reshape((seq_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
321
322        let scale = (head_dim as f64).sqrt();
323        let scores = q.matmul(&k.permute((0, 2, 1))?)?;  // [n_heads, seq, seq]
324        let scores = (scores / scale)?;
325
326        // Add variate attention bias using precomputed vbias data
327        let scores = apply_var_attn_bias(&scores, is_masked, seq_len, &blk.vbias_obs, &blk.vbias_mask, &self.device)?;
328
329        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
330        let out = attn.matmul(&v)?; // [n_heads, seq, head_dim]
331        let out = out.permute((1, 0, 2))?.contiguous()?.reshape((seq_len, d_model))?;
332
333        zsfm_nn::linear_nobias(&out, &blk.attn_o_w)
334    }
335}
336
337// ---------------------------------------------------------------------------
338// Ops
339// ---------------------------------------------------------------------------
340
341/// QK per-head RMSNorm: normalize each head independently.
342/// x: [seq, n_heads, head_dim] → normalize over head_dim per head
343fn qk_norm_heads(
344    x: &Tensor,
345    weight: &Tensor, // [head_dim]
346    seq_len: usize,
347    n_heads: usize,
348    head_dim: usize,
349) -> Result<Tensor> {
350    let x_flat = x.reshape((seq_len * n_heads, head_dim))?;
351    let normed = zsfm_nn::rms_norm(&x_flat, Some(weight), 1e-6)?;
352    Ok(normed.reshape((seq_len, n_heads, head_dim))?)
353}
354
355/// Add variate attention bias to scores [n_heads, seq, seq].
356/// Build [n_heads, 1, seq_len] variate bias via per-key mask selection:
357///   bias[head, key] = vbias_mask[head] if is_masked[key] else vbias_obs[head]
358///                   = vbias_obs[head] + (vbias_mask[head] - vbias_obs[head]) * m[key]
359fn apply_var_attn_bias(
360    scores: &Tensor,
361    is_masked: &[u8],
362    seq_len: usize,
363    vbias_obs: &Tensor,   // [n_heads, 1, 1]
364    vbias_mask: &Tensor,  // [n_heads, 1, 1]
365    device: &Device,
366) -> Result<Tensor> {
367    let mask_f: Vec<f32> = is_masked.iter().map(|&m| m as f32).collect();
368    let mask_t = Tensor::from_vec(mask_f, (1usize, 1, seq_len), device)?;
369    let delta = (vbias_mask - vbias_obs)?;
370    let bias = vbias_obs.broadcast_add(&delta.broadcast_mul(&mask_t)?)?;
371    Ok(scores.broadcast_add(&bias)?)
372}
373
374/// RoPE using precomputed cos/sin tables [1, seq, half].
375/// x is [n_heads, seq, head_dim].
376fn apply_rope_with_tables(
377    x: &Tensor,
378    cos_t: &Tensor,
379    sin_t: &Tensor,
380    head_dim: usize,
381) -> Result<Tensor> {
382    let half = head_dim / 2;
383    let x1 = x.narrow(D::Minus1, 0, half)?.contiguous()?;
384    let x2 = x.narrow(D::Minus1, half, half)?.contiguous()?;
385    let rot1 = (x1.broadcast_mul(cos_t)? - x2.broadcast_mul(sin_t)?)?;
386    let rot2 = (x1.broadcast_mul(sin_t)? + x2.broadcast_mul(cos_t)?)?;
387    Ok(Tensor::cat(&[&rot1, &rot2], D::Minus1)?.contiguous()?)
388}
389
390// ---------------------------------------------------------------------------
391// zsfm-core::Forecaster
392// ---------------------------------------------------------------------------
393
394impl zsfm_core::Forecaster for MoiraiModel {
395    type Config = MoiraiConfig;
396
397    fn load(gguf_path: &Path, config: MoiraiConfig) -> Result<Self> {
398        MoiraiModel::load(gguf_path, config)
399    }
400
401    /// Moirai is channel-independent: each variate is forecast separately (no cross-variate
402    /// attention), point-forecast only, so `mask` is unused. Returns a single "quantile" row
403    /// (the point forecast) with one column per variate. Variates are forecast in parallel
404    /// via `rayon` — each is a fully independent forward pass through the same read-only
405    /// weights (the only shared mutable state, `rope_cache`, is behind a `Mutex`).
406    fn forecast(
407        &self,
408        context: &[Vec<f32>],
409        _mask: &[Vec<bool>],
410        horizon: usize,
411    ) -> Result<zsfm_core::QuantileMatrix> {
412        anyhow::ensure!(!context.is_empty(), "context must have at least one variate");
413        let variates: Vec<Vec<f32>> = context
414            .par_iter()
415            .enumerate()
416            .map(|(vi, ctx)| -> Result<Vec<f32>> {
417                anyhow::ensure!(!ctx.is_empty(), "variate {vi} context must not be empty");
418                MoiraiModel::forecast(self, ctx, horizon).with_context(|| format!("forecast variate {vi}"))
419            })
420            .collect::<Result<Vec<_>>>()?;
421        Ok(vec![variates])
422    }
423}
zsfm_moirai/infer/mod.rs

zsfm_moirai/infer/
mod.rs
