zsfm_flowstate/infer/
mod.rs

1//! FlowState-R1 inference on the Burn backend (Flex CPU) for tensor
2//! stages, with the original verbatim host kernels for the SSM core.
3//!
4//! Projections (embed, B/C, decoder, basis), activations and norms run on
5//! Burn tensors; the diagonal SSM scan (`ssm_scan_step`), discretization
6//! math, RevIN, Legendre basis and quantile recalibration are engine-agnostic
7//! host code.
8//! Weight loading reads GGUF straight via [`zsfm_burn::GgufWeights`].
9
10use std::collections::HashMap;
11use std::path::Path;
12use std::sync::Mutex;
13
14use anyhow::{Context, Result};
15use burn::tensor::{activation, Device, Tensor, TensorData};
16use simdeez::prelude::*;
17use zsfm_burn::linear::burn_linear_nd;
18use zsfm_burn::norm::burn_layer_norm_nd;
19use zsfm_burn::norm::burn_rms_norm_nd;
20use zsfm_burn::{open_gguf_weights, GgufWeights};
21
22use crate::config::FlowStateConfig;
23
24// ---------------------------------------------------------------------------
25// Config
26// ---------------------------------------------------------------------------
27
28#[derive(Debug, Clone)]
29pub struct InferConfig {
30    pub num_layers: usize,
31    pub embed_dim: usize,
32    pub state_dim: usize,
33    pub n_inputs: usize, // 2 for with_missing (value + mask)
34    pub decoder_dim: usize,
35    pub decoder_patch_len: usize,
36    pub quantiles: Vec<f32>,
37    pub basis_range: [f32; 2], // e.g. [-1.0, 0.95] for "legs"
38    pub context_length: usize,
39    pub eps: f32,
40}
41
42/// `FlowStateConfig` already resolved every default (via serde) or failed to parse if a
43/// genuinely required field was missing, so this mapping is infallible and needs no further
44/// defaulting of its own.
45impl From<&FlowStateConfig> for InferConfig {
46    fn from(c: &FlowStateConfig) -> Self {
47        InferConfig {
48            num_layers: c.encoder_num_layers as usize,
49            embed_dim: c.embedding_feature_dim as usize,
50            state_dim: c.encoder_state_dim as usize,
51            n_inputs: c.n_inputs() as usize,
52            decoder_dim: c.decoder_dim as usize,
53            decoder_patch_len: c.decoder_patch_len as usize,
54            quantiles: c.quantiles.clone(),
55            basis_range: c.basis_range(),
56            context_length: c.context_length as usize,
57            eps: 1e-5,
58        }
59    }
60}
61
62impl InferConfig {
63    pub fn quantiles(&self) -> &[f32] {
64        &self.quantiles
65    }
66    pub fn median_index(&self) -> usize {
67        self.quantiles
68            .iter()
69            .position(|&q| (q - 0.5).abs() < 1e-6)
70            .unwrap_or(self.quantiles.len() / 2)
71    }
72}
73
74fn device() -> Device {
75    Device::flex()
76}
77
78fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
79    weights.load_2d(name)
80}
81
82fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
83    weights.load_1d(name)
84}
85
86fn load_f32_vec(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
87    weights.load_flat(name)
88}
89
90fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
91    t.to_data()
92        .try_to_vec()
93        .map_err(|e| anyhow::anyhow!("{e:?}"))
94}
95
96fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
97    burn_linear_nd(x, w, b)
98}
99
100fn layer_norm<const D: usize>(x: Tensor<D>, w: Tensor<1>, b: Tensor<1>, eps: f32) -> Tensor<D> {
101    burn_layer_norm_nd(x, w, b, eps)
102}
103
104fn rms<const D: usize>(x: Tensor<D>, w: Tensor<1>, eps: f64) -> Tensor<D> {
105    burn_rms_norm_nd(x, w, eps as f32)
106}
107
108fn selu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
109    const SCALE: f32 = 1.0507009873554804934193349852946;
110    const ALPHA_SCALE: f32 = 1.0507009873554804934193349852946 * 1.6732632423543772848170429916717;
111    let pos = activation::relu(x.clone());
112    let neg = x - pos.clone();
113    let selu_pos = pos * SCALE;
114    let selu_neg = (neg.exp() - 1.0) * ALPHA_SCALE;
115    selu_pos + selu_neg
116}
117
118fn sigmoid<const D: usize>(x: Tensor<D>) -> Tensor<D> {
119    activation::sigmoid(x)
120}
121
122simd_runtime_generate!(
123    fn ssm_scan_step(
124        state_r: &mut [f32],
125        state_i: &mut [f32],
126        a_r: &[f32],
127        a_i: &[f32],
128        bu_r: &[f32],
129        bu_i: &[f32],
130    ) {
131        let mut sr = &mut state_r[..];
132        let mut si = &mut state_i[..];
133        let mut ar = &a_r[..];
134        let mut ai = &a_i[..];
135        let mut br = &bu_r[..];
136        let mut bi = &bu_i[..];
137
138        while sr.len() >= S::Vf32::WIDTH {
139            let sr_v = S::Vf32::load_from_slice(sr);
140            let si_v = S::Vf32::load_from_slice(si);
141            let ar_v = S::Vf32::load_from_slice(ar);
142            let ai_v = S::Vf32::load_from_slice(ai);
143            let br_v = S::Vf32::load_from_slice(br);
144            let bi_v = S::Vf32::load_from_slice(bi);
145
146            let new_r = ar_v.mul_add(sr_v, ai_v.neg_mul_add(si_v, br_v));
147            let new_i = ar_v.mul_add(si_v, ai_v.mul_add(sr_v, bi_v));
148
149            new_r.copy_to_slice(sr);
150            new_i.copy_to_slice(si);
151
152            sr = &mut sr[S::Vf32::WIDTH..];
153            si = &mut si[S::Vf32::WIDTH..];
154            ar = &ar[S::Vf32::WIDTH..];
155            ai = &ai[S::Vf32::WIDTH..];
156            br = &br[S::Vf32::WIDTH..];
157            bi = &bi[S::Vf32::WIDTH..];
158        }
159
160        for j in 0..sr.len() {
161            let nr = ar[j] * sr[j] - ai[j] * si[j] + br[j];
162            let ni = ar[j] * si[j] + ai[j] * sr[j] + bi[j];
163            sr[j] = nr;
164            si[j] = ni;
165        }
166    }
167);
168
169struct BurnS5W {
170    log_lambda_real: Vec<f32>,
171    lambda_imag: Vec<f32>,
172    b_r: Tensor<2>,
173    b_i: Tensor<2>,
174    c_r: Tensor<2>,
175    c_i: Tensor<2>,
176    d: Tensor<1>,
177    log_delta: Vec<f32>,
178}
179
180struct BurnBlockW {
181    ssm: BurnS5W,
182    out_weight: Tensor<2>,
183    out_bias: Tensor<1>,
184    norm_weight: Tensor<1>,
185    norm_bias: Tensor<1>,
186    disc_cache: Mutex<HashMap<u32, (Vec<f32>, Vec<f32>, Tensor<2>, Tensor<2>)>>,
187}
188
189pub struct BurnFlowStateModel {
190    config: InferConfig,
191    embed_w: Tensor<2>,
192    embed_b: Tensor<1>,
193    blocks: Vec<BurnBlockW>,
194    decoder_w: Tensor<2>,
195    decoder_b: Tensor<1>,
196    legendre_cache: Mutex<HashMap<usize, Tensor<2>>>,
197}
198
199impl BurnFlowStateModel {
200    pub fn load(gguf_path: &Path, config: InferConfig) -> Result<Self> {
201        let weights = open_gguf_weights(gguf_path)
202            .with_context(|| format!("open {}", gguf_path.display()))?;
203
204        let embed_w = load_t(&weights, "embed.weight")?;
205        let embed_b = load_v(&weights, "embed.bias")?;
206        let mut blocks = Vec::with_capacity(config.num_layers);
207        for n in 0..config.num_layers {
208            blocks.push(BurnBlockW {
209                ssm: BurnS5W {
210                    log_lambda_real: load_f32_vec(
211                        &weights,
212                        &format!("blk.{n}.ssm.log_lambda_real"),
213                    )?,
214                    lambda_imag: load_f32_vec(&weights, &format!("blk.{n}.ssm.lambda_imag"))?,
215                    b_r: load_t(&weights, &format!("blk.{n}.ssm.b_r"))?,
216                    b_i: load_t(&weights, &format!("blk.{n}.ssm.b_i"))?,
217                    c_r: load_t(&weights, &format!("blk.{n}.ssm.c_r"))?,
218                    c_i: load_t(&weights, &format!("blk.{n}.ssm.c_i"))?,
219                    d: load_v(&weights, &format!("blk.{n}.ssm.d"))?,
220                    log_delta: load_f32_vec(&weights, &format!("blk.{n}.ssm.log_delta"))?,
221                },
222                out_weight: load_t(&weights, &format!("blk.{n}.out.weight"))?,
223                out_bias: load_v(&weights, &format!("blk.{n}.out.bias"))?,
224                norm_weight: load_v(&weights, &format!("blk.{n}.norm.weight"))?,
225                norm_bias: load_v(&weights, &format!("blk.{n}.norm.bias"))?,
226                disc_cache: Mutex::new(HashMap::new()),
227            });
228        }
229        Ok(Self {
230            config,
231            embed_w,
232            embed_b,
233            blocks,
234            decoder_w: load_t(&weights, "decoder.weight")?,
235            decoder_b: load_v(&weights, "decoder.bias")?,
236            legendre_cache: Mutex::new(HashMap::new()),
237        })
238    }
239
240    pub fn quantiles(&self) -> &[f32] {
241        self.config.quantiles()
242    }
243
244    pub fn median_index(&self) -> usize {
245        self.config.median_index()
246    }
247
248    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> Result<Vec<Vec<f32>>> {
249        let cfg = &self.config;
250        let ctx_len = context.len().min(cfg.context_length);
251        let start = context.len().saturating_sub(ctx_len);
252        let context = &context[start..];
253        let seq_len = context.len();
254        let (normed_values, final_mean, final_std) = causal_revin_norm(context, cfg.eps);
255        let mut input_data = vec![0.0f32; seq_len * cfg.n_inputs];
256        for t in 0..seq_len {
257            input_data[t * cfg.n_inputs] = normed_values[t];
258            if cfg.n_inputs > 1 {
259                input_data[t * cfg.n_inputs + 1] = 0.0;
260            }
261        }
262        let dev = device();
263        let input_t =
264            Tensor::<2>::from_data(TensorData::new(input_data, [seq_len, cfg.n_inputs]), &dev);
265        let mut hidden = linear(input_t, self.embed_w.clone(), Some(self.embed_b.clone()));
266        let scale_factor = cfg.decoder_patch_len as f32 / prediction_length as f32;
267        let scratch_len = seq_len * cfg.state_dim;
268        let mut scan_r = vec![0.0f32; scratch_len];
269        let mut scan_i = vec![0.0f32; scratch_len];
270        let mut state_r = vec![0.0f32; cfg.state_dim];
271        let mut state_i = vec![0.0f32; cfg.state_dim];
272        let num_layers = self.blocks.len();
273        for (i, block) in self.blocks.iter().enumerate() {
274            let is_last = i == num_layers - 1;
275            hidden = self.apply_s5_layer(
276                hidden,
277                block,
278                scale_factor,
279                is_last,
280                &mut scan_r,
281                &mut scan_i,
282                &mut state_r,
283                &mut state_i,
284            )?;
285        }
286        let n_q = cfg.quantiles.len();
287        let coeffs = linear(hidden, self.decoder_w.clone(), Some(self.decoder_b.clone()));
288        let coeffs: Tensor<2> = coeffs.reshape([n_q, cfg.decoder_dim]);
289        let basis = {
290            let mut cache = self
291                .legendre_cache
292                .lock()
293                .map_err(|_| anyhow::anyhow!("cache mutex poisoned"))?;
294            if !cache.contains_key(&prediction_length) {
295                let raw = legendre_basis(
296                    prediction_length,
297                    cfg.decoder_dim,
298                    cfg.basis_range,
299                    scale_factor,
300                    cfg.decoder_patch_len,
301                );
302                let flat: Vec<f32> = raw.into_iter().flatten().collect();
303                cache.insert(
304                    prediction_length,
305                    Tensor::<2>::from_data(
306                        TensorData::new(flat, [prediction_length, cfg.decoder_dim]),
307                        &dev,
308                    ),
309                );
310            }
311            cache[&prediction_length].clone()
312        };
313        let out_t = coeffs.matmul(basis.transpose());
314        let out_raw: Vec<f32> = to_host(&out_t)?;
315        let mut denormed = vec![vec![0.0f32; prediction_length]; n_q];
316        for q in 0..n_q {
317            for p in 0..prediction_length {
318                denormed[q][p] = out_raw[q * prediction_length + p] * final_std + final_mean;
319            }
320        }
321        let mut output = vec![vec![0.0f32; prediction_length]; n_q];
322        for p in 0..prediction_length {
323            let mut sorted: Vec<f32> = (0..n_q).map(|q| denormed[q][p]).collect();
324            sorted.sort_by(|a, b| a.total_cmp(b));
325            for (qi, &prob) in cfg.quantiles.iter().enumerate() {
326                let idx = (n_q - 1) as f32 * prob;
327                let lower = idx.floor() as usize;
328                let upper = idx.ceil() as usize;
329                let weight = idx - lower as f32;
330                output[qi][p] = sorted[lower] * (1.0 - weight) + sorted[upper] * weight;
331            }
332        }
333        Ok(output)
334    }
335
336    #[allow(clippy::too_many_arguments)]
337    fn apply_s5_layer(
338        &self,
339        x: Tensor<2>,
340        block: &BurnBlockW,
341        scale_factor: f32,
342        is_last: bool,
343        scan_r: &mut Vec<f32>,
344        scan_i: &mut Vec<f32>,
345        state_r: &mut Vec<f32>,
346        state_i: &mut Vec<f32>,
347    ) -> Result<Tensor<2>> {
348        let cfg = &self.config;
349        let state_dim = cfg.state_dim;
350        let embed_dim = cfg.embed_dim;
351        let seq_len = x.dims()[0];
352        let skip: Tensor<2> = if is_last {
353            x.clone().narrow(0, seq_len - 1, 1)
354        } else {
355            x.clone()
356        };
357        let (a_bar_r, a_bar_i, b_bar_r, b_bar_i) = {
358            let key = scale_factor.to_bits();
359            let mut cache = block
360                .disc_cache
361                .lock()
362                .map_err(|_| anyhow::anyhow!("cache mutex poisoned"))?;
363            if !cache.contains_key(&key) {
364                let result = discretize(&block.ssm, scale_factor, state_dim, embed_dim)?;
365                cache.insert(key, result);
366            }
367            let (ar, ai, brt, bit) = &cache[&key];
368            (ar.clone(), ai.clone(), brt.clone(), bit.clone())
369        };
370        let bu_r = x.clone().matmul(b_bar_r.transpose());
371        let bu_i = x.matmul(b_bar_i.transpose());
372        let bu_r_data: Vec<f32> = to_host(&bu_r)?;
373        let bu_i_data: Vec<f32> = to_host(&bu_i)?;
374        state_r[..state_dim].fill(0.0);
375        state_i[..state_dim].fill(0.0);
376        let (h_r_t, h_i_t) = if is_last {
377            for t in 0..seq_len {
378                let bu_r_t = &bu_r_data[t * state_dim..(t + 1) * state_dim];
379                let bu_i_t = &bu_i_data[t * state_dim..(t + 1) * state_dim];
380                ssm_scan_step(
381                    &mut state_r[..state_dim],
382                    &mut state_i[..state_dim],
383                    &a_bar_r,
384                    &a_bar_i,
385                    bu_r_t,
386                    bu_i_t,
387                );
388            }
389            let dev = device();
390            (
391                Tensor::<2>::from_data(
392                    TensorData::new(state_r[..state_dim].to_vec(), [1, state_dim]),
393                    &dev,
394                ),
395                Tensor::<2>::from_data(
396                    TensorData::new(state_i[..state_dim].to_vec(), [1, state_dim]),
397                    &dev,
398                ),
399            )
400        } else {
401            let needed = seq_len * state_dim;
402            if scan_r.len() < needed {
403                scan_r.resize(needed, 0.0);
404            }
405            if scan_i.len() < needed {
406                scan_i.resize(needed, 0.0);
407            }
408            for t in 0..seq_len {
409                let bu_r_t = &bu_r_data[t * state_dim..(t + 1) * state_dim];
410                let bu_i_t = &bu_i_data[t * state_dim..(t + 1) * state_dim];
411                ssm_scan_step(
412                    &mut state_r[..state_dim],
413                    &mut state_i[..state_dim],
414                    &a_bar_r,
415                    &a_bar_i,
416                    bu_r_t,
417                    bu_i_t,
418                );
419                let row = t * state_dim;
420                scan_r[row..row + state_dim].copy_from_slice(&state_r[..state_dim]);
421                scan_i[row..row + state_dim].copy_from_slice(&state_i[..state_dim]);
422            }
423            let dev = device();
424            (
425                Tensor::<2>::from_data(
426                    TensorData::new(scan_r[..needed].to_vec(), [seq_len, state_dim]),
427                    &dev,
428                ),
429                Tensor::<2>::from_data(
430                    TensorData::new(scan_i[..needed].to_vec(), [seq_len, state_dim]),
431                    &dev,
432                ),
433            )
434        };
435        let y_from_cr = h_r_t.matmul(block.ssm.c_r.clone().transpose());
436        let y_from_ci = h_i_t.matmul(block.ssm.c_i.clone().transpose());
437        let y_raw = y_from_cr - y_from_ci;
438        let d: Tensor<1> = block.ssm.d.clone();
439        let y_t = y_raw + skip.clone() * d.unsqueeze_dim::<2>(0);
440        let y_selu = selu(y_t);
441        let gate_pre = linear(
442            y_selu.clone(),
443            block.out_weight.clone(),
444            Some(block.out_bias.clone()),
445        );
446        let gate = sigmoid(gate_pre);
447        let y_gated = y_selu * gate;
448        let y_normed = layer_norm(
449            y_gated,
450            block.norm_weight.clone(),
451            block.norm_bias.clone(),
452            self.config.eps,
453        );
454        Ok(y_normed + skip)
455    }
456}
457
458fn discretize(
459    ssm: &BurnS5W,
460    scale_factor: f32,
461    state_dim: usize,
462    embed_dim: usize,
463) -> Result<(Vec<f32>, Vec<f32>, Tensor<2>, Tensor<2>)> {
464    let mut a_r = vec![0.0f32; state_dim];
465    let mut a_i = vec![0.0f32; state_dim];
466    let mut coeff_r = vec![0.0f32; state_dim];
467    let mut coeff_i = vec![0.0f32; state_dim];
468    for s in 0..state_dim {
469        let lam_r = -ssm.log_lambda_real[s].exp();
470        let lam_i = ssm.lambda_imag[s];
471        let delta = scale_factor * ssm.log_delta[s].exp();
472        let exp_r = lam_r * delta;
473        let exp_i = lam_i * delta;
474        let mag = exp_r.exp();
475        a_r[s] = mag * exp_i.cos();
476        a_i[s] = mag * exp_i.sin();
477        let num_r = a_r[s] - 1.0;
478        let num_i = a_i[s];
479        let denom_sq = lam_r * lam_r + lam_i * lam_i;
480        if denom_sq > 1e-20 {
481            coeff_r[s] = (num_r * lam_r + num_i * lam_i) / denom_sq;
482            coeff_i[s] = (num_i * lam_r - num_r * lam_i) / denom_sq;
483        } else {
484            coeff_r[s] = delta;
485            coeff_i[s] = 0.0;
486        }
487    }
488    let b_r_data: Vec<f32> = to_host(&ssm.b_r)?;
489    let b_i_data: Vec<f32> = to_host(&ssm.b_i)?;
490    let mut b_bar_r = vec![0.0f32; state_dim * embed_dim];
491    let mut b_bar_i = vec![0.0f32; state_dim * embed_dim];
492    for s in 0..state_dim {
493        for e in 0..embed_dim {
494            let br = b_r_data[s * embed_dim + e];
495            let bi = b_i_data[s * embed_dim + e];
496            b_bar_r[s * embed_dim + e] = coeff_r[s] * br - coeff_i[s] * bi;
497            b_bar_i[s * embed_dim + e] = coeff_r[s] * bi + coeff_i[s] * br;
498        }
499    }
500    let dev = device();
501    let b_bar_r_t = Tensor::<2>::from_data(TensorData::new(b_bar_r, [state_dim, embed_dim]), &dev);
502    let b_bar_i_t = Tensor::<2>::from_data(TensorData::new(b_bar_i, [state_dim, embed_dim]), &dev);
503    Ok((a_r, a_i, b_bar_r_t, b_bar_i_t))
504}
505
506fn causal_revin_norm(x: &[f32], eps: f32) -> (Vec<f32>, f32, f32) {
507    let n = x.len();
508    let mut normed = vec![0.0f32; n];
509    let mut cum_sum = 0.0f32;
510    let mut cum_sq_diff = 0.0f32;
511    let mut final_mean = 0.0f32;
512    let mut final_std = 1.0f32;
513    for t in 0..n {
514        let count = (t + 1) as f32;
515        cum_sum += x[t];
516        let mean_t = cum_sum / count;
517        cum_sq_diff += (x[t] - mean_t) * (x[t] - mean_t);
518        let var_t = (cum_sq_diff / count).max(0.0);
519        let std_t = (var_t + eps).sqrt();
520        normed[t] = (x[t] - mean_t) / std_t;
521        if t == n - 1 {
522            final_mean = mean_t;
523            final_std = std_t;
524        }
525    }
526    (normed, final_mean, final_std)
527}
528
529fn legendre_basis(
530    n_points: usize,
531    degree: usize,
532    range: [f32; 2],
533    scale: f32,
534    pred_dist: usize,
535) -> Vec<Vec<f32>> {
536    let dt = scale * (range[1] - range[0]) / pred_dist as f32;
537    let t: Vec<f32> = (1..=n_points).map(|i| range[0] + i as f32 * dt).collect();
538    let mut basis = vec![vec![0.0f32; degree + 1]; n_points];
539    for (p, &x) in t.iter().enumerate() {
540        basis[p][0] = 1.0;
541        if degree >= 1 {
542            basis[p][1] = x;
543        }
544        for k in 1..degree {
545            let kf = k as f32;
546            basis[p][k + 1] =
547                ((2.0 * kf + 1.0) * x * basis[p][k] - kf * basis[p][k - 1]) / (kf + 1.0);
548        }
549        for d in 0..degree {
550            basis[p][d] /= 4.0;
551        }
552        basis[p].truncate(degree);
553    }
554    basis
555}
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
561
562    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
563        let mut x = seed;
564        (0..n)
565            .map(|_| {
566                x = x
567                    .wrapping_mul(6364136223846793005)
568                    .wrapping_add(1442695040888963407);
569                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
570            })
571            .collect()
572    }
573
574    fn tiny_config() -> InferConfig {
575        InferConfig {
576            num_layers: 1,
577            embed_dim: 8,
578            state_dim: 8,
579            n_inputs: 2,
580            decoder_dim: 4,
581            decoder_patch_len: 2,
582            quantiles: vec![0.25, 0.5, 0.75],
583            basis_range: [-1.0, 0.95],
584            context_length: 16,
585            eps: 1e-5,
586        }
587    }
588
589    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
590        *seed += 1;
591        let data: Vec<f32> = pseudo(rows * cols, *seed)
592            .iter()
593            .map(|v| v * scale)
594            .collect();
595        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
596        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
597    }
598
599    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
600        *seed += 1;
601        let bytes: Vec<u8> = pseudo(n, *seed)
602            .iter()
603            .flat_map(|v| v.to_le_bytes())
604            .collect();
605        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
606    }
607
608    #[test]
609    fn burn_forecast_golden() {
610        let dir = std::env::temp_dir().join(format!("flowstate-burn-{}", std::process::id()));
611        let _ = std::fs::remove_dir_all(&dir);
612        std::fs::create_dir_all(&dir).unwrap();
613        let path = dir.join("model.gguf");
614        let mut w = GGUFWriter::new();
615        w.add_metadata(
616            "general.architecture",
617            GGUFMetaValue::String("flowstate".into()),
618        );
619        let mut seed = 3100u64;
620        let s = 0.05;
621        put2(&mut w, "embed.weight", 8, 2, &mut seed, s);
622        put1(&mut w, "embed.bias", 8, &mut seed);
623        put1(&mut w, "blk.0.ssm.log_lambda_real", 8, &mut seed);
624        put1(&mut w, "blk.0.ssm.lambda_imag", 8, &mut seed);
625        put2(&mut w, "blk.0.ssm.b_r", 8, 8, &mut seed, s);
626        put2(&mut w, "blk.0.ssm.b_i", 8, 8, &mut seed, s);
627        put2(&mut w, "blk.0.ssm.c_r", 8, 8, &mut seed, s);
628        put2(&mut w, "blk.0.ssm.c_i", 8, 8, &mut seed, s);
629        put1(&mut w, "blk.0.ssm.d", 8, &mut seed);
630        put1(&mut w, "blk.0.ssm.log_delta", 8, &mut seed);
631        put2(&mut w, "blk.0.out.weight", 8, 8, &mut seed, s);
632        put1(&mut w, "blk.0.out.bias", 8, &mut seed);
633        put1(&mut w, "blk.0.norm.weight", 8, &mut seed);
634        put1(&mut w, "blk.0.norm.bias", 8, &mut seed);
635        put2(&mut w, "decoder.weight", 12, 8, &mut seed, s);
636        put1(&mut w, "decoder.bias", 12, &mut seed);
637        let f = std::fs::File::create(&path).unwrap();
638        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
639
640        let cfg = tiny_config();
641        let burn = BurnFlowStateModel::load(&path, cfg).unwrap();
642        let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
643        let t0 = std::time::Instant::now();
644        let a = burn.forecast(&ctx, 2).unwrap();
645        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
646        let b = burn.forecast(&ctx, 2).unwrap();
647        assert_eq!(a.len(), 3);
648        assert!(a
649            .iter()
650            .all(|r| r.len() == 2 && r.iter().all(|v| v.is_finite())));
651        assert_eq!(a, b, "forecast must be deterministic");
652        println!("flowstate synthetic: burn {burn_ms:.1}ms out={a:?}");
653        // Golden recorded from the verified Burn port (ex-candle parity 0.00e0).
654        let golden: Vec<Vec<f32>> = vec![
655            vec![11.657719, 11.33724],
656            vec![11.73962, 11.358313],
657            vec![11.769588, 11.490122],
658        ];
659        assert_eq!(a.len(), golden.len());
660        for (ra, re) in a.iter().zip(golden.iter()) {
661            assert_eq!(ra.len(), re.len());
662            for (x, y) in ra.iter().zip(re.iter()) {
663                assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
664            }
665        }
666        let _ = std::fs::remove_dir_all(&dir);
667    }
668}
zsfm_flowstate/infer/mod.rs

zsfm_flowstate/infer/
mod.rs
