Skip to main content

cortiq_engine/
ltxte.rs

1//! The LTX-2.5 prompt encoder: Gemma-4 12B, the two aggregate projections,
2//! and the embeddings connectors — everything between a prompt string and
3//! the context the DiT cross-attends to.
4//!
5//! Three stages, all read from one `ltx-2.5-av` container:
6//!
7//! 1. **Gemma-4 12B** (`te.model.*`), 48 layers over a 1024-token window.
8//!    Two layer kinds alternate: forty *sliding* layers (head 256, 16 query
9//!    heads over 8 key heads, θ = 10⁴) and eight *full* layers every sixth
10//!    (head 512, 16 query heads over one key head, θ = 10⁶ with only the
11//!    first quarter of each head rotated, and the value projection **is**
12//!    the key projection). Attention is unscaled — the q/k RMS-norms carry
13//!    the scale — and every layer ends multiplied by its `layer_scalar`.
14//! 2. **The aggregate projections** (`te.text_embedding_projection.*`).
15//!    The features are not the last hidden state: all forty-nine layer
16//!    outputs are RMS-normalized per token *per layer*, concatenated into
17//!    188160 numbers, rescaled by `sqrt(out_dim / hidden)`, and projected —
18//!    once to 4096 for video, once to 2048 for audio.
19//! 3. **The connectors** (`dit.{video,audio}_embeddings_connector.*`), eight
20//!    gated-attention blocks with 1-D RoPE. Padded positions are replaced by
21//!    128 learnable registers tiled across the window first, which is why
22//!    the DiT needs no prompt mask: after the substitution every position
23//!    carries signal.
24
25use crate::ltxdit::{Attn, Lin, Rope, gelu_tanh, rms_plain, rows};
26use crate::pool::Pool;
27use crate::qtensor::QTensor;
28use cortiq_core::CmfModel;
29use std::sync::Arc;
30
31const EPS: f64 = 1e-6;
32
33/// RMS normalization with a plain (not `1 + w`) learned weight — Gemma-4's,
34/// unlike Gemma-2's and Gemma-3's.
35fn rms(x: &[f32], w: &[f32], dst: &mut [f32]) {
36    let ss = x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64;
37    let inv = 1.0 / (ss + EPS).sqrt();
38    for ((d, &v), &g) in dst.iter_mut().zip(x).zip(w) {
39        *d = (v as f64 * inv) as f32 * g;
40    }
41}
42
43/// RMS normalization with no learned weight (the value norm).
44fn rms_nw(x: &mut [f32]) {
45    let ss = x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64;
46    let inv = 1.0 / (ss + EPS).sqrt();
47    for v in x.iter_mut() {
48        *v = (*v as f64 * inv) as f32;
49    }
50}
51
52/// Half-rotation RoPE tables: `cos`/`sin` of `[T, head_dim/2]`, applied as
53/// `x·cos + rotate_half(x)·sin`.
54struct Rot {
55    cos: Vec<f32>,
56    sin: Vec<f32>,
57    half: usize,
58}
59
60impl Rot {
61    /// `inv_freq[j] = base^(-2j/head_dim)` for the rotated prefix, zero for
62    /// the rest — a zero frequency is an identity rotation, which is how the
63    /// "proportional" variant leaves three quarters of a full-attention head
64    /// unrotated.
65    fn build(seq: usize, head_dim: usize, base: f64, rotary: f64) -> Rot {
66        let half = head_dim / 2;
67        let rope_angles = (rotary * head_dim as f64 / 2.0) as usize;
68        let inv: Vec<f64> = (0..half)
69            .map(|j| {
70                if j < rope_angles {
71                    1.0 / base.powf((2 * j) as f64 / head_dim as f64)
72                } else {
73                    0.0
74                }
75            })
76            .collect();
77        let mut cos = vec![0f32; seq * half];
78        let mut sin = vec![0f32; seq * half];
79        for p in 0..seq {
80            for (j, &f) in inv.iter().enumerate() {
81                let a = p as f64 * f;
82                cos[p * half + j] = a.cos() as f32;
83                sin[p * half + j] = a.sin() as f32;
84            }
85        }
86        Rot { cos, sin, half }
87    }
88
89    fn apply(&self, t: usize, row: &mut [f32]) {
90        let h = self.half;
91        let (c, s) = (&self.cos[t * h..(t + 1) * h], &self.sin[t * h..(t + 1) * h]);
92        for i in 0..h {
93            let (a, b) = (row[i], row[i + h]);
94            row[i] = a * c[i] - b * s[i];
95            row[i + h] = b * c[i] + a * s[i];
96        }
97    }
98}
99
100struct GemmaLayer {
101    q: Lin,
102    k: Lin,
103    v: Option<Lin>,
104    o: Lin,
105    q_norm: Vec<f32>,
106    k_norm: Vec<f32>,
107    in_norm: Vec<f32>,
108    post_attn_norm: Vec<f32>,
109    pre_ff_norm: Vec<f32>,
110    post_ff_norm: Vec<f32>,
111    gate: Lin,
112    up: Lin,
113    down: Lin,
114    scalar: f32,
115    head_dim: usize,
116    q_heads: usize,
117    kv_heads: usize,
118    sliding: bool,
119}
120
121/// The full prompt encoder.
122pub struct LtxTextEncoder {
123    embed: QTensor,
124    mapped_bytes: u64,
125    layers: Vec<GemmaLayer>,
126    norm: Vec<f32>,
127    video_agg: Lin,
128    audio_agg: Lin,
129    v_conn: Connector,
130    a_conn: Connector,
131    hidden: usize,
132    embed_scale: f32,
133    pub max_len: usize,
134    pub bos: u32,
135    pub pad: u32,
136}
137
138struct Connector {
139    blocks: Vec<(Attn, Lin, Lin)>,
140    registers: Vec<f32>,
141    dim: usize,
142    heads: usize,
143    dh: usize,
144    max_pos: f64,
145}
146
147fn vecf(model: &Arc<CmfModel>, name: &str) -> Result<Vec<f32>, String> {
148    crate::dit::cmf_f32(model, name)
149}
150
151impl LtxTextEncoder {
152    pub fn from_cmf(model: &Arc<CmfModel>) -> Result<LtxTextEncoder, String> {
153        let cfg_bytes = ["te.gemma_config_json", "ltx.gemma_config_json"]
154            .iter()
155            .find_map(|n| model.tensor(n).map(|e| model.entry_bytes(e)));
156        let cfg: serde_json::Value = match cfg_bytes {
157            Some(b) => serde_json::from_slice(b).map_err(|e| format!("gemma config: {e}"))?,
158            None => serde_json::Value::Null,
159        };
160        let tc = cfg.get("text_config").cloned().unwrap_or(serde_json::Value::Null);
161        let g = |k: &str, d: f64| tc.get(k).and_then(|v| v.as_f64()).unwrap_or(d);
162        let hidden = g("hidden_size", 3840.0) as usize;
163        let n_layers = g("num_hidden_layers", 48.0) as usize;
164        let head_dim = g("head_dim", 256.0) as usize;
165        let global_head_dim = g("global_head_dim", 512.0) as usize;
166        let types: Vec<String> = tc
167            .get("layer_types")
168            .and_then(|v| v.as_array())
169            .map(|a| a.iter().map(|s| s.as_str().unwrap_or("sliding_attention").to_string()).collect())
170            .unwrap_or_else(|| {
171                // the released layout: every sixth layer is a full-attention one
172                (0..n_layers)
173                    .map(|i| {
174                        if i % 6 == 5 { "full_attention".into() } else { "sliding_attention".into() }
175                    })
176                    .collect()
177            });
178
179        let mut layers = Vec::with_capacity(n_layers);
180        for i in 0..n_layers {
181            let p = format!("te.model.layers.{i}");
182            let sliding = types[i] != "full_attention";
183            let hd = if sliding { head_dim } else { global_head_dim };
184            let q = Lin::load(model, &format!("{p}.self_attn.q_proj"), false)?;
185            let k = Lin::load(model, &format!("{p}.self_attn.k_proj"), false)?;
186            let v = match model.tensor(&format!("{p}.self_attn.v_proj.weight")) {
187                Some(_) => Some(Lin::load(model, &format!("{p}.self_attn.v_proj"), false)?),
188                None => None,
189            };
190            let q_rows = model
191                .tensor(&format!("{p}.self_attn.q_proj.weight"))
192                .ok_or_else(|| format!("missing {p}.self_attn.q_proj.weight"))?
193                .shape[0];
194            let k_rows = model
195                .tensor(&format!("{p}.self_attn.k_proj.weight"))
196                .ok_or("missing k_proj")?
197                .shape[0];
198            layers.push(GemmaLayer {
199                q,
200                k,
201                v,
202                o: Lin::load(model, &format!("{p}.self_attn.o_proj"), false)?,
203                q_norm: vecf(model, &format!("{p}.self_attn.q_norm.weight"))?,
204                k_norm: vecf(model, &format!("{p}.self_attn.k_norm.weight"))?,
205                in_norm: vecf(model, &format!("{p}.input_layernorm.weight"))?,
206                post_attn_norm: vecf(model, &format!("{p}.post_attention_layernorm.weight"))?,
207                pre_ff_norm: vecf(model, &format!("{p}.pre_feedforward_layernorm.weight"))?,
208                post_ff_norm: vecf(model, &format!("{p}.post_feedforward_layernorm.weight"))?,
209                gate: Lin::load(model, &format!("{p}.mlp.gate_proj"), false)?,
210                up: Lin::load(model, &format!("{p}.mlp.up_proj"), false)?,
211                down: Lin::load(model, &format!("{p}.mlp.down_proj"), false)?,
212                scalar: vecf(model, &format!("{p}.layer_scalar"))?[0],
213                head_dim: hd,
214                q_heads: q_rows / hd,
215                kv_heads: k_rows / hd,
216                sliding,
217            });
218        }
219
220        let conn = |prefix: &str, dim: usize, heads: usize, dh: usize| -> Result<Connector, String> {
221            let mut blocks = Vec::new();
222            let mut i = 0usize;
223            while model
224                .tensor(&format!("{prefix}.transformer_1d_blocks.{i}.attn1.to_q.weight"))
225                .is_some()
226            {
227                let p = format!("{prefix}.transformer_1d_blocks.{i}");
228                blocks.push((
229                    Attn::load(model, &format!("{p}.attn1"), heads, dh)?,
230                    Lin::load(model, &format!("{p}.ff.net.0.proj"), true)?,
231                    Lin::load(model, &format!("{p}.ff.net.2"), true)?,
232                ));
233                i += 1;
234            }
235            Ok(Connector {
236                blocks,
237                registers: vecf(model, &format!("{prefix}.learnable_registers"))?,
238                dim,
239                heads,
240                dh,
241                max_pos: 4096.0,
242            })
243        };
244
245        Ok(LtxTextEncoder {
246            embed: QTensor::from_model(model, "te.model.embed_tokens.weight")?,
247            mapped_bytes: model.primary_bytes().len() as u64,
248            layers,
249            norm: vecf(model, "te.model.norm.weight")?,
250            video_agg: Lin::load(model, "te.text_embedding_projection.video_aggregate_embed", true)?,
251            audio_agg: Lin::load(model, "te.text_embedding_projection.audio_aggregate_embed", true)?,
252            v_conn: conn("dit.video_embeddings_connector", 4096, 32, 128)?,
253            a_conn: conn("dit.audio_embeddings_connector", 2048, 32, 64)?,
254            hidden,
255            embed_scale: (hidden as f64).sqrt() as f32,
256            max_len: 1024,
257            bos: tc.get("bos_token_id").and_then(|v| v.as_u64()).unwrap_or(2) as u32,
258            pad: tc.get("pad_token_id").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
259        })
260    }
261
262    /// Left-pad the prompt to the encoder's window, prepending BOS — the
263    /// layout the reference tokenizer produces.
264    pub fn pad_ids(&self, ids: &[u32]) -> (Vec<u32>, Vec<f32>) {
265        let mut body = Vec::with_capacity(self.max_len);
266        if ids.first() != Some(&self.bos) {
267            body.push(self.bos);
268        }
269        body.extend_from_slice(ids);
270        body.truncate(self.max_len);
271        let padlen = self.max_len - body.len();
272        let mut out = vec![self.pad; padlen];
273        out.extend_from_slice(&body);
274        let mut mask = vec![0f32; padlen];
275        mask.extend(std::iter::repeat_n(1f32, body.len()));
276        (out, mask)
277    }
278
279    /// Every one of the 49 hidden states, in HF's order: the scaled
280    /// embedding first, then each layer's output.
281    pub fn hidden_states(&self, ids: &[u32], mask: &[f32], pool: Option<&Pool>) -> Vec<Vec<f32>> {
282        // On a device that cannot keep the whole container wired, the
283        // encoder is the wrong thing to give the GPU: it lives in a
284        // different part of the file than the transformer, so putting both
285        // on the books makes the driver evict between commits and the
286        // *denoising loop* pays for it, step after step. The encoder runs
287        // once per render — let it have the CPU.
288        let _pause = self.crowds_the_device().then(crate::gpu::pause_gpu);
289        self.hidden_states_inner(ids, mask, pool)
290    }
291
292    #[cfg(target_os = "macos")]
293    fn crowds_the_device(&self) -> bool {
294        // Two tests, and the buffer one is what actually bites: a container
295        // larger than the biggest buffer the device will make needs more
296        // than one window, and the encoder's weights and the transformer's
297        // are in different windows. Whichever limit the mapping is past,
298        // the prompt phase is better off on the CPU.
299        let one_window = crate::gpu_metal::max_buffer_bytes();
300        let wired = crate::gpu_metal::working_set_bytes();
301        (one_window > 0 && self.mapped_bytes > one_window)
302            || (wired > 0 && self.mapped_bytes > wired)
303    }
304
305    #[cfg(not(target_os = "macos"))]
306    fn crowds_the_device(&self) -> bool {
307        false
308    }
309
310    fn hidden_states_inner(&self, ids: &[u32], mask: &[f32], pool: Option<&Pool>) -> Vec<Vec<f32>> {
311        let t = ids.len();
312        let d = self.hidden;
313        let mut x = vec![0f32; t * d];
314        for (i, &id) in ids.iter().enumerate() {
315            self.embed.row_f32(id as usize, &mut x[i * d..(i + 1) * d]);
316            for v in x[i * d..(i + 1) * d].iter_mut() {
317                *v *= self.embed_scale;
318            }
319        }
320        let mut out = vec![x.clone()];
321        // one rotation table per (head_dim, theta) pair in use
322        let rot_slide = Rot::build(t, self.layers[0].head_dim, 10000.0, 1.0);
323        let full = self.layers.iter().find(|l| !l.sliding);
324        let rot_full = full.map(|l| Rot::build(t, l.head_dim, 1_000_000.0, 0.25));
325
326        for layer in &self.layers {
327            let rot = if layer.sliding { &rot_slide } else { rot_full.as_ref().unwrap() };
328            let mut h = vec![0f32; t * d];
329            for i in 0..t {
330                rms(&x[i * d..(i + 1) * d], &layer.in_norm, &mut h[i * d..(i + 1) * d]);
331            }
332            let attn = self.attention(layer, &h, t, mask, rot, pool);
333            for i in 0..t {
334                let mut n = vec![0f32; d];
335                rms(&attn[i * d..(i + 1) * d], &layer.post_attn_norm, &mut n);
336                for (v, &a) in x[i * d..(i + 1) * d].iter_mut().zip(&n) {
337                    *v += a;
338                }
339            }
340            let mut h2 = vec![0f32; t * d];
341            for i in 0..t {
342                rms(&x[i * d..(i + 1) * d], &layer.pre_ff_norm, &mut h2[i * d..(i + 1) * d]);
343            }
344            let mut g = layer.gate.apply(&h2, t, pool);
345            let u = layer.up.apply(&h2, t, pool);
346            for (a, &b) in g.iter_mut().zip(&u) {
347                *a = gelu_tanh(*a) * b;
348            }
349            let ff = layer.down.apply(&g, t, pool);
350            for i in 0..t {
351                let mut n = vec![0f32; d];
352                rms(&ff[i * d..(i + 1) * d], &layer.post_ff_norm, &mut n);
353                for (v, &a) in x[i * d..(i + 1) * d].iter_mut().zip(&n) {
354                    *v += a;
355                }
356            }
357            for v in x.iter_mut() {
358                *v *= layer.scalar;
359            }
360            out.push(x.clone());
361        }
362        // HF's hidden-state tuple is (embedding, layer 0 … layer 46,
363        // norm(layer 47)): the last entry is the *normalized* final state,
364        // not the raw one. The feature extractor reads all forty-nine, so
365        // getting this wrong poisons a forty-ninth of every token's
366        // features — and it is the entry with the largest magnitude.
367        if let Some(last) = out.last_mut() {
368            let mut n = vec![0f32; t * d];
369            for i in 0..t {
370                rms(&last[i * d..(i + 1) * d], &self.norm, &mut n[i * d..(i + 1) * d]);
371            }
372            *last = n;
373        }
374        out
375    }
376
377    fn attention(
378        &self,
379        l: &GemmaLayer,
380        h: &[f32],
381        t: usize,
382        mask: &[f32],
383        rot: &Rot,
384        pool: Option<&Pool>,
385    ) -> Vec<f32> {
386        let hd = l.head_dim;
387        let qi = l.q_heads * hd;
388        let ki = l.kv_heads * hd;
389        let mut q = l.q.apply(h, t, pool);
390        let mut k = l.k.apply(h, t, pool);
391        let mut v = match &l.v {
392            Some(p) => p.apply(h, t, pool),
393            None => k.clone(),
394        };
395        for i in 0..t {
396            for hh in 0..l.q_heads {
397                let r = &mut q[i * qi + hh * hd..i * qi + (hh + 1) * hd];
398                let mut n = vec![0f32; hd];
399                rms(r, &l.q_norm, &mut n);
400                r.copy_from_slice(&n);
401                rot.apply(i, r);
402            }
403            for hh in 0..l.kv_heads {
404                let r = &mut k[i * ki + hh * hd..i * ki + (hh + 1) * hd];
405                let mut n = vec![0f32; hd];
406                rms(r, &l.k_norm, &mut n);
407                r.copy_from_slice(&n);
408                rot.apply(i, r);
409                rms_nw(&mut v[i * ki + hh * hd..i * ki + (hh + 1) * hd]);
410            }
411        }
412        // The mask is the same for every head: causal, bounded by the
413        // sliding window, and closed over the prompt's left padding. Build
414        // it once, then run both halves of attention as GEMMs per head.
415        let window = if l.sliding { 1024usize } else { usize::MAX };
416        let mut bias = vec![0f32; t * t];
417        for i in 0..t {
418            for j in 0..t {
419                let blocked =
420                    j > i || (window != usize::MAX && i - j >= window) || mask[j] == 0.0;
421                bias[i * t + j] = if blocked { f32::NEG_INFINITY } else { 0.0 };
422            }
423        }
424        // A left-pad row attends to nothing at all. The reference's masked
425        // softmax leaves it a uniform average; we leave it zero. Either is
426        // dead weight — the feature extractor masks these rows and the
427        // connector overwrites them with registers — but a row of all
428        // -inf would come back NaN, so mark it and zero it after.
429        let dead: Vec<bool> = (0..t).map(|i| mask[i] == 0.0).collect();
430        let mut out = vec![0f32; t * qi];
431        let mut qh = vec![0f32; t * hd];
432        let mut kh = vec![0f32; t * hd];
433        let mut vh = vec![0f32; t * hd];
434        let mut sc = vec![0f32; t * t];
435        let mut oh = vec![0f32; t * hd];
436        for hh in 0..l.q_heads {
437            let kv = hh * l.kv_heads / l.q_heads;
438            for i in 0..t {
439                qh[i * hd..(i + 1) * hd].copy_from_slice(&q[i * qi + hh * hd..][..hd]);
440                kh[i * hd..(i + 1) * hd].copy_from_slice(&k[i * ki + kv * hd..][..hd]);
441                vh[i * hd..(i + 1) * hd].copy_from_slice(&v[i * ki + kv * hd..][..hd]);
442            }
443            crate::fcd_ops::gemm_nt(&qh, &kh, &mut sc, t, hd, t, pool);
444            let sp = crate::ltxdit::Shared(sc.as_mut_ptr());
445            rows(pool, t, &|s, e| {
446                let r = unsafe { sp.at(s * t, (e - s) * t) };
447                for (row, i) in r.chunks_exact_mut(t).zip(s..e) {
448                    if dead[i] {
449                        row.iter_mut().for_each(|x| *x = 0.0);
450                        continue;
451                    }
452                    for (x, b) in row.iter_mut().zip(&bias[i * t..(i + 1) * t]) {
453                        *x += *b;
454                    }
455                    crate::ltxdit::softmax(row);
456                }
457            });
458            oh.iter_mut().for_each(|x| *x = 0.0);
459            crate::fcd_ops::gemm_dx(&sc, &vh, &mut oh, t, hd, t, pool);
460            for i in 0..t {
461                out[i * qi + hh * hd..i * qi + (hh + 1) * hd]
462                    .copy_from_slice(&oh[i * hd..(i + 1) * hd]);
463            }
464        }
465        l.o.apply(&out, t, pool)
466    }
467
468    /// Hidden states → the two context tensors the DiT reads.
469    pub fn encode_ids(
470        &self,
471        ids: &[u32],
472        mask: &[f32],
473        pool: Option<&Pool>,
474    ) -> (Vec<f32>, Vec<f32>, usize) {
475        // The pause covers the *whole* prompt phase — the layers, the
476        // aggregate projections and the connectors. Leaving the connectors
477        // outside it was enough to put the far half of the container back on
478        // the driver's books and cost seconds a call.
479        let _pause = self.crowds_the_device().then(crate::gpu::pause_gpu);
480        let hs = self.hidden_states_inner(ids, mask, pool);
481        let (t, d, l) = (ids.len(), self.hidden, hs.len());
482        // per-token, per-layer RMS over the hidden dimension, concatenated
483        // layer-last: [T, d·L]
484        let mut feats = vec![0f32; t * d * l];
485        for (li, layer) in hs.iter().enumerate() {
486            for i in 0..t {
487                let row = &layer[i * d..(i + 1) * d];
488                let var = row.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / d as f64;
489                let inv = 1.0 / (var + 1e-6).sqrt();
490                let keep = mask[i] != 0.0;
491                for j in 0..d {
492                    feats[i * d * l + j * l + li] = if keep { (row[j] as f64 * inv) as f32 } else { 0.0 };
493                }
494            }
495        }
496        // The aggregate projection is 4096x188160 — one weight buffer past
497        // what a GPU binding may address, and it runs twice per prompt, not
498        // per step. Keep it on the CPU rather than teach the device to page
499        // a 2 GiB binding for a millisecond of work.
500        let project = |agg: &Lin, out_dim: usize| -> Vec<f32> {
501            let scale = ((out_dim as f64) / (d as f64)).sqrt() as f32;
502            let scaled: Vec<f32> = feats.iter().map(|&v| v * scale).collect();
503            crate::gpu::cpu_scope(|| agg.apply(&scaled, t, pool))
504        };
505        let vfeat = project(&self.video_agg, 4096);
506        let afeat = project(&self.audio_agg, 2048);
507        // the connectors want valid tokens first; the prompt arrives left-padded
508        let order: Vec<usize> = (0..t)
509            .filter(|&i| mask[i] != 0.0)
510            .chain((0..t).filter(|&i| mask[i] == 0.0))
511            .collect();
512        let valid = mask.iter().filter(|&&m| m != 0.0).count();
513        let reorder = |x: &[f32], dim: usize| -> Vec<f32> {
514            let mut o = vec![0f32; t * dim];
515            for (new, &old) in order.iter().enumerate() {
516                o[new * dim..(new + 1) * dim].copy_from_slice(&x[old * dim..(old + 1) * dim]);
517            }
518            o
519        };
520        let v = self.v_conn.run(&reorder(&vfeat, 4096), t, valid, pool);
521        let a = self.a_conn.run(&reorder(&afeat, 2048), t, valid, pool);
522        (v, a, t)
523    }
524}
525
526impl Connector {
527    /// Eight gated-attention blocks over the whole window. Padded positions
528    /// are first replaced by the learnable registers, tiled across the
529    /// window, which is what makes the mask vanish: after the substitution
530    /// every position is signal and attention is unmasked.
531    fn run(&self, x: &[f32], t: usize, valid: usize, pool: Option<&Pool>) -> Vec<f32> {
532        let d = self.dim;
533        let regs = self.registers.len() / d;
534        let mut h = x.to_vec();
535        for i in valid..t {
536            let r = i % regs;
537            h[i * d..(i + 1) * d].copy_from_slice(&self.registers[r * d..(r + 1) * d]);
538        }
539        let pos: Vec<Vec<f64>> = (0..t).map(|i| vec![i as f64]).collect();
540        let pe = Rope::build(&pos, &[self.max_pos], d, self.heads, 10000.0);
541        for (attn, ff_in, ff_out) in &self.blocks {
542            let mut n = vec![0f32; t * d];
543            for i in 0..t {
544                rms_plain(&h[i * d..(i + 1) * d], &mut n[i * d..(i + 1) * d]);
545            }
546            let a = attn.forward(&n, t, &n, t, Some(&pe), Some(&pe), None, pool);
547            for (v, &y) in h.iter_mut().zip(&a) {
548                *v += y;
549            }
550            let mut n2 = vec![0f32; t * d];
551            for i in 0..t {
552                rms_plain(&h[i * d..(i + 1) * d], &mut n2[i * d..(i + 1) * d]);
553            }
554            let mut g = ff_in.apply(&n2, t, pool);
555            for v in g.iter_mut() {
556                *v = gelu_tanh(*v);
557            }
558            let f = ff_out.apply(&g, t, pool);
559            for (v, &y) in h.iter_mut().zip(&f) {
560                *v += y;
561            }
562        }
563        let mut out = vec![0f32; t * d];
564        for i in 0..t {
565            rms_plain(&h[i * d..(i + 1) * d], &mut out[i * d..(i + 1) * d]);
566        }
567        let _ = self.dh;
568        out
569    }
570}