zsfm_sundial/infer/
mod.rs

1pub mod rope;
2
3use anyhow::Result;
4use candle_core::{DType, Device, Tensor, D};
5use candle_core::quantized::gguf_file;
6use candle_nn::ops;
7use std::collections::HashMap;
8use std::sync::Mutex;
9use std::fs::File;
10use std::io::BufReader;
11use std::path::{Path, PathBuf};
12
13use rope::RopeCache;
14use zsfm_nn::linear;
15
16const N_HEADS: usize = 12;
17const HEAD_DIM: usize = 64;
18const EMBED_IN: usize = 32;  // 2 * input_token_len
19const PATCH_SIZE: usize = 16;
20const TIME_DIM: usize = 256;
21const MAX_SEQ: usize = 4096;
22const ROPE_THETA: f64 = 10000.0;
23
24struct EmbedW {
25    hidden_w: Tensor, hidden_b: Tensor,
26    output_w: Tensor, output_b: Tensor,
27    skip_w: Tensor,   skip_b: Tensor,
28}
29
30struct AttnW {
31    qkv_w: Tensor, qkv_b: Tensor,
32    o_w: Tensor,
33}
34
35struct NormW { w: Tensor, b: Tensor }
36
37struct BlockW {
38    attn: AttnW,
39    attn_norm: NormW,
40    ffn_norm: NormW,
41    gate_w: Tensor,
42    up_w: Tensor,
43    down_w: Tensor,
44}
45
46struct FlowResW {
47    ln: NormW,
48    mlp1_w: Tensor, mlp1_b: Tensor,
49    mlp2_w: Tensor, mlp2_b: Tensor,
50    adaln_w: Tensor, adaln_b: Tensor,
51}
52
53struct FlowW {
54    t1_w: Tensor, t1_b: Tensor,
55    t2_w: Tensor, t2_b: Tensor,
56    cond_w: Tensor, cond_b: Tensor,
57    in_w: Tensor, in_b: Tensor,
58    res: Vec<FlowResW>,
59    out_w: Tensor, out_b: Tensor,
60    out_adaln_w: Tensor, out_adaln_b: Tensor,
61}
62
63/// Sundial reads its own architecture (layer count, flow depth, output length, sampling steps)
64/// straight out of the GGUF metadata written at conversion time — there is no separate
65/// `InferConfig`/`config.json` needed for inference (unlike every other model in the workspace).
66/// [`SundialModelBuilder`] only exposes the two things a caller can actually tune at load time:
67/// the compute device and an ODE step-count override.
68pub struct SundialModel {
69    device: Device,
70    embed: EmbedW,
71    blocks: Vec<BlockW>,
72    norm: NormW,
73    flow: FlowW,
74    rope: RopeCache,
75    n_steps: usize,
76    output_len: usize,
77    // Precomputed: sinusoidal_embed → t1_proj → silu → t2_proj for steps 0..=n_steps
78    t_emb_table: Vec<Tensor>,
79    causal_mask_cache: Mutex<HashMap<usize, Tensor>>,
80}
81
82// ---------------------------------------------------------------------------
83// Builder
84// ---------------------------------------------------------------------------
85
86/// Fluent constructor for [`SundialModel`].
87///
88/// ```no_run
89/// use zsfm_sundial::SundialModel;
90///
91/// # fn main() -> anyhow::Result<()> {
92/// let model = SundialModel::builder("sundial.gguf").steps(20).build()?;
93/// # Ok(()) }
94/// ```
95pub struct SundialModelBuilder {
96    gguf_path: PathBuf,
97    device: Device,
98    steps_override: Option<usize>,
99}
100
101impl SundialModelBuilder {
102    fn new(gguf_path: impl Into<PathBuf>) -> Self {
103        Self { gguf_path: gguf_path.into(), device: Device::Cpu, steps_override: None }
104    }
105
106    /// Override the GGUF-embedded ODE step count (default: use model metadata, typically 50).
107    /// 10-20 steps recommended; latency scales linearly with this value.
108    pub fn steps(mut self, n: usize) -> Self {
109        self.steps_override = Some(n);
110        self
111    }
112
113    pub fn device(mut self, device: Device) -> Self {
114        self.device = device;
115        self
116    }
117
118    pub fn build(self) -> Result<SundialModel> {
119        SundialModel::load(&self.gguf_path, &self.device, self.steps_override)
120    }
121}
122
123fn get_u32(content: &gguf_file::Content, key: &str) -> Option<u32> {
124    match content.metadata.get(key) {
125        Some(gguf_file::Value::U32(v)) => Some(*v),
126        Some(gguf_file::Value::U64(v)) => Some(*v as u32),
127        _ => None,
128    }
129}
130
131fn layer_norm(x: &Tensor, nw: &NormW) -> Result<Tensor> {
132    zsfm_nn::layer_norm(x, &nw.w, &nw.b, 1e-5)
133}
134
135fn layer_norm_no_params(x: &Tensor) -> Result<Tensor> {
136    let mean = x.mean_keepdim(D::Minus1)?;
137    let diff = x.broadcast_sub(&mean)?;
138    let var = diff.sqr()?.mean_keepdim(D::Minus1)?;
139    let std = var.affine(1.0, 1e-5)?.sqrt()?;
140    Ok(diff.broadcast_div(&std)?)
141}
142
143fn silu(x: &Tensor) -> Result<Tensor> {
144    ops::silu(x).map_err(anyhow::Error::from)
145}
146
147fn make_causal_mask(seq: usize, device: &Device) -> Result<Tensor> {
148    zsfm_nn::make_causal_mask(seq, 0, device)
149}
150
151fn sinusoidal_embed(t: f32, dim: usize) -> Vec<f32> {
152    let half = dim / 2;
153    let freqs: Vec<f32> = (0..half)
154        .map(|i| (-(10000.0f32.ln() * i as f32 / half as f32)).exp())
155        .collect();
156    let args: Vec<f32> = freqs.iter().map(|&f| t * f).collect();
157    let mut emb = Vec::with_capacity(dim);
158    for &a in &args { emb.push(a.cos()); }
159    for &a in &args { emb.push(a.sin()); }
160    emb
161}
162
163impl SundialModel {
164    /// Start building a [`SundialModel`] — see [`SundialModelBuilder`].
165    pub fn builder(gguf_path: impl Into<PathBuf>) -> SundialModelBuilder {
166        SundialModelBuilder::new(gguf_path)
167    }
168
169    pub fn load(path: &Path, device: &Device, steps_override: Option<usize>) -> Result<Self> {
170        let f = File::open(path)
171            .map_err(|e| anyhow::anyhow!("open {}: {}", path.display(), e))?;
172        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, f);
173        let content = gguf_file::Content::read(&mut reader)
174            .map_err(|e| anyhow::anyhow!("read gguf: {}", e))?;
175
176        let n_layers = get_u32(&content, "sundial1.block_count").unwrap_or(12) as usize;
177        let gguf_n_steps = get_u32(&content, "sundial1.flow.num_sampling_steps").unwrap_or(50) as usize;
178        let n_steps = steps_override.unwrap_or(gguf_n_steps);
179        let flow_depth = get_u32(&content, "sundial1.flow.depth").unwrap_or(3) as usize;
180        let output_len = get_u32(&content, "sundial1.output_token_len").unwrap_or(720) as usize;
181
182        macro_rules! lt {
183            ($name:expr) => {{
184                let name: &str = $name;
185                let qt = content.tensor(&mut reader, name, device)
186                    .map_err(|e| anyhow::anyhow!("load {}: {}", name, e))?;
187                qt.dequantize(device)
188                    .map_err(|e| anyhow::anyhow!("dequantize {}: {}", name, e))?
189            }};
190        }
191
192        let embed = EmbedW {
193            hidden_w: lt!("embed.hidden.weight"),
194            hidden_b: lt!("embed.hidden.bias"),
195            output_w: lt!("embed.output.weight"),
196            output_b: lt!("embed.output.bias"),
197            skip_w:   lt!("embed.skip.weight"),
198            skip_b:   lt!("embed.skip.bias"),
199        };
200
201        let mut blocks = Vec::with_capacity(n_layers);
202        for n in 0..n_layers {
203            let p = |s: &str| format!("blk.{n}.{s}");
204            let q_w = lt!(&p("attn_q.weight"));
205            let k_w = lt!(&p("attn_k.weight"));
206            let v_w = lt!(&p("attn_v.weight"));
207            let q_b = lt!(&p("attn_q.bias"));
208            let k_b = lt!(&p("attn_k.bias"));
209            let v_b = lt!(&p("attn_v.bias"));
210            let qkv_w = Tensor::cat(&[&q_w, &k_w, &v_w], 0)
211                .map_err(|e| anyhow::anyhow!("qkv_w cat blk.{n}: {e}"))?;
212            let qkv_b = Tensor::cat(&[&q_b, &k_b, &v_b], 0)
213                .map_err(|e| anyhow::anyhow!("qkv_b cat blk.{n}: {e}"))?;
214            blocks.push(BlockW {
215                attn: AttnW {
216                    qkv_w,
217                    qkv_b,
218                    o_w: lt!(&p("attn_out.weight")),
219                },
220                attn_norm: NormW {
221                    w: lt!(&p("attn_norm.weight")), b: lt!(&p("attn_norm.bias")),
222                },
223                ffn_norm: NormW {
224                    w: lt!(&p("ffn_norm.weight")), b: lt!(&p("ffn_norm.bias")),
225                },
226                gate_w: lt!(&p("ffn_gate.weight")),
227                up_w:   lt!(&p("ffn_up.weight")),
228                down_w: lt!(&p("ffn_down.weight")),
229            });
230        }
231
232        let norm = NormW {
233            w: lt!("norm.weight"),
234            b: lt!("norm.bias"),
235        };
236
237        let mut res_blocks = Vec::with_capacity(flow_depth);
238        for k in 0..flow_depth {
239            let fp = |s: &str| format!("flow.res.{k}.{s}");
240            res_blocks.push(FlowResW {
241                ln: NormW { w: lt!(&fp("ln.weight")), b: lt!(&fp("ln.bias")) },
242                mlp1_w: lt!(&fp("mlp1.weight")), mlp1_b: lt!(&fp("mlp1.bias")),
243                mlp2_w: lt!(&fp("mlp2.weight")), mlp2_b: lt!(&fp("mlp2.bias")),
244                adaln_w: lt!(&fp("adaln.weight")), adaln_b: lt!(&fp("adaln.bias")),
245            });
246        }
247
248        let flow = FlowW {
249            t1_w: lt!("flow.t_proj1.weight"), t1_b: lt!("flow.t_proj1.bias"),
250            t2_w: lt!("flow.t_proj2.weight"), t2_b: lt!("flow.t_proj2.bias"),
251            cond_w: lt!("flow.cond.weight"),  cond_b: lt!("flow.cond.bias"),
252            in_w:  lt!("flow.in_proj.weight"), in_b: lt!("flow.in_proj.bias"),
253            res: res_blocks,
254            out_w:      lt!("flow.out_linear.weight"),
255            out_b:      lt!("flow.out_linear.bias"),
256            out_adaln_w: lt!("flow.out_adaln.weight"),
257            out_adaln_b: lt!("flow.out_adaln.bias"),
258        };
259
260        let rope = RopeCache::new(HEAD_DIM, MAX_SEQ, ROPE_THETA, device)?;
261
262        // Precompute time embeddings for all ODE steps (sinusoidal → t1_proj → silu → t2_proj)
263        let t_emb_table: Vec<Tensor> = (0..=n_steps)
264            .map(|i| {
265                let t_scaled = i as f32 / n_steps as f32 * 1000.0;
266                let t_raw = Tensor::from_vec(sinusoidal_embed(t_scaled, TIME_DIM), (1, TIME_DIM), device)?;
267                let t_h = silu(&linear(&t_raw, &flow.t1_w, Some(&flow.t1_b))?)?;
268                linear(&t_h, &flow.t2_w, Some(&flow.t2_b))
269            })
270            .collect::<Result<Vec<_>>>()
271            .map_err(|e| anyhow::anyhow!("t_emb_table: {e}"))?;
272
273        Ok(Self {
274            device: device.clone(),
275            embed, blocks, norm, flow, rope, n_steps, output_len, t_emb_table,
276            causal_mask_cache: Mutex::new(HashMap::new()),
277        })
278    }
279
280    pub fn forecast(&self, context: &[f32], device: &Device) -> Result<Vec<f32>> {
281        let n = context.len();
282        // ReVIN: whole-series population mean/std
283        let mean = context.iter().sum::<f32>() / n as f32;
284        let var = context.iter().map(|&x| (x - mean) * (x - mean)).sum::<f32>() / n as f32;
285        let std = (var + 1e-5).sqrt();
286        let normed: Vec<f32> = context.iter().map(|&x| (x - mean) / std).collect();
287
288        // Pad to multiple of patch_size and build embedding input [values | mask]
289        let n_patches = (n + PATCH_SIZE - 1) / PATCH_SIZE;
290        let mut embed_in = vec![0.0f32; n_patches * EMBED_IN];
291        for p in 0..n_patches {
292            let start = p * PATCH_SIZE;
293            let end = (start + PATCH_SIZE).min(n);
294            for i in start..end {
295                embed_in[p * EMBED_IN + (i - start)] = normed[i];
296                embed_in[p * EMBED_IN + PATCH_SIZE + (i - start)] = 1.0;
297            }
298        }
299
300        let x = Tensor::from_vec(embed_in, (1, n_patches, EMBED_IN), device)?;
301        let mut h = self.embed_forward(&x)?;
302
303        let mask = {
304            let mut cache = self.causal_mask_cache.lock().unwrap();
305            if !cache.contains_key(&n_patches) {
306                cache.insert(n_patches, make_causal_mask(n_patches, &self.device)?);
307            }
308            cache[&n_patches].clone()
309        };
310        for block in &self.blocks {
311            h = self.block_forward(&h, block, &mask)?;
312        }
313        h = layer_norm(&h, &self.norm)?;
314
315        // Last patch hidden state as flow condition
316        let cond = h.narrow(1, n_patches - 1, 1)?.squeeze(1)?;  // [1, hidden]
317
318        // Flow matching: Euler from t=0 (noise) to t=1 (data)
319        let output = self.flow_sample(&cond, device)?;
320
321        let vals = output.to_vec2::<f32>()?;
322        Ok(vals[0].iter().map(|&v| v * std + mean).collect())
323    }
324
325    fn embed_forward(&self, x: &Tensor) -> Result<Tensor> {
326        let h = silu(&linear(x, &self.embed.hidden_w, Some(&self.embed.hidden_b))?)?;
327        let out = linear(&h, &self.embed.output_w, Some(&self.embed.output_b))?;
328        let skip = linear(x, &self.embed.skip_w, Some(&self.embed.skip_b))?;
329        Ok((out + skip)?)
330    }
331
332    fn block_forward(&self, x: &Tensor, blk: &BlockW, mask: &Tensor) -> Result<Tensor> {
333        let h = layer_norm(x, &blk.attn_norm)?;
334        let h = self.attn_forward(&h, &blk.attn, mask)?;
335        let x = (x + &h)?;
336        let h = layer_norm(&x, &blk.ffn_norm)?;
337        let h = self.ffn_forward(&h, blk)?;
338        Ok((x + h)?)
339    }
340
341    fn attn_forward(&self, x: &Tensor, attn: &AttnW, mask: &Tensor) -> Result<Tensor> {
342        let (b, seq, _) = x.dims3()?;
343        let d = N_HEADS * HEAD_DIM;
344        let qkv = linear(x, &attn.qkv_w, Some(&attn.qkv_b))?;
345        let q = qkv.narrow(D::Minus1, 0, d)?.contiguous()?;
346        let k = qkv.narrow(D::Minus1, d, d)?.contiguous()?;
347        let v = qkv.narrow(D::Minus1, 2 * d, d)?.contiguous()?;
348
349        // [b, seq, H*D] → [b, H, seq, D]
350        let q = q.reshape((b, seq, N_HEADS, HEAD_DIM))?.permute((0, 2, 1, 3))?.contiguous()?;
351        let k = k.reshape((b, seq, N_HEADS, HEAD_DIM))?.permute((0, 2, 1, 3))?.contiguous()?;
352        let v = v.reshape((b, seq, N_HEADS, HEAD_DIM))?.permute((0, 2, 1, 3))?.contiguous()?;
353
354        let q = self.rope.apply(&q, 0)?;
355        let k = self.rope.apply(&k, 0)?;
356
357        let scale = 1.0 / (HEAD_DIM as f64).sqrt();
358        let scores = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?.affine(scale, 0.0)?;
359        let scores = scores.broadcast_add(mask)?;
360        let aw = ops::softmax(&scores, D::Minus1)?;
361
362        let out = aw.matmul(&v)?;
363        let out = out.permute((0, 2, 1, 3))?.reshape((b, seq, N_HEADS * HEAD_DIM))?;
364        linear(&out, &attn.o_w, None)
365    }
366
367    fn ffn_forward(&self, x: &Tensor, blk: &BlockW) -> Result<Tensor> {
368        let gate = silu(&linear(x, &blk.gate_w, None)?)?;
369        let up = linear(x, &blk.up_w, None)?;
370        let h = (gate * &up)?;
371        linear(&h, &blk.down_w, None)
372    }
373
374    // Evaluate the flow network at a given state x and timestep-conditioned activations c_act.
375    fn flow_net_eval(&self, x: &Tensor, c_act: &Tensor) -> Result<Tensor> {
376        let mut h = linear(x, &self.flow.in_w, Some(&self.flow.in_b))?;
377        for res in &self.flow.res {
378            let adaln = linear(c_act, &res.adaln_w, Some(&res.adaln_b))?;
379            let third = adaln.dim(D::Minus1)? / 3;
380            let shift = adaln.narrow(D::Minus1, 0, third)?;
381            let scale = adaln.narrow(D::Minus1, third, third)?;
382            let gate  = adaln.narrow(D::Minus1, 2 * third, third)?;
383            let h_norm = layer_norm(&h, &res.ln)?;
384            let h_mod = h_norm.broadcast_mul(&scale.affine(1.0, 1.0)?)?.broadcast_add(&shift)?;
385            let h_mlp = silu(&linear(&h_mod, &res.mlp1_w, Some(&res.mlp1_b))?)?;
386            let h_mlp = linear(&h_mlp, &res.mlp2_w, Some(&res.mlp2_b))?;
387            h = (h + (gate * h_mlp)?)?;
388        }
389        let adaln = linear(c_act, &self.flow.out_adaln_w, Some(&self.flow.out_adaln_b))?;
390        let half = adaln.dim(D::Minus1)? / 2;
391        let shift = adaln.narrow(D::Minus1, 0, half)?;
392        let scale = adaln.narrow(D::Minus1, half, half)?;
393        let h_norm = layer_norm_no_params(&h)?;
394        let h_mod = h_norm.broadcast_mul(&scale.affine(1.0, 1.0)?)?.broadcast_add(&shift)?;
395        linear(&h_mod, &self.flow.out_w, Some(&self.flow.out_b))
396    }
397
398    // Heun's method (2nd-order Runge-Kutta) for the flow ODE.
399    // Each step does 2 network evals but converges quadratically, so n_steps/2 fewer steps
400    // are needed compared to Euler for equivalent accuracy. For OT-FM (linear paths) the
401    // default 50-step Euler run can be replaced by ~10 Heun steps with near-identical output.
402    fn flow_sample(&self, cond: &Tensor, device: &Device) -> Result<Tensor> {
403        let cond_emb = linear(cond, &self.flow.cond_w, Some(&self.flow.cond_b))?;
404        let dt = 1.0f64 / self.n_steps as f64;
405        let mut x = Tensor::zeros((1, self.output_len), DType::F32, device)?;
406
407        // c[i] = silu(t_emb_table[i] + cond_emb); c2 at step i == c1 at step i+1,
408        // so compute once and carry forward to halve the add+silu cost.
409        let mut c_cur = silu(&(&self.t_emb_table[0] + &cond_emb)?)?;
410        for i in 0..self.n_steps {
411            let k1 = self.flow_net_eval(&x, &c_cur)?;
412
413            let x_pred = (x.clone() + k1.affine(dt, 0.0)?)?;
414
415            let c_next = silu(&(&self.t_emb_table[i + 1] + &cond_emb)?)?;
416            let k2 = self.flow_net_eval(&x_pred, &c_next)?;
417
418            x = (x + ((&k1 + &k2)? * (dt * 0.5))?)?;
419            c_cur = c_next;
420        }
421        Ok(x)
422    }
423}
424
425// ---------------------------------------------------------------------------
426// zsfm-core::Forecaster
427// ---------------------------------------------------------------------------
428
429impl zsfm_core::Forecaster for SundialModel {
430    /// No config.json is consulted for inference — see the struct docs.
431    type Config = ();
432
433    fn load(gguf_path: &Path, _config: ()) -> Result<Self> {
434        SundialModel::load(gguf_path, &Device::Cpu, None)
435    }
436
437    /// Sundial is univariate-only and point-forecast-only (flow-matching samples a single path,
438    /// not a quantile distribution) — `mask` is unused and the returned matrix has exactly one
439    /// row (the point forecast) and one variate.
440    fn forecast(
441        &self,
442        context: &[Vec<f32>],
443        _mask: &[Vec<bool>],
444        horizon: usize,
445    ) -> Result<zsfm_core::QuantileMatrix> {
446        anyhow::ensure!(context.len() == 1, "SundialModel only supports univariate forecasting (1 variate)");
447        let raw = SundialModel::forecast(self, &context[0], &self.device)?;
448        let point: Vec<f32> = raw.into_iter().take(horizon).collect();
449        Ok(vec![vec![point]])
450    }
451}
zsfm_sundial/infer/mod.rs

zsfm_sundial/infer/
mod.rs
