Skip to main content

cortiq_engine/
mmh3.rs

1//! MiniMax-H3: the packed-token audio-video DiT.
2//!
3//! One stream of tokens — `[text | audio | video]` — denoised by fifty
4//! blocks of full self-attention, 3-axis RoPE and adaLN modulation, with
5//! the video and audio latents riding different flow schedules and a
6//! separate output head each.
7//!
8//! Three things make it unlike the image DiT next door:
9//!
10//! * **The sequence is packed, not batched.** Text, audio rows and video
11//!   rows sit in one sequence and attend to each other; there is no
12//!   cross-attention. Modulation is what tells a token which modality it
13//!   is: every block emits six modulation vectors for each of three
14//!   modality tags at each of the (at most two, for t2va) distinct
15//!   timesteps, and a segment table says which row each token reads.
16//!
17//! * **adaLN arrives as a curve, not a matrix.** The released weight is
18//!   [96768, 2688] per block — 13 B parameters, 40% of the model — for a
19//!   map whose input is one number. `cortiq animate-pack` collapses it
20//!   onto a rank-24 basis of the timestep curve that already carries the
21//!   Turbo LoRA, so what lands here is [96768, 24] and a [1025, 24]
22//!   table to interpolate. Measured against the full matrix: rms 8.7e-5
23//!   on a signal of rms 0.46.
24//!
25//! * **Two clocks.** The sampler hands in the video sigma; the audio
26//!   stream's own sigma is a closed-form remap of it (shift 12 → 3), and
27//!   the two are integrated separately. `forward` returns both
28//!   velocities unscaled, each on its own schedule — the reference
29//!   returns the audio one pre-multiplied by d(σ_a)/d(σ_v) so that a
30//!   single-schedule sampler is approximately right, which at four steps
31//!   it is not.
32//!
33//! Parity: `tools/mk_mmh3_toy.py` builds a toy checkpoint carrying the
34//! release's real tensor names and a golden forward from ComfyUI's own
35//! module; `tools/mmh3_toy_gate.sh` diffs this port against it.
36
37use crate::dit::Proj;
38use crate::ltxlora::{LoraBank, LoraBranch};
39
40/// Per-phase microseconds of the DiT block, under `CMF_MMH3_PROF=1`:
41/// 0 norm+modulate · 1 qkv GEMM · 2 qk-norm+RoPE · 3 attention ·
42/// 4 out GEMM+residual · 5 fc1 GEMM · 6 SwiGLU · 7 fc2 GEMM.
43pub static MMH3_PROF: [std::sync::atomic::AtomicU64; 9] = [
44    std::sync::atomic::AtomicU64::new(0),
45    std::sync::atomic::AtomicU64::new(0),
46    std::sync::atomic::AtomicU64::new(0),
47    std::sync::atomic::AtomicU64::new(0),
48    std::sync::atomic::AtomicU64::new(0),
49    std::sync::atomic::AtomicU64::new(0),
50    std::sync::atomic::AtomicU64::new(0),
51    std::sync::atomic::AtomicU64::new(0),
52    std::sync::atomic::AtomicU64::new(0),
53];
54
55pub(crate) fn mmh3_prof_on() -> bool {
56    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
57    *ON.get_or_init(|| std::env::var("CMF_MMH3_PROF").is_ok())
58}
59
60/// One line per phase, sorted by cost — the map that says which kernel
61/// to write next.
62pub fn mmh3_prof_report() -> Option<String> {
63    if !mmh3_prof_on() {
64        return None;
65    }
66    const NAMES: [&str; 9] = [
67        "norm+mod",
68        "qkv gemm",
69        "qknorm+rope",
70        "attention",
71        "out gemm",
72        "fc1 gemm",
73        "swiglu",
74        "fc2 gemm",
75        "residual",
76    ];
77    let mut v: Vec<(u64, &str)> = MMH3_PROF
78        .iter()
79        .map(|a| a.load(std::sync::atomic::Ordering::Relaxed))
80        .zip(NAMES)
81        .collect();
82    let total: u64 = v.iter().map(|(us, _)| *us).sum();
83    v.sort_by(|a, b| b.0.cmp(&a.0));
84    let mut out = format!("mmh3 phases (total {:.1} s):", total as f64 / 1e6);
85    for (us, name) in v {
86        out.push_str(&format!(
87            "\n  {name:<14} {:>7.1} s  {:>5.1}%",
88            us as f64 / 1e6,
89            100.0 * us as f64 / total.max(1) as f64
90        ));
91    }
92    Some(out)
93}
94
95use crate::pool::Pool;
96use cortiq_core::CmfModel;
97use std::sync::Arc;
98
99/// Frames a video latent token spans, cycling with period 5 — the
100/// checkpoint's temporal grid, which is NOT uniform.
101const FRAME_PER_TOKEN: [f64; 5] = [1.0, 4.0, 4.0, 4.0, 4.0];
102const FRAME_RESCALE: f64 = 5.0 / 3.0;
103
104/// Modality tags, as the adaLN row layout orders them.
105const TAG_VIDEO: usize = 0;
106const TAG_TEXT: usize = 1;
107const TAG_AUDIO: usize = 2;
108const MODALITIES: usize = 3;
109/// shift/scale/gate for attention, then the same three for the MLP.
110const EXPAND: usize = 6;
111/// Where a visual condition row sits on the schedule: all but arrived.
112const VISUAL_COND_TIMESTEP: f64 = 0.999;
113/// An audio condition is not noised at all.
114const AUDIO_COND_TIMESTEP: f64 = 1.0;
115
116// ── the flow-schedule remap ─────────────────────────────────────────
117
118/// σ on `to`'s schedule at the same point of the shared base grid.
119pub fn time_shift_sigma(sigma: f64, from_shift: f64, to_shift: f64) -> f64 {
120    let base = sigma / (from_shift + sigma * (1.0 - from_shift));
121    to_shift * base / (1.0 + (to_shift - 1.0) * base)
122}
123
124/// d(σ_to)/d(σ_from) at the same base-grid point.
125pub fn time_shift_slope(sigma: f64, from_shift: f64, to_shift: f64) -> f64 {
126    let base = sigma / (from_shift + sigma * (1.0 - from_shift));
127    (to_shift * (1.0 + (from_shift - 1.0) * base).powi(2))
128        / (from_shift * (1.0 + (to_shift - 1.0) * base).powi(2))
129}
130
131// ── the packed layout ───────────────────────────────────────────────
132
133#[derive(Clone, Copy, PartialEq, Debug)]
134pub enum Kind {
135    Text,
136    /// A keyframe's latent, re-injected every step and never denoised.
137    Cond,
138    /// A reference block's video rows — an image, or a clip's frames.
139    RefImg,
140    /// A reference block's audio rows.
141    RefAudio,
142    Audio,
143    Video,
144}
145
146/// One contiguous run of rows sharing a modality tag and a timestep.
147#[derive(Clone, Copy, Debug)]
148pub struct Segment {
149    pub start: usize,
150    pub stop: usize,
151    pub kind: Kind,
152}
153
154/// The static structure of one shape signature: where each stream sits
155/// in the sequence and what 3-D position every row carries.
156pub struct Layout {
157    pub seq_len: usize,
158    pub segments: Vec<Segment>,
159    /// [seq_len, 3] — (t, h, w), f64 because the axes are fractional.
160    pub pos: Vec<[f64; 3]>,
161    pub text_len: usize,
162    pub audio_t: usize,
163    pub latent_t: usize,
164    pub lat_h: usize,
165    pub lat_w: usize,
166    /// Rows per latent frame after the 2×2 patch.
167    pub frame_rows: usize,
168    /// Per-token modality tag for the text span. A vision block inside
169    /// the prompt carries the VIDEO tag, not the text one.
170    pub text_tags: Vec<u8>,
171}
172
173/// One reference block, in the order it was given.
174#[derive(Clone, Debug)]
175pub enum Ref {
176    /// A still, at its own latent size.
177    Image { lat_h: usize, lat_w: usize },
178    /// Standalone audio, `t` latent frames of it.
179    Audio { t: usize },
180    /// Frames, with an optional soundtrack that packs immediately
181    /// before them and shares their origin.
182    Video {
183        latent_t: usize,
184        lat_h: usize,
185        lat_w: usize,
186        audio_t: usize,
187    },
188}
189
190/// Channel-major stereo rows: `t` frames per channel, the two channels
191/// pinned to the grid's extreme w coordinates so RoPE can tell them
192/// apart, h flat at zero.
193fn push_audio_grid(pos: &mut Vec<[f64; 3]>, cursor: f64, t: usize, w_low: f64, w_high: f64) {
194    for ch in 0..2 {
195        let w = if ch == 0 { w_low } else { w_high };
196        for i in 0..t {
197            pos.push([cursor + i as f64, 0.0, w]);
198        }
199    }
200}
201
202/// `linspace((1 − ratio)/2, (1 + ratio)/2, dim/patch, endpoint=False) · 32`
203fn axis_from_sqrt_area(dim: usize, patch: usize, sqrt_area: f64) -> Vec<f64> {
204    let ratio = dim as f64 / sqrt_area;
205    let n = dim / patch;
206    (0..n)
207        .map(|i| (i as f64 * (ratio / n as f64) + (1.0 - ratio) / 2.0) * 32.0)
208        .collect()
209}
210
211impl Layout {
212    /// t2va: `[text | audio | video]`, the target streams last and in
213    /// that order. Keyframe and reference blocks would slot between the
214    /// text and the audio; this port does text-to-video only.
215    pub fn t2va(
216        text_len: usize,
217        latent_t: usize,
218        lat_h: usize,
219        lat_w: usize,
220        audio_t: usize,
221    ) -> Self {
222        Self::build(text_len, latent_t, lat_h, lat_w, audio_t, &[], &[])
223    }
224
225    /// `fl2va`: keyframe condition rows sit between the text and the
226    /// audio, sharing the TARGET spatial grid, each pinned to the time
227    /// coordinate of the frame it stands for — the first frame at the
228    /// text's end, the last one a whole clip further on, minus one
229    /// span. They never advance the cursor, so audio and video still
230    /// start where they would have.
231    ///
232    /// `frames` gives each keyframe's pixel index and the clip's total,
233    /// and `text_tags` the per-token modality of the prompt span.
234    pub fn fl2va(
235        text_len: usize,
236        latent_t: usize,
237        lat_h: usize,
238        lat_w: usize,
239        audio_t: usize,
240        frames: &[(usize, usize)],
241        text_tags: &[u8],
242    ) -> Self {
243        Self::build(text_len, latent_t, lat_h, lat_w, audio_t, frames, text_tags)
244    }
245
246    /// `ref2va`: reference images, audio and clips ahead of the target
247    /// streams. Unlike a keyframe, a reference ADVANCES the cursor —
248    /// each block occupies its own stretch of the time axis, and the
249    /// target audio and video begin after the last of them.
250    pub fn ref2va(
251        text_len: usize,
252        latent_t: usize,
253        lat_h: usize,
254        lat_w: usize,
255        audio_t: usize,
256        refs: &[Ref],
257        text_tags: &[u8],
258    ) -> Self {
259        Self::build_full(
260            text_len,
261            latent_t,
262            lat_h,
263            lat_w,
264            audio_t,
265            &[],
266            refs,
267            text_tags,
268        )
269    }
270
271    fn build(
272        text_len: usize,
273        latent_t: usize,
274        lat_h: usize,
275        lat_w: usize,
276        audio_t: usize,
277        frames: &[(usize, usize)],
278        text_tags: &[u8],
279    ) -> Self {
280        Self::build_full(
281            text_len,
282            latent_t,
283            lat_h,
284            lat_w,
285            audio_t,
286            frames,
287            &[],
288            text_tags,
289        )
290    }
291
292    #[allow(clippy::too_many_arguments)]
293    fn build_full(
294        text_len: usize,
295        latent_t: usize,
296        lat_h: usize,
297        lat_w: usize,
298        audio_t: usize,
299        frames: &[(usize, usize)],
300        refs: &[Ref],
301        text_tags: &[u8],
302    ) -> Self {
303        let area = ((lat_h * lat_w) as f64).sqrt();
304        let h_axis = axis_from_sqrt_area(lat_h, 2, area);
305        let w_axis = axis_from_sqrt_area(lat_w, 2, area);
306        let frame_rows = h_axis.len() * w_axis.len();
307
308        let mut pos: Vec<[f64; 3]> = Vec::new();
309        let mut segments = Vec::new();
310
311        segments.push(Segment {
312            start: 0,
313            stop: text_len,
314            kind: Kind::Text,
315        });
316        for i in 0..text_len {
317            pos.push([i as f64, 0.0, 0.0]);
318        }
319
320        // Both target streams share this origin: the text runs out at
321        // `text_len` and audio and video start together from there —
322        // unless references push it along.
323        let mut cursor = text_len as f64;
324        let (w_low_t, w_high_t) = (w_axis[0], w_axis[w_axis.len() - 1]);
325
326        for r in refs {
327            match r {
328                Ref::Image { lat_h, lat_w } => {
329                    let (rh, rw) = (
330                        axis_from_sqrt_area(*lat_h, 2, ((lat_h * lat_w) as f64).sqrt()),
331                        axis_from_sqrt_area(*lat_w, 2, ((lat_h * lat_w) as f64).sqrt()),
332                    );
333                    let start = pos.len();
334                    for &h in &rh {
335                        for &w in &rw {
336                            pos.push([cursor, h, w]);
337                        }
338                    }
339                    segments.push(Segment {
340                        start,
341                        stop: pos.len(),
342                        kind: Kind::RefImg,
343                    });
344                    cursor += 1.0;
345                }
346                Ref::Audio { t } => {
347                    if *t > 0 {
348                        let start = pos.len();
349                        push_audio_grid(&mut pos, cursor, *t, w_low_t, w_high_t);
350                        segments.push(Segment {
351                            start,
352                            stop: pos.len(),
353                            kind: Kind::RefAudio,
354                        });
355                    }
356                    cursor += *t as f64;
357                }
358                Ref::Video {
359                    latent_t: vt,
360                    lat_h: rh_,
361                    lat_w: rw_,
362                    audio_t: rt,
363                } => {
364                    let area = ((rh_ * rw_) as f64).sqrt();
365                    let rh = axis_from_sqrt_area(*rh_, 2, area);
366                    let rw = axis_from_sqrt_area(*rw_, 2, area);
367                    // The block's audio packs immediately BEFORE its
368                    // frames, both from the same origin, and takes its
369                    // w extremes from the block's own grid.
370                    if *rt > 0 {
371                        let start = pos.len();
372                        push_audio_grid(&mut pos, cursor, *rt, rw[0], rw[rw.len() - 1]);
373                        segments.push(Segment {
374                            start,
375                            stop: pos.len(),
376                            kind: Kind::RefAudio,
377                        });
378                    }
379                    let start = pos.len();
380                    let mut t_coord = cursor;
381                    for k in 0..*vt {
382                        for &h in &rh {
383                            for &w in &rw {
384                                pos.push([t_coord, h, w]);
385                            }
386                        }
387                        t_coord += FRAME_RESCALE * FRAME_PER_TOKEN[k % 5];
388                    }
389                    segments.push(Segment {
390                        start,
391                        stop: pos.len(),
392                        kind: Kind::RefImg,
393                    });
394                    let spans: f64 = (0..*vt)
395                        .map(|k| FRAME_RESCALE * FRAME_PER_TOKEN[k % 5])
396                        .sum();
397                    cursor += (*rt as f64).max(spans);
398                }
399            }
400        }
401        let cursor = cursor;
402
403        // Keyframes, in the order given.
404        let spans: f64 = (0..latent_t)
405            .map(|k| FRAME_RESCALE * FRAME_PER_TOKEN[k % 5])
406            .sum();
407        for &(pixel_index, frame_count) in frames {
408            let cond_t = if pixel_index == 0 {
409                cursor
410            } else if frame_count > 0 && pixel_index == frame_count - 1 {
411                cursor + spans - FRAME_RESCALE
412            } else {
413                panic!("only the first and last frame can anchor a keyframe");
414            };
415            let start = pos.len();
416            for &h in &h_axis {
417                for &w in &w_axis {
418                    pos.push([cond_t, h, w]);
419                }
420            }
421            segments.push(Segment {
422                start,
423                stop: pos.len(),
424                kind: Kind::Cond,
425            });
426        }
427
428        // Audio is channel-major stereo: every latent frame once per
429        // channel, the two channels pinned to the frame grid's extreme
430        // w coordinates so they are distinguishable under RoPE.
431        let a_start = pos.len();
432        push_audio_grid(&mut pos, cursor, audio_t, w_low_t, w_high_t);
433        segments.push(Segment {
434            start: a_start,
435            stop: pos.len(),
436            kind: Kind::Audio,
437        });
438
439        // Video: the t axis advances by the per-token frame spans, not
440        // by one; h/w come from the shared frame grid.
441        let v_start = pos.len();
442        let mut t_coord = cursor;
443        for k in 0..latent_t {
444            for &h in &h_axis {
445                for &w in &w_axis {
446                    pos.push([t_coord, h, w]);
447                }
448            }
449            t_coord += FRAME_RESCALE * FRAME_PER_TOKEN[k % 5];
450        }
451        segments.push(Segment {
452            start: v_start,
453            stop: pos.len(),
454            kind: Kind::Video,
455        });
456
457        Self {
458            seq_len: pos.len(),
459            segments,
460            pos,
461            text_len,
462            audio_t,
463            latent_t,
464            lat_h,
465            lat_w,
466            frame_rows,
467            text_tags: if text_tags.is_empty() {
468                vec![TAG_TEXT as u8; text_len]
469            } else {
470                text_tags.to_vec()
471            },
472        }
473    }
474
475    /// How many keyframe condition rows the layout carries.
476    pub fn cond_rows(&self) -> usize {
477        self.segments
478            .iter()
479            .filter(|s| s.kind == Kind::Cond)
480            .map(|s| s.stop - s.start)
481            .sum()
482    }
483
484    fn segment(&self, kind: Kind) -> Segment {
485        *self
486            .segments
487            .iter()
488            .find(|s| s.kind == kind)
489            .expect("every layout carries all three streams")
490    }
491}
492
493// ── weights ─────────────────────────────────────────────────────────
494
495struct Adaln {
496    /// [out, rank] — the curve basis, already carrying the LoRA.
497    w: Proj,
498    b: Vec<f32>,
499    /// [grid, rank]
500    table: Vec<f32>,
501    rank: usize,
502    out: usize,
503}
504
505impl Adaln {
506    fn load(model: &Arc<CmfModel>, prefix: &str) -> Result<Self, String> {
507        let w = Proj::from_model(model, &format!("{prefix}.weight"))?;
508        let table = crate::dit::cmf_f32(model, &format!("{prefix}.table"))?;
509        let b = crate::dit::cmf_f32(model, &format!("{prefix}.bias"))?;
510        let out = w.rows();
511        let rank = table.len() / CURVE_GRID;
512        Ok(Self {
513            w,
514            b,
515            table,
516            rank,
517            out,
518        })
519    }
520
521    /// The modulation vectors at `ts`: `[ts.len() · MODALITIES, expand ·
522    /// hidden]` laid out so row `t·MODALITIES + tag`, chunk `e`, starts
523    /// at `(t·MODALITIES + tag)·expand·hidden + e·hidden`.
524    fn eval(&self, ts: &[f64], pool: Option<&Pool>) -> Vec<f32> {
525        let g = CURVE_GRID;
526        let mut coords = vec![0f32; ts.len() * self.rank];
527        for (i, &t) in ts.iter().enumerate() {
528            // t → fractional grid index; out-of-range clamps to the ends,
529            // and the last interval is kept whole so t = 1 does not read
530            // past the table.
531            let p = (t.clamp(0.0, 1.0) * (g - 1) as f64) as f32;
532            let i0 = (p.floor() as usize).min(g - 2);
533            let f = p - i0 as f32;
534            for k in 0..self.rank {
535                let (a, b) = (
536                    self.table[i0 * self.rank + k],
537                    self.table[(i0 + 1) * self.rank + k],
538                );
539                coords[i * self.rank + k] = a + (b - a) * f;
540            }
541        }
542        let mut out = vec![0f32; ts.len() * self.out];
543        self.w.matmat(&coords, ts.len(), &mut out, pool);
544        for row in out.chunks_exact_mut(self.out) {
545            for (v, &bv) in row.iter_mut().zip(&self.b) {
546                *v += bv;
547            }
548        }
549        out
550    }
551}
552
553struct Block {
554    norm1: Vec<f32>,
555    norm2: Vec<f32>,
556    qkv: Proj, // [3·heads·hd, hidden]
557    out: Proj, // [hidden, heads·hd]
558    q_norm: Vec<f32>,
559    k_norm: Vec<f32>,
560    fc1: Proj, // [2·ffn, hidden]
561    fc2: Proj, // [hidden, ffn]
562    adaln: Option<Adaln>,
563    /// Runtime low-rank branches, one per projection an adapter names.
564    /// A branch is why the fused device paths below stand down: they
565    /// keep qkv, the attention output and the FFN's middle on the card,
566    /// and the branch has to read exactly those.
567    lora: BlockLora,
568}
569
570/// The four projections an adapter can reach in a packed block. `adaln`
571/// is deliberately absent: this container carries the modulation as a
572/// rank-24 curve, and folding an adaLN update into it needs the time
573/// embedding the packer had (`animate-pack --lora --time-embedder`).
574#[derive(Default)]
575struct BlockLora {
576    qkv: Option<LoraBranch>,
577    out: Option<LoraBranch>,
578    fc1: Option<LoraBranch>,
579    fc2: Option<LoraBranch>,
580}
581
582/// A projection and its branch, in one device submission where the platform
583/// has that kernel and the shapes fit — the branch reads the activation the
584/// base GEMM already uploaded and accumulates into the output it just wrote,
585/// so it costs no transfer of its own. Falls back to `matmat` + `add`, which
586/// is correct everywhere and pays a round trip for the branch.
587///
588/// The router and the probe both need the branch and the base separated to
589/// measure one against the other, which the fused kernel does not do — so
590/// when either is asked for, a branch takes the split path until it HAS been
591/// measured, and the fused one from the next step onward. Without that the
592/// probe reports 0.0000 for every branch and reads as "this adapter does
593/// nothing", which is what it did in the first run of it.
594fn proj_with_lora(
595    p: &Proj,
596    br: &Option<LoraBranch>,
597    x: &[f32],
598    n: usize,
599    out: &mut [f32],
600    pool: Option<&Pool>,
601) {
602    #[cfg(target_os = "macos")]
603    if let Some(l) = br.as_ref().filter(|l| l.live()) {
604        if let Some((model, idx)) = p.q4tp_mapped() {
605            let (rows, cols) = (p.rows(), p.cols());
606            if n >= 32
607                && cols % 32 == 0
608                && crate::gpu::enabled_here()
609                && !crate::gpu::mm_killed()
610                && (!crate::ltxlora::wants_measurement() || l.resonance() != 0.0)
611                && crate::gpu_metal::q4tp_matmat_lora(model, idx, x, n, rows, cols, out, &l.side())
612            {
613                return;
614            }
615        }
616    }
617    p.matmat(x, n, out, pool);
618    if let Some(l) = br {
619        l.add(x, n, out, pool);
620    }
621}
622
623impl BlockLora {
624    /// Live, not merely present: a branch the router has switched off
625    /// gives the fused device path back for the rest of the render.
626    fn live(br: &Option<LoraBranch>) -> bool {
627        br.as_ref().is_some_and(|l| l.live())
628    }
629    fn ffn_any(&self) -> bool {
630        Self::live(&self.fc1) || Self::live(&self.fc2)
631    }
632}
633
634pub(crate) const CURVE_GRID: usize = 1025;
635
636pub struct MiniMaxH3 {
637    video_patch: Proj,
638    video_patch_b: Vec<f32>,
639    audio_patch: Proj,
640    audio_patch_b: Vec<f32>,
641    condition: Proj,
642    condition_b: Vec<f32>,
643    refiner: Vec<Block>,
644    refiner_norm: Vec<f32>,
645    blocks: Vec<Block>,
646    final_norm: Vec<f32>,
647    final_adaln: Adaln,
648    video_out: Proj,
649    video_out_b: Vec<f32>,
650    audio_out: Proj,
651    audio_out_b: Vec<f32>,
652    inv_freq: Vec<f32>,
653    pool: Option<Arc<Pool>>,
654    pub hidden: usize,
655    pub heads: usize,
656    pub head_dim: usize,
657    pub ffn: usize,
658    pub latents_dim: usize,
659    pub audio_dim: usize,
660    pub text_dim: usize,
661    pub shift_video: f64,
662    pub shift_audio: f64,
663    eps: f64,
664    qk_eps: f64,
665    final_eps: f64,
666    /// How much of a keyframe latent survives the noise blend, and
667    /// therefore where its rows sit on the schedule. `VISUAL_COND_
668    /// TIMESTEP` is the reference's default; 1.0 turns the blend off.
669    pub cond_aug: f64,
670    /// The same, for a reference soundtrack. The reference's default is
671    /// 1.0 — an audio condition is not noised at all.
672    pub cond_aug_audio: f64,
673    /// Adapter keys that found a projection, for the caller's report.
674    lora_bound: std::collections::BTreeSet<String>,
675}
676
677fn rms_norm_into(x: &[f32], w: &[f32], eps: f64, dst: &mut [f32]) {
678    let ss = x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64;
679    let inv = 1.0 / (ss + eps).sqrt();
680    for ((d, &v), &g) in dst.iter_mut().zip(x).zip(w) {
681        *d = (v as f64 * inv) as f32 * g;
682    }
683}
684
685fn silu(v: f32) -> f32 {
686    v / (1.0 + (-v).exp())
687}
688
689impl MiniMaxH3 {
690    pub fn from_cmf(model: &Arc<CmfModel>) -> Result<Self, String> {
691        Self::from_cmf_lora(model, None)
692    }
693
694    /// The same, with an adapter consulted for every projection.
695    ///
696    /// The base weights are q4tp and stay untouched: a rank-32 update
697    /// cannot be folded into a four-bit ladder without dequantizing the
698    /// whole DiT, so the branch rides beside the base GEMM the way
699    /// [`crate::ltxlora`] does it for LTX.
700    pub fn from_cmf_lora(
701        model: &Arc<CmfModel>,
702        bank: Option<&LoraBank>,
703    ) -> Result<Self, String> {
704        let cfg: serde_json::Value = serde_json::from_slice(
705            model
706                .tensor_bytes("dit.config_json")
707                .map_err(|e| e.to_string())?,
708        )
709        .map_err(|e| format!("dit.config_json: {e}"))?;
710        let u = |k: &str| cfg[k].as_u64().unwrap_or(0) as usize;
711        let f = |k: &str, d: f64| cfg[k].as_f64().unwrap_or(d);
712        let n_blocks = u("num_layers");
713        let n_refiner = u("token_refiner_num_layers");
714        let f32v = |n: &str| crate::dit::cmf_f32(model, n);
715
716        // Which of the adapter's names actually found a projection here.
717        let bound: std::cell::RefCell<std::collections::BTreeSet<String>> = Default::default();
718        let load_block = |prefix: &str, with_adaln: bool| -> Result<Block, String> {
719            let qkv = Proj::from_model(model, &format!("dit.{prefix}.attn.qkv_proj.weight"))?;
720            let out = Proj::from_model(model, &format!("dit.{prefix}.attn.out_proj.weight"))?;
721            let fc1 = Proj::from_model(model, &format!("dit.{prefix}.mlp.fc1.weight"))?;
722            let fc2 = Proj::from_model(model, &format!("dit.{prefix}.mlp.fc2.weight"))?;
723            // A branch is bound against the SHAPE of the projection it will
724            // ride: `add` writes n·out floats through a raw pointer, so an
725            // adapter for another model has to be refused by name, not
726            // discovered as a corrupted panel.
727            let lora = match bank {
728                None => BlockLora::default(),
729                Some(k) => {
730                    let mut take = |suffix: &str, p: &Proj| -> Result<Option<LoraBranch>, String> {
731                        let key = format!("{prefix}.{suffix}");
732                        let br = k.branch_for(&format!("dit.{key}"), p.rows(), p.cols())?;
733                        if br.is_some() {
734                            bound.borrow_mut().insert(key);
735                        }
736                        Ok(br)
737                    };
738                    BlockLora {
739                        qkv: take("attn.qkv_proj", &qkv)?,
740                        out: take("attn.out_proj", &out)?,
741                        fc1: take("mlp.fc1", &fc1)?,
742                        fc2: take("mlp.fc2", &fc2)?,
743                    }
744                }
745            };
746            Ok(Block {
747                norm1: f32v(&format!("dit.{prefix}.norm1.weight"))?,
748                norm2: f32v(&format!("dit.{prefix}.norm2.weight"))?,
749                qkv,
750                out,
751                q_norm: f32v(&format!("dit.{prefix}.attn.q_norm.weight"))?,
752                k_norm: f32v(&format!("dit.{prefix}.attn.k_norm.weight"))?,
753                fc1,
754                fc2,
755                adaln: if with_adaln {
756                    Some(Adaln::load(model, &format!("dit.{prefix}.adaln"))?)
757                } else {
758                    None
759                },
760                lora,
761            })
762        };
763        let blocks = (0..n_blocks)
764            .map(|i| load_block(&format!("blocks.{i}"), true))
765            .collect::<Result<Vec<_>, _>>()?;
766        let refiner = (0..n_refiner)
767            .map(|i| load_block(&format!("token_refiner.blocks.{i}"), false))
768            .collect::<Result<Vec<_>, _>>()?;
769
770        Ok(Self {
771            video_patch: Proj::from_model(model, "dit.video_patch_proj.weight")?,
772            video_patch_b: f32v("dit.video_patch_proj.bias")?,
773            audio_patch: Proj::from_model(model, "dit.audio_patch_proj.weight")?,
774            audio_patch_b: f32v("dit.audio_patch_proj.bias")?,
775            condition: Proj::from_model(model, "dit.condition_proj.weight")?,
776            condition_b: f32v("dit.condition_proj.bias")?,
777            refiner,
778            refiner_norm: f32v("dit.token_refiner.final_norm.weight")?,
779            blocks,
780            final_norm: f32v("dit.final_layer.norm.weight")?,
781            final_adaln: Adaln::load(model, "dit.final_layer.adaln")?,
782            video_out: Proj::from_model(model, "dit.final_layer.video_out.weight")?,
783            video_out_b: f32v("dit.final_layer.video_out.bias")?,
784            audio_out: Proj::from_model(model, "dit.final_layer.audio_out.weight")?,
785            audio_out_b: f32v("dit.final_layer.audio_out.bias")?,
786            inv_freq: f32v("dit.rope_inv_freq")?,
787            pool: Pool::from_env(),
788            hidden: u("hidden_size"),
789            heads: u("num_attention_heads"),
790            head_dim: u("attention_head_dim"),
791            ffn: u("ffn_hidden_size"),
792            latents_dim: u("latents_dim"),
793            audio_dim: u("audio_latents_dim"),
794            text_dim: u("text_dim"),
795            shift_video: f("sigma_shift_video", 12.0),
796            shift_audio: f("sigma_shift_audio", 3.0),
797            eps: f("norm_eps", 1e-5),
798            qk_eps: f("qk_norm_eps", 1e-5),
799            final_eps: f("final_norm_eps", 1e-5),
800            cond_aug: VISUAL_COND_TIMESTEP,
801            cond_aug_audio: AUDIO_COND_TIMESTEP,
802            lora_bound: bound.into_inner(),
803        })
804    }
805
806    /// How many of the adapter's branches found a projection.
807    pub fn lora_bound(&self) -> usize {
808        self.lora_bound.len()
809    }
810
811    /// Whether this adapter key landed on a projection of ours.
812    pub fn lora_binds(&self, key: &str) -> bool {
813        self.lora_bound.contains(key)
814    }
815
816    /// Under `CMF_LORA_PROBE=1`: every branch by measured contribution,
817    /// loudest first, and which ones the router switched off. This is the
818    /// map that says where an adapter actually lives — for the Realism
819    /// adapter it is not uniform across the fifty blocks.
820    pub fn lora_report(&self) -> Option<String> {
821        if !crate::ltxlora::probe_report_on() {
822            return None;
823        }
824        let mut rows: Vec<(f32, bool, String)> = Vec::new();
825        let mut walk = |blocks: &[Block], pre: &str| {
826            for (i, b) in blocks.iter().enumerate() {
827                for (name, br) in [
828                    ("attn.qkv_proj", &b.lora.qkv),
829                    ("attn.out_proj", &b.lora.out),
830                    ("mlp.fc1", &b.lora.fc1),
831                    ("mlp.fc2", &b.lora.fc2),
832                ] {
833                    if let Some(l) = br {
834                        rows.push((l.resonance(), l.live(), format!("{pre}{i}.{name}")));
835                    }
836                }
837            }
838        };
839        walk(&self.blocks, "blocks.");
840        walk(&self.refiner, "token_refiner.blocks.");
841        if rows.is_empty() {
842            return None;
843        }
844        rows.sort_by(|a, b| b.0.total_cmp(&a.0));
845        let live = rows.iter().filter(|r| r.1).count();
846        let mut s = format!(
847            "lora branches by contribution ‖sΔY‖/‖Y‖ ({} of {} live):\n",
848            live,
849            rows.len()
850        );
851        for (r, on, name) in &rows {
852            s.push_str(&format!(
853                "  {:>8.4}  {}  {name}\n",
854                r,
855                if *on { "on " } else { "off" }
856            ));
857        }
858        Some(s)
859    }
860
861    /// Qwen3-VL states `[n, text_dim]` → refined text embeds
862    /// `[n, hidden]`. Prompt-only, so the caller does this once per
863    /// generation rather than once per step.
864    pub fn refine_text(&self, states: &[f32], n: usize) -> Vec<f32> {
865        let pool = self.pool.as_deref();
866        let mut h = vec![0f32; n * self.hidden];
867        self.condition.matmat(states, n, &mut h, pool);
868        for row in h.chunks_exact_mut(self.hidden) {
869            for (v, &b) in row.iter_mut().zip(&self.condition_b) {
870                *v += b;
871            }
872        }
873        let ids: Vec<[f64; 3]> = Vec::new();
874        for blk in &self.refiner {
875            self.block_forward(blk, &mut h, n, None, &ids, &[]);
876        }
877        let mut out = vec![0f32; n * self.hidden];
878        for (o, x) in out
879            .chunks_exact_mut(self.hidden)
880            .zip(h.chunks_exact(self.hidden))
881        {
882            rms_norm_into(x, &self.refiner_norm, self.final_eps, o);
883        }
884        out
885    }
886
887    /// The rotation angles of every row: `[n, 48]`. Each of the three
888    /// position axes contributes `inv_freq.len()` angles, and the pair
889    /// (j, j+48) of the first 96 head dims rotates by angle j.
890    fn rope_angles(&self, pos: &[[f64; 3]]) -> Vec<f32> {
891        let k = self.inv_freq.len();
892        let mut out = Vec::with_capacity(pos.len() * 3 * k);
893        for p in pos {
894            for axis in 0..3 {
895                for j in 0..k {
896                    out.push((p[axis] * self.inv_freq[j] as f64) as f32);
897                }
898            }
899        }
900        out
901    }
902
903    /// Per-head RMSNorm then the partial split-half rotation, in place.
904    /// `angles` carries 48 angles a row — three axes × 16 frequencies —
905    /// and pair (j, j+48) of the head's first 96 dims turns by angle j.
906    /// Dims 96..128 are not rotated at all, which is the checkpoint's
907    /// own `rope_inv_freq_len` arithmetic and not a truncation.
908    /// In place on a STRIDED view of the fused qkv buffer: `stride` is
909    /// the row pitch and `off` the plane's start. Normalizing q and k
910    /// through a scratch copy cost four full passes over `n·heads·hd`
911    /// per block — 10 G element copies over a render — for arithmetic
912    /// that touches each element once.
913    #[allow(clippy::too_many_arguments)]
914    fn norm_rope_w(
915        &self,
916        v: &mut [f32],
917        n: usize,
918        heads: usize,
919        w: &[f32],
920        angles: &[f32],
921        stride: usize,
922        off: usize,
923    ) {
924        let hd = self.head_dim;
925        let pairs = if angles.is_empty() {
926            0
927        } else {
928            angles.len() / n
929        };
930        let pool = self.pool.as_deref();
931        let ptr = SendPtr(v.as_mut_ptr());
932        let work = |lo: usize, hi: usize| {
933            for p in lo..hi {
934                for h in 0..heads {
935                    // SAFETY: workers own disjoint token ranges, and the
936                    // heads of one token are disjoint within it.
937                    let x = unsafe { ptr.row(p * stride + off + h * hd, hd) };
938                    let ss = x.iter().map(|&a| (a as f64) * (a as f64)).sum::<f64>() / hd as f64;
939                    let inv = 1.0 / (ss + self.qk_eps).sqrt();
940                    for (d, &g) in x.iter_mut().zip(w) {
941                        *d = (*d as f64 * inv) as f32 * g;
942                    }
943                    for j in 0..pairs {
944                        let a = angles[p * pairs + j];
945                        let (s, c) = a.sin_cos();
946                        let (lo_v, hi_v) = (x[j], x[j + pairs]);
947                        x[j] = lo_v * c - hi_v * s;
948                        x[j + pairs] = lo_v * s + hi_v * c;
949                    }
950                }
951            }
952        };
953        match pool {
954            Some(pl) => pl.run_rows(n, &work),
955            None => work(0, n),
956        }
957    }
958
959    /// Full bidirectional attention over the packed sequence.
960    /// `nr` = (rope angles, q norm weights, k norm weights, eps) when the
961    /// DEVICE should apply qk-norm and RoPE. Passing it means the caller
962    /// skipped `norm_rope_w`, which is the only reason the qkv panel had
963    /// to come back to the host at all.
964    fn attention(
965        &self,
966        qkv: &[f32],
967        n: usize,
968        attn: &mut [f32],
969        nr: Option<(&[f32], &[f32], &[f32], f32)>,
970    ) {
971        let t_repack = std::time::Instant::now();
972        let (nh, hd) = (self.heads, self.head_dim);
973        let inner = nh * hd;
974        let scale = 1.0 / (hd as f32).sqrt();
975        let pool = self.pool.as_deref();
976        // Device path: scores, softmax and the PV product all stay in
977        // device buffers (`dit_qk` → `dit_softmax` → `dit_pv`), so the
978        // n×n score plane — 144 MB at render size — never crosses the
979        // bus. Measured share of a denoise step before this: 41.5%.
980        // The kernels exist and Lumina's DiT already rides them; this
981        // block is what puts MiniMax on the same road. CMF_MMH3_ATTN=cpu
982        // forces the host loop (the A/B that proved the parity).
983        if std::env::var("CMF_MMH3_ATTN").as_deref() != Ok("cpu")
984            && crate::gpu::enabled_here()
985            && n >= 256
986        {
987            // The qkv panel goes up in one piece and is split into
988            // head-major planes ON the card. The host repack below cost
989            // 4.4 s of a 7.3 s attention phase — more than the device
990            // work it fed. Measured at render size: repack 4.4 → 1.5 s,
991            // render 130.0 → 111.8 s, frames bit-identical.
992            // `CMF_MMH3_ATTN=repack` forces the host form.
993            // The device can do qk-norm and RoPE itself (kernel and
994            // plumbing are in; pass Some((angles, q_norm, k_norm, eps))).
995            // `attention` cannot reach them: they live in the caller,
996            // and handing them over means changing this signature — the
997            // last step of task #33, and the point where the panel stops
998            // coming home at all.
999            if std::env::var("CMF_MMH3_ATTN").as_deref() != Ok("repack")
1000                && crate::gpu::dit_attention_packed(qkv, nh, n, hd, scale, nr, attn)
1001            {
1002                // No stamp: the caller's slot-3 stopwatch already spans
1003                // this whole call, and CMF_DIT_ATTN_PROF breaks the
1004                // device half into its three walls. Two overlapping
1005                // stopwatches on the same work is how a 19.2 s step
1006                // came to read as 26.1.
1007                return;
1008            }
1009            assert!(
1010                nr.is_none(),
1011                "device qk-norm was requested but the device path refused; \
1012                 the host loops below expect q/k already normalized"
1013            );
1014            let mut qh = vec![0f32; nh * n * hd];
1015            let mut kh = vec![0f32; nh * n * hd];
1016            let mut vh = vec![0f32; nh * n * hd];
1017            {
1018                let (pq, pk, pv) = (
1019                    SendPtr(qh.as_mut_ptr()),
1020                    SendPtr(kh.as_mut_ptr()),
1021                    SendPtr(vh.as_mut_ptr()),
1022                );
1023                pool_rows(pool, n, &|lo, hi| {
1024                    for p in lo..hi {
1025                        let base = p * 3 * inner;
1026                        for h in 0..nh {
1027                            let dst = (h * n + p) * hd;
1028                            // SAFETY: workers own disjoint token ranges,
1029                            // and each token writes its own head slots.
1030                            unsafe {
1031                                pq.row(dst, hd)
1032                                    .copy_from_slice(&qkv[base + h * hd..base + (h + 1) * hd]);
1033                                pk.row(dst, hd).copy_from_slice(
1034                                    &qkv[base + inner + h * hd..base + inner + (h + 1) * hd],
1035                                );
1036                                pv.row(dst, hd).copy_from_slice(
1037                                    &qkv[base + 2 * inner + h * hd
1038                                        ..base + 2 * inner + (h + 1) * hd],
1039                                );
1040                            }
1041                        }
1042                    }
1043                });
1044            }
1045            // Split the phase: the repack above is host work, this call
1046            // is device work. The last round wrote a tensor-core QK on
1047            // the assumption that the GEMMs dominate and the step did
1048            // not move — so measure the halves before writing anything
1049            // else. Slot 2 (qknorm+rope, unused on this path) takes the
1050            // repack; slot 3 keeps the device call.
1051            Self::prof(2, t_repack);
1052            // No stamp on the device call: the caller's slot-3 stopwatch
1053            // already spans it, and stamping both double-counts.
1054            if crate::gpu::dit_attention(&qh, &kh, &vh, nh, nh, n, hd, scale, attn) {
1055                return;
1056            }
1057        }
1058        let mut qh = vec![0f32; n * hd];
1059        let mut kh = vec![0f32; n * hd];
1060        let mut vt = vec![0f32; hd * n];
1061        let mut scores = vec![0f32; n * n];
1062        let mut oh = vec![0f32; n * hd];
1063        for h in 0..nh {
1064            for p in 0..n {
1065                let base = p * 3 * inner;
1066                let qs = &qkv[base + h * hd..base + (h + 1) * hd];
1067                for (d, &val) in qs.iter().enumerate() {
1068                    qh[p * hd + d] = val * scale;
1069                }
1070                kh[p * hd..(p + 1) * hd]
1071                    .copy_from_slice(&qkv[base + inner + h * hd..base + inner + (h + 1) * hd]);
1072                let vs = &qkv[base + 2 * inner + h * hd..base + 2 * inner + (h + 1) * hd];
1073                for (d, &val) in vs.iter().enumerate() {
1074                    vt[d * n + p] = val;
1075                }
1076            }
1077            crate::fcd_ops::gemm_nt(&qh, &kh, &mut scores, n, hd, n, pool);
1078            let sp = SendPtr(scores.as_mut_ptr());
1079            let soft = |lo: usize, hi: usize| {
1080                for r in lo..hi {
1081                    // SAFETY: workers own disjoint score rows.
1082                    softmax_inplace(unsafe { sp.row(r * n, n) });
1083                }
1084            };
1085            match pool {
1086                Some(pl) => pl.run_rows(n, &soft),
1087                None => soft(0, n),
1088            }
1089            crate::fcd_ops::gemm_nt(&scores, &vt, &mut oh, n, n, hd, pool);
1090            for p in 0..n {
1091                attn[p * inner + h * hd..p * inner + (h + 1) * hd]
1092                    .copy_from_slice(&oh[p * hd..(p + 1) * hd]);
1093            }
1094        }
1095    }
1096
1097    /// Where a denoise step's wall actually goes, in microseconds, under
1098    /// `CMF_MMH3_PROF=1`. Optimizing a 60-second step without this is
1099    /// guesswork, and guesswork on a rented card is expensive.
1100    fn prof(slot: usize, t: std::time::Instant) {
1101        if !mmh3_prof_on() {
1102            return;
1103        }
1104        MMH3_PROF[slot].fetch_add(
1105            t.elapsed().as_micros() as u64,
1106            std::sync::atomic::Ordering::Relaxed,
1107        );
1108    }
1109
1110    /// One block. `mods` is the block's modulation buffer and `rows` the
1111    /// per-token row index into it; both empty for the refiner, which is
1112    /// unmodulated and unrotated.
1113    fn block_forward(
1114        &self,
1115        blk: &Block,
1116        x: &mut [f32],
1117        n: usize,
1118        mods: Option<&[f32]>,
1119        pos: &[[f64; 3]],
1120        rows: &[u32],
1121    ) {
1122        let hs = self.hidden;
1123        let pool = self.pool.as_deref();
1124        let inner = self.heads * self.head_dim;
1125        let angles = if pos.is_empty() {
1126            Vec::new()
1127        } else {
1128            self.rope_angles(pos)
1129        };
1130
1131        let t = std::time::Instant::now();
1132        let mut xn = vec![0f32; n * hs];
1133        self.norm_rows(&mut xn, x, n, hs, &blk.norm1, self.eps);
1134        if let (Some(m), false) = (mods, rows.is_empty()) {
1135            self.modulate(&mut xn, hs, m, rows, 0, 1);
1136        }
1137        Self::prof(0, t);
1138        let t = std::time::Instant::now();
1139        // The projection writes its panel into a device buffer and the
1140        // split reads it THERE — with qk-norm on the device too, nothing
1141        // touches that panel on the host, so 160 MB down and the same
1142        // back up per block simply stop happening. Measured at render
1143        // size: 104.9 → 98.1 s, frames bit-identical.
1144        // `CMF_MMH3_FUSEQKV=0` sends the panel home again.
1145        let t_qkv = std::time::Instant::now();
1146        // Ask whether the device can TAKE the norm, not merely whether a
1147        // device exists. Skipping the host loop for a backend with no
1148        // packed kernel hands the attention unnormalized q/k and there is
1149        // no way back from there — which is why the refusal below used to
1150        // be an assert.
1151        let qk_gpu = std::env::var("CMF_MMH3_QKNORM").as_deref() != Ok("cpu")
1152            && crate::gpu::enabled_here()
1153            && crate::gpu::dit_attention_packed_available()
1154            && n >= 256;
1155        // An adapter on qkv needs the panel the fused paths never let
1156        // out of the card; one on `out` needs the attention output. Each
1157        // stands down only for the projection it owns — a qkv-only
1158        // adapter still gets device attention through the chain below.
1159        if qk_gpu
1160            && !BlockLora::live(&blk.lora.qkv)
1161            && std::env::var("CMF_MMH3_FUSEQKV").as_deref() != Ok("0")
1162        {
1163            // Best case first: qkv, attention and the output projection
1164            // with nothing crossing the bus between them. It refuses at
1165            // the door when anything is missing, so falling through to
1166            // the chain below never repeats work.
1167            let fuse_out = !BlockLora::live(&blk.lora.out)
1168                && std::env::var("CMF_MMH3_FUSEOUT").as_deref() != Ok("0");
1169            if std::env::var("CMF_GPU_DEBUG").is_ok() {
1170                static ONCE: std::sync::Once = std::sync::Once::new();
1171                ONCE.call_once(|| {
1172                    eprintln!(
1173                        "mmh3 fuse-out gate: qkv_q={} out_q={} qkv_map={} out_map={}",
1174                        matches!(&blk.qkv, Proj::Q(_)),
1175                        matches!(&blk.out, Proj::Q(_)),
1176                        matches!(&blk.qkv, Proj::Q(q) if q.mapped_q4tp().is_some()),
1177                        matches!(&blk.out, Proj::Q(o) if o.mapped_q4tp().is_some()),
1178                    )
1179                });
1180            }
1181            if let (true, Proj::Q(q), Proj::Q(o)) = (fuse_out, &blk.qkv, &blk.out) {
1182                if let (Some((m, i)), Some((_, oi))) = (q.mapped_q4tp(), o.mapped_q4tp()) {
1183                    let mut proj = vec![0f32; n * hs];
1184                    if crate::gpu::dit_qkv_attn_out(
1185                        m,
1186                        i,
1187                        oi,
1188                        &xn,
1189                        n,
1190                        hs,
1191                        self.heads,
1192                        self.head_dim,
1193                        1.0 / (self.head_dim as f32).sqrt(),
1194                        (
1195                            &angles[..],
1196                            &blk.q_norm[..],
1197                            &blk.k_norm[..],
1198                            self.qk_eps as f32,
1199                        ),
1200                        &mut proj,
1201                    ) {
1202                        Self::prof(1, t_qkv);
1203                        let t_res = std::time::Instant::now();
1204                        self.residual(x, hs, &proj, mods, rows, 2);
1205                        Self::prof(8, t_res);
1206                        self.ffn_tail(blk, x, n, mods, rows);
1207                        return;
1208                    }
1209                }
1210            }
1211            if let Proj::Q(q) = &blk.qkv {
1212                if let Some((m, i)) = q.mapped_q4tp() {
1213                    let mut attn = vec![0f32; n * inner];
1214                    if crate::gpu::dit_qkv_attention(
1215                        m,
1216                        i,
1217                        &xn,
1218                        n,
1219                        hs,
1220                        self.heads,
1221                        self.head_dim,
1222                        1.0 / (self.head_dim as f32).sqrt(),
1223                        (
1224                            &angles[..],
1225                            &blk.q_norm[..],
1226                            &blk.k_norm[..],
1227                            self.qk_eps as f32,
1228                        ),
1229                        &mut attn,
1230                    ) {
1231                        // Stamp the same slots the unfused path does, or
1232                        // the profile reads 5.5 s for a step the device
1233                        // spends 6+ in: an instrument with a blind spot
1234                        // is worse than none, and this one has cost two
1235                        // wrong conclusions already.
1236                        Self::prof(1, t_qkv);
1237                        let t_out = std::time::Instant::now();
1238                        let mut proj = vec![0f32; n * hs];
1239                        proj_with_lora(&blk.out, &blk.lora.out, &attn, n, &mut proj, pool);
1240                        Self::prof(4, t_out);
1241                        let t_res = std::time::Instant::now();
1242                        self.residual(x, hs, &proj, mods, rows, 2);
1243                        Self::prof(8, t_res);
1244                        self.ffn_tail(blk, x, n, mods, rows);
1245                        return;
1246                    }
1247                }
1248            }
1249        }
1250        let mut qkv = vec![0f32; n * 3 * inner];
1251        proj_with_lora(&blk.qkv, &blk.lora.qkv, &xn, n, &mut qkv, pool);
1252        Self::prof(1, t);
1253        // q and k are the first two thirds of every row; normalize and
1254        // rotate them where they lie, leaving v alone.
1255        // qk-norm and RoPE ride the same device pass that scatters the
1256        // panel head-major, so the host neither walks the panel nor
1257        // needs it back. Parity holds over a whole render: 1, 2 and 4
1258        // steps all match this loop's output exactly (delta 0.000).
1259        //
1260        // A 7.7% "drift" was measured first and was an artefact — the
1261        // reference had been rendered by an EARLIER binary, so the
1262        // comparison carried every change since, not this one. Same
1263        // binary, both arms, or the number means nothing.
1264        // `CMF_MMH3_QKNORM=cpu` restores the host loop.
1265        let qk_on_gpu = qk_gpu;
1266        let t = std::time::Instant::now();
1267        if !qk_on_gpu {
1268            for (which, w) in [(0usize, &blk.q_norm), (1usize, &blk.k_norm)] {
1269                self.norm_rope_w(
1270                    &mut qkv,
1271                    n,
1272                    self.heads,
1273                    w,
1274                    &angles,
1275                    3 * inner,
1276                    which * inner,
1277                );
1278            }
1279        }
1280        Self::prof(2, t);
1281        let t = std::time::Instant::now();
1282        let mut attn = vec![0f32; n * inner];
1283        self.attention(
1284            &qkv,
1285            n,
1286            &mut attn,
1287            qk_on_gpu.then_some((
1288                &angles[..],
1289                &blk.q_norm[..],
1290                &blk.k_norm[..],
1291                self.qk_eps as f32,
1292            )),
1293        );
1294        Self::prof(3, t);
1295        let t = std::time::Instant::now();
1296        let mut proj = vec![0f32; n * hs];
1297        proj_with_lora(&blk.out, &blk.lora.out, &attn, n, &mut proj, pool);
1298        Self::prof(4, t);
1299        // Own slot: the residual is host-side elementwise work with
1300        // modulation, and billing it to the projection hid which of the
1301        // two actually costs (they were 9.8 s together at 512×288).
1302        let t = std::time::Instant::now();
1303        self.residual(x, hs, &proj, mods, rows, 2);
1304        Self::prof(8, t);
1305
1306        self.ffn_tail(blk, x, n, mods, rows);
1307    }
1308
1309    /// The FFN half of a block: norm, modulate, SwiGLU, residual. Its
1310    /// own function so the attention half can return early — the fused
1311    /// path (`dit_qkv_attention`) finishes attention on the device and
1312    /// has nowhere to jump to otherwise.
1313    fn ffn_tail(&self, blk: &Block, x: &mut [f32], n: usize, mods: Option<&[f32]>, rows: &[u32]) {
1314        let hs = self.hidden;
1315        let pool = self.pool.as_deref();
1316        let mut xn = vec![0f32; n * hs];
1317        let mut proj = vec![0f32; n * hs];
1318        let t = std::time::Instant::now();
1319        self.norm_rows(&mut xn, x, n, hs, &blk.norm2, self.eps);
1320        if let (Some(m), false) = (mods, rows.is_empty()) {
1321            self.modulate(&mut xn, hs, m, rows, 3, 4);
1322        }
1323        Self::prof(0, t);
1324        let t = std::time::Instant::now();
1325        // Device-resident FFN: fc1 → SwiGLU → fc2 without the
1326        // intermediate crossing the bus. At render size that panel is
1327        // hundreds of megabytes each way, per block, per step.
1328        // CMF_MMH3_FFN=cpu forces the host chain below.
1329        if std::env::var("CMF_MMH3_FFN").as_deref() != Ok("cpu")
1330            && !blk.lora.ffn_any()
1331            && crate::gpu::enabled_here()
1332            && n >= 64
1333        {
1334            if let (Proj::Q(q1), Proj::Q(q2)) = (&blk.fc1, &blk.fc2) {
1335                if let (Some((m, i1)), Some((_, i2))) = (q1.mapped_q4tp(), q2.mapped_q4tp()) {
1336                    let mut fout = vec![0f32; n * hs];
1337                    if crate::gpu::q4tp_ffn_packed(m, i1, i2, &xn, n, hs, self.ffn, None, &mut fout)
1338                    {
1339                        Self::prof(5, t);
1340                        self.residual(x, hs, &fout, mods, rows, 5);
1341                        return;
1342                    }
1343                }
1344            }
1345        }
1346        let mut gu = vec![0f32; n * 2 * self.ffn];
1347        proj_with_lora(&blk.fc1, &blk.lora.fc1, &xn, n, &mut gu, pool);
1348        Self::prof(5, t);
1349        let t = std::time::Instant::now();
1350        // SwiGLU: fc1's output is [gate | up] per row.
1351        let ffn = self.ffn;
1352        let mut act = vec![0f32; n * ffn];
1353        let ap = SendPtr(act.as_mut_ptr());
1354        pool_rows(pool, n, &|lo, hi| {
1355            for p in lo..hi {
1356                let row = &gu[p * 2 * ffn..(p + 1) * 2 * ffn];
1357                let (g, up) = row.split_at(ffn);
1358                // SAFETY: workers own disjoint token ranges.
1359                for (o, (&a, &b)) in unsafe { ap.row(p * ffn, ffn) }
1360                    .iter_mut()
1361                    .zip(g.iter().zip(up))
1362                {
1363                    *o = silu(a) * b;
1364                }
1365            }
1366        });
1367        Self::prof(6, t);
1368        let t = std::time::Instant::now();
1369        proj_with_lora(&blk.fc2, &blk.lora.fc2, &act, n, &mut proj, pool);
1370        Self::prof(7, t);
1371        self.residual(x, hs, &proj, mods, rows, 5);
1372    }
1373
1374    /// RMSNorm every row of `src` into `dst`, across the pool. One
1375    /// block does this four times over `n·hidden`; on a 1 879-token
1376    /// pack that is 40 M elements a block, and it was running on one
1377    /// thread while forty-seven sat idle.
1378    fn norm_rows(&self, dst: &mut [f32], src: &[f32], n: usize, hs: usize, w: &[f32], eps: f64) {
1379        let ptr = SendPtr(dst.as_mut_ptr());
1380        pool_rows(self.pool.as_deref(), n, &|lo, hi| {
1381            for p in lo..hi {
1382                // SAFETY: workers own disjoint token ranges.
1383                rms_norm_into(&src[p * hs..(p + 1) * hs], w, eps, unsafe {
1384                    ptr.row(p * hs, hs)
1385                });
1386            }
1387        });
1388    }
1389
1390    /// `x = x·(1 + scale[row]) + shift[row]`, per token, across the pool.
1391    fn modulate(
1392        &self,
1393        x: &mut [f32],
1394        hs: usize,
1395        mods: &[f32],
1396        rows: &[u32],
1397        shift_e: usize,
1398        scale_e: usize,
1399    ) {
1400        let stride = EXPAND * hs;
1401        let ptr = SendPtr(x.as_mut_ptr());
1402        pool_rows(self.pool.as_deref(), rows.len(), &|lo, hi| {
1403            for p in lo..hi {
1404                let base = rows[p] as usize * stride;
1405                let shift = &mods[base + shift_e * hs..base + (shift_e + 1) * hs];
1406                let scale = &mods[base + scale_e * hs..base + (scale_e + 1) * hs];
1407                // SAFETY: workers own disjoint token ranges.
1408                for ((v, &sc), &sh) in unsafe { ptr.row(p * hs, hs) }
1409                    .iter_mut()
1410                    .zip(scale)
1411                    .zip(shift)
1412                {
1413                    *v = *v * (1.0 + sc) + sh;
1414                }
1415            }
1416        });
1417    }
1418
1419    /// The gated residual — or a plain one where there is no modulation
1420    /// (the token refiner).
1421    fn residual(
1422        &self,
1423        x: &mut [f32],
1424        hs: usize,
1425        other: &[f32],
1426        mods: Option<&[f32]>,
1427        rows: &[u32],
1428        gate_e: usize,
1429    ) {
1430        let n = x.len() / hs;
1431        let stride = EXPAND * hs;
1432        let ptr = SendPtr(x.as_mut_ptr());
1433        let gated = mods.filter(|_| !rows.is_empty());
1434        pool_rows(self.pool.as_deref(), n, &|lo, hi| {
1435            for p in lo..hi {
1436                // SAFETY: workers own disjoint token ranges.
1437                let row = unsafe { ptr.row(p * hs, hs) };
1438                let src = &other[p * hs..(p + 1) * hs];
1439                match gated {
1440                    Some(m) => {
1441                        let base = rows[p] as usize * stride;
1442                        let gate = &m[base + gate_e * hs..base + (gate_e + 1) * hs];
1443                        for ((v, &g), &o) in row.iter_mut().zip(gate).zip(src) {
1444                            *v += g * o;
1445                        }
1446                    }
1447                    None => {
1448                        for (v, &o) in row.iter_mut().zip(src) {
1449                            *v += o;
1450                        }
1451                    }
1452                }
1453            }
1454        });
1455    }
1456
1457    /// One denoise evaluation.
1458    ///
1459    /// `video` is `[latents_dim, latent_t, lat_h, lat_w]` and `audio` is
1460    /// `[audio_dim, 2, audio_t]`, both in the reference's channel-major
1461    /// order. `text` is the refined `[n, hidden]` stream. Returns the
1462    /// two velocities, EACH ON ITS OWN SCHEDULE and unscaled.
1463    pub fn forward(
1464        &self,
1465        layout: &Layout,
1466        text: &[f32],
1467        video: &[f32],
1468        audio: &[f32],
1469        sigma_v: f64,
1470        cond: &[Vec<f32>],
1471    ) -> (Vec<f32>, Vec<f32>) {
1472        let hs = self.hidden;
1473        let pool = self.pool.as_deref();
1474        let sigma_v = sigma_v.max(1e-6);
1475        let t_v = 1.0 - sigma_v;
1476        let t_a = 1.0 - time_shift_sigma(sigma_v, self.shift_video, self.shift_audio);
1477
1478        // Distinct timesteps, sorted — the adaLN row index is a position
1479        // in this list, so the order is part of the contract. A keyframe
1480        // pins its rows near 1: they are conditions, not noise being
1481        // removed.
1482        let has_cond = layout
1483            .segments
1484            .iter()
1485            .any(|s| matches!(s.kind, Kind::Cond | Kind::RefImg));
1486        let has_ref_audio = layout.segments.iter().any(|s| s.kind == Kind::RefAudio);
1487        // A reference soundtrack pins to the AUDIO clock's condition
1488        // timestep, which is its own number.
1489        let t_cond_a = t_a.max(self.cond_aug_audio);
1490        // The condition rows' timestep IS the noise-augmentation figure:
1491        // the reference blends `aug` of the latent with `1 − aug` of
1492        // noise and then tells the block the row sits at `aug`. Turning
1493        // the blend off means aug = 1, and the timestep moves with it.
1494        let t_cond = t_v.max(self.cond_aug);
1495        let mut ts = vec![t_v, t_a];
1496        if has_cond {
1497            ts.push(t_cond);
1498        }
1499        if has_ref_audio {
1500            ts.push(t_cond_a);
1501        }
1502        ts.sort_by(|a, b| a.partial_cmp(b).unwrap());
1503        ts.dedup();
1504        let row_of = |t: f64| ts.iter().position(|&x| x == t).unwrap();
1505        let (row_v, row_a) = (row_of(t_v), row_of(t_a));
1506        let row_c = if has_cond { row_of(t_cond) } else { 0 };
1507        let row_ca = if has_ref_audio { row_of(t_cond_a) } else { 0 };
1508
1509        // Per-token modulation row: t_row · MODALITIES + tag. The text
1510        // span is not uniform once a vision block is in it — those
1511        // positions carry the VIDEO tag.
1512        let mut rows = vec![0u32; layout.seq_len];
1513        for s in &layout.segments {
1514            match s.kind {
1515                Kind::Text => {
1516                    for (i, v) in rows[s.start..s.stop].iter_mut().enumerate() {
1517                        let tag = *layout.text_tags.get(i).unwrap_or(&(TAG_TEXT as u8));
1518                        *v = (row_v * MODALITIES + tag as usize) as u32;
1519                    }
1520                }
1521                // A reference's rows are conditions too: same timestep
1522                // near 1, and the modality of whichever stream they
1523                // belong to.
1524                Kind::Cond | Kind::RefImg => {
1525                    for v in rows[s.start..s.stop].iter_mut() {
1526                        *v = (row_c * MODALITIES + TAG_VIDEO) as u32;
1527                    }
1528                }
1529                Kind::RefAudio => {
1530                    for v in rows[s.start..s.stop].iter_mut() {
1531                        *v = (row_ca * MODALITIES + TAG_AUDIO) as u32;
1532                    }
1533                }
1534                Kind::Video => {
1535                    for v in rows[s.start..s.stop].iter_mut() {
1536                        *v = (row_v * MODALITIES + TAG_VIDEO) as u32;
1537                    }
1538                }
1539                Kind::Audio => {
1540                    for v in rows[s.start..s.stop].iter_mut() {
1541                        *v = (row_a * MODALITIES + TAG_AUDIO) as u32;
1542                    }
1543                }
1544            }
1545        }
1546
1547        // ── embed ──
1548        let v_rows = patchify_video(
1549            video,
1550            self.latents_dim,
1551            layout.latent_t,
1552            layout.lat_h,
1553            layout.lat_w,
1554        );
1555        let v_n = v_rows.len() / (self.latents_dim * 4);
1556        let a_rows = pack_audio(audio, self.audio_dim, layout.audio_t);
1557        // Condition rows go through the SAME patch projection as the
1558        // target, so they are patchified the same way — one frame each.
1559        let vd = self.latents_dim * 4;
1560        let cond_rows: Vec<Vec<f32>> = cond
1561            .iter()
1562            .enumerate()
1563            .map(|(i, z)| {
1564                let mut r = patchify_video(z, self.latents_dim, 1, layout.lat_h, layout.lat_w);
1565                if self.cond_aug < 1.0 {
1566                    // The reference draws this from a torch generator
1567                    // reseeded per condition; ours is its own stream, so
1568                    // the 0.1% it contributes differs — deliberately, and
1569                    // it is 0.1% of a unit normal.
1570                    let noise = crate::videogen::gauss_pub(r.len(), 0x5EED ^ i as u64);
1571                    let a = self.cond_aug as f32;
1572                    for (v, n) in r.iter_mut().zip(&noise) {
1573                        *v = a * *v + (1.0 - a) * n;
1574                    }
1575                }
1576                r
1577            })
1578            .collect();
1579
1580        let mut h = vec![0f32; layout.seq_len * hs];
1581        let mut ci = 0usize;
1582        for s in layout.segments.iter().filter(|s| s.kind == Kind::Cond) {
1583            let n = s.stop - s.start;
1584            let r = cond_rows.get(ci).unwrap_or_else(|| {
1585                panic!(
1586                    "layout has {} cond segments, {} latents given",
1587                    ci + 1,
1588                    cond_rows.len()
1589                )
1590            });
1591            self.video_patch
1592                .matmat(r, n, &mut h[s.start * hs..s.stop * hs], pool);
1593            for row in h[s.start * hs..s.stop * hs].chunks_exact_mut(hs) {
1594                for (v, &b) in row.iter_mut().zip(&self.video_patch_b) {
1595                    *v += b;
1596                }
1597            }
1598            ci += 1;
1599        }
1600        let vseg = layout.segment(Kind::Video);
1601        let aseg = layout.segment(Kind::Audio);
1602        let tseg = layout.segment(Kind::Text);
1603        h[tseg.start * hs..tseg.stop * hs].copy_from_slice(&text[..(tseg.stop - tseg.start) * hs]);
1604        self.video_patch
1605            .matmat(&v_rows, v_n, &mut h[vseg.start * hs..vseg.stop * hs], pool);
1606        self.audio_patch.matmat(
1607            &a_rows,
1608            aseg.stop - aseg.start,
1609            &mut h[aseg.start * hs..aseg.stop * hs],
1610            pool,
1611        );
1612        for row in h[vseg.start * hs..vseg.stop * hs].chunks_exact_mut(hs) {
1613            for (v, &b) in row.iter_mut().zip(&self.video_patch_b) {
1614                *v += b;
1615            }
1616        }
1617        for row in h[aseg.start * hs..aseg.stop * hs].chunks_exact_mut(hs) {
1618            for (v, &b) in row.iter_mut().zip(&self.audio_patch_b) {
1619                *v += b;
1620            }
1621        }
1622
1623        // ── blocks ──
1624        for (i, blk) in self.blocks.iter().enumerate() {
1625            let mods = blk.adaln.as_ref().unwrap().eval(&ts, pool);
1626            self.block_forward(blk, &mut h, layout.seq_len, Some(&mods), &layout.pos, &rows);
1627            // Per-block watch on the row that dies: growing magnitude is
1628            // an overflow, a sudden jump from finite to NaN is a kernel.
1629            if let Ok(w) = std::env::var("CMF_MMH3_WATCHROW") {
1630                if let Ok(r) = w.parse::<usize>() {
1631                    let row = &h[r * hs..(r + 1) * hs];
1632                    let amax = row
1633                        .iter()
1634                        .filter(|v| v.is_finite())
1635                        .fold(0f32, |a, v| a.max(v.abs()));
1636                    let bad = row.iter().filter(|v| !v.is_finite()).count();
1637                    if bad > 0 || i == 0 || amax > 1e3 {
1638                        eprintln!("  watch blk {i}: row {r} absmax {amax:.3e} nonfinite {bad}");
1639                    }
1640                }
1641            }
1642            if std::env::var_os("CMF_DIT_PROGRESS").is_some() {
1643                eprint!("\r  block {}/{}", i + 1, self.blocks.len());
1644            }
1645        }
1646
1647        // ── heads ──
1648        if std::env::var("CMF_MMH3_NANPROBE").is_ok() {
1649            let bad: Vec<usize> = (0..h.len() / hs)
1650                .filter(|r| h[r * hs..(r + 1) * hs].iter().any(|v| !v.is_finite()))
1651                .collect();
1652            if !bad.is_empty() {
1653                eprintln!(
1654                    "  nanprobe rows: {} of {} bad, first 8 {:?}, video seg {}..{}, audio seg {}..{}",
1655                    bad.len(),
1656                    h.len() / hs,
1657                    &bad[..bad.len().min(8)],
1658                    vseg.start,
1659                    vseg.stop,
1660                    aseg.start,
1661                    aseg.stop
1662                );
1663            }
1664        }
1665        let fm = self.final_adaln.eval(&ts, pool);
1666        let mut video_out = vec![0f32; (vseg.stop - vseg.start) * vd];
1667        let mut audio_out = vec![0f32; (aseg.stop - aseg.start) * self.audio_dim];
1668        for (seg, row, w, b, dst, dim) in [
1669            (
1670                vseg,
1671                row_v,
1672                &self.video_out,
1673                &self.video_out_b,
1674                &mut video_out,
1675                vd,
1676            ),
1677            (
1678                aseg,
1679                row_a,
1680                &self.audio_out,
1681                &self.audio_out_b,
1682                &mut audio_out,
1683                self.audio_dim,
1684            ),
1685        ] {
1686            let n = seg.stop - seg.start;
1687            let mut hn = vec![0f32; n * hs];
1688            for (o, src) in hn
1689                .chunks_exact_mut(hs)
1690                .zip(h[seg.start * hs..seg.stop * hs].chunks_exact(hs))
1691            {
1692                rms_norm_into(src, &self.final_norm, self.final_eps, o);
1693            }
1694            // The final layer's adaLN has one modality, so the row IS
1695            // the timestep index.
1696            let shift = &fm[row * 2 * hs..row * 2 * hs + hs];
1697            let scale = &fm[row * 2 * hs + hs..(row + 1) * 2 * hs];
1698            for r in hn.chunks_exact_mut(hs) {
1699                for ((v, &sc), &sh) in r.iter_mut().zip(scale).zip(shift) {
1700                    *v = *v * (1.0 + sc) + sh;
1701                }
1702            }
1703            // `CMF_MMH3_NANPROBE=1`: which side of the head's GEMM a NaN
1704            // is on. The stream that dies is the one to instrument, and
1705            // "the input was already bad" and "this GEMM made it bad" are
1706            // different bugs in different files.
1707            let probe = std::env::var("CMF_MMH3_NANPROBE").is_ok();
1708            if probe {
1709                let bad_in = hn.iter().filter(|v| !v.is_finite()).count();
1710                eprintln!(
1711                    "  nanprobe dim={dim} n={n} in_bad={bad_in} in_absmax={:.4}",
1712                    hn.iter()
1713                        .filter(|v| v.is_finite())
1714                        .fold(0f32, |a, v| a.max(v.abs()))
1715                );
1716            }
1717            w.matmat(&hn, n, dst, pool);
1718            if probe {
1719                let bad_out = dst.iter().filter(|v| !v.is_finite()).count();
1720                eprintln!(
1721                    "  nanprobe dim={dim} out_bad={bad_out}/{} out_absmax={:.4}",
1722                    dst.len(),
1723                    dst.iter()
1724                        .filter(|v| v.is_finite())
1725                        .fold(0f32, |a, v| a.max(v.abs()))
1726                );
1727            }
1728            for r in dst.chunks_exact_mut(dim) {
1729                for (v, &bv) in r.iter_mut().zip(b.iter()) {
1730                    *v += bv;
1731                }
1732            }
1733        }
1734
1735        // The reference predicts toward the data and the sampler steps
1736        // σ down, hence the sign; the audio velocity is returned on its
1737        // OWN clock rather than pre-scaled by d(σ_a)/d(σ_v).
1738        let video = unpatchify_video(
1739            &video_out,
1740            self.latents_dim,
1741            layout.latent_t,
1742            layout.lat_h,
1743            layout.lat_w,
1744        );
1745        let audio = unpack_audio(&audio_out, self.audio_dim, layout.audio_t);
1746        (
1747            video.iter().map(|&v| -v).collect(),
1748            audio.iter().map(|&v| -v).collect(),
1749        )
1750    }
1751}
1752
1753// ── modulation helpers ──────────────────────────────────────────────
1754
1755/// Rows of `n` items split across pool workers (serial without a pool).
1756fn pool_rows(pool: Option<&Pool>, n: usize, f: &(dyn Fn(usize, usize) + Sync)) {
1757    match pool {
1758        Some(p) => p.run_rows(n, f),
1759        None => f(0, n),
1760    }
1761}
1762
1763// ── stream (un)packing ──────────────────────────────────────────────
1764
1765/// `[C, T, H, W]` → `[T·(H/2)·(W/2), C·4]`, the 2×2 spatial patch
1766/// flattened channel-major-outer as `einsum("nctrhpwq->nthwcrpq")`.
1767pub fn patchify_video(x: &[f32], c: usize, t: usize, h: usize, w: usize) -> Vec<f32> {
1768    let (ph, pw) = (h / 2, w / 2);
1769    let mut out = vec![0f32; t * ph * pw * c * 4];
1770    let mut i = 0;
1771    for ti in 0..t {
1772        for hi in 0..ph {
1773            for wi in 0..pw {
1774                for ci in 0..c {
1775                    for p in 0..2 {
1776                        for q in 0..2 {
1777                            out[i] = x[((ci * t + ti) * h + hi * 2 + p) * w + wi * 2 + q];
1778                            i += 1;
1779                        }
1780                    }
1781                }
1782            }
1783        }
1784    }
1785    out
1786}
1787
1788/// The inverse of `patchify_video`.
1789pub fn unpatchify_video(rows: &[f32], c: usize, t: usize, h: usize, w: usize) -> Vec<f32> {
1790    let (ph, pw) = (h / 2, w / 2);
1791    let mut out = vec![0f32; c * t * h * w];
1792    let mut i = 0;
1793    for ti in 0..t {
1794        for hi in 0..ph {
1795            for wi in 0..pw {
1796                for ci in 0..c {
1797                    for p in 0..2 {
1798                        for q in 0..2 {
1799                            out[((ci * t + ti) * h + hi * 2 + p) * w + wi * 2 + q] = rows[i];
1800                            i += 1;
1801                        }
1802                    }
1803                }
1804            }
1805        }
1806    }
1807    out
1808}
1809
1810/// `[C, 2, T]` → `[2·T, C]`, channel-major: channel 0's frames then
1811/// channel 1's.
1812pub fn pack_audio(x: &[f32], c: usize, t: usize) -> Vec<f32> {
1813    let mut out = vec![0f32; 2 * t * c];
1814    for ch in 0..2 {
1815        for ti in 0..t {
1816            for ci in 0..c {
1817                out[(ch * t + ti) * c + ci] = x[(ci * 2 + ch) * t + ti];
1818            }
1819        }
1820    }
1821    out
1822}
1823
1824/// The inverse of `pack_audio`.
1825pub fn unpack_audio(rows: &[f32], c: usize, t: usize) -> Vec<f32> {
1826    let mut out = vec![0f32; c * 2 * t];
1827    for ch in 0..2 {
1828        for ti in 0..t {
1829            for ci in 0..c {
1830                out[(ci * 2 + ch) * t + ti] = rows[(ch * t + ti) * c + ci];
1831            }
1832        }
1833    }
1834    out
1835}
1836
1837// ── small shared bits ───────────────────────────────────────────────
1838
1839struct SendPtr(*mut f32);
1840unsafe impl Send for SendPtr {}
1841unsafe impl Sync for SendPtr {}
1842impl SendPtr {
1843    /// SAFETY: caller guarantees disjoint `[off, off+len)` per worker.
1844    #[allow(clippy::mut_from_ref)]
1845    unsafe fn row(&self, off: usize, len: usize) -> &mut [f32] {
1846        unsafe { std::slice::from_raw_parts_mut(self.0.add(off), len) }
1847    }
1848}
1849
1850fn softmax_inplace(row: &mut [f32]) {
1851    let mx = row.iter().cloned().fold(f32::MIN, f32::max);
1852    let mut den = 0f32;
1853    for r in row.iter_mut() {
1854        *r = (*r - mx).exp();
1855        den += *r;
1856    }
1857    if den > 0.0 {
1858        let inv = 1.0 / den;
1859        for r in row.iter_mut() {
1860            *r *= inv;
1861        }
1862    }
1863}
1864
1865#[cfg(test)]
1866mod tests {
1867    use super::*;
1868
1869    #[test]
1870    fn schedule_remap_is_an_involution() {
1871        for &s in &[1e-3, 0.25, 0.5, 0.8, 0.972_973, 1.0] {
1872            let a = time_shift_sigma(s, 12.0, 3.0);
1873            let back = time_shift_sigma(a, 3.0, 12.0);
1874            assert!((back - s).abs() < 1e-9, "{s} -> {a} -> {back}");
1875        }
1876    }
1877
1878    #[test]
1879    fn slope_matches_a_finite_difference() {
1880        let h = 1e-6;
1881        for &s in &[0.2, 0.5, 0.9] {
1882            let num = (time_shift_sigma(s + h, 12.0, 3.0) - time_shift_sigma(s - h, 12.0, 3.0))
1883                / (2.0 * h);
1884            let got = time_shift_slope(s, 12.0, 3.0);
1885            assert!((num - got).abs() < 1e-5, "{s}: {num} vs {got}");
1886        }
1887    }
1888
1889    #[test]
1890    fn patchify_round_trips() {
1891        let (c, t, h, w) = (3usize, 2usize, 4usize, 6usize);
1892        let x: Vec<f32> = (0..c * t * h * w).map(|i| i as f32).collect();
1893        let rows = patchify_video(&x, c, t, h, w);
1894        assert_eq!(rows.len(), t * (h / 2) * (w / 2) * c * 4);
1895        assert_eq!(unpatchify_video(&rows, c, t, h, w), x);
1896    }
1897
1898    #[test]
1899    fn audio_pack_round_trips() {
1900        let (c, t) = (5usize, 7usize);
1901        let x: Vec<f32> = (0..c * 2 * t).map(|i| i as f32).collect();
1902        let rows = pack_audio(&x, c, t);
1903        assert_eq!(unpack_audio(&rows, c, t), x);
1904    }
1905
1906    #[test]
1907    fn keyframes_sit_between_the_text_and_the_audio() {
1908        let (tl, lt, lh, lw, at) = (8usize, 3usize, 8usize, 12usize, 5usize);
1909        let base = Layout::t2va(tl, lt, lh, lw, at);
1910        // First and last frame of a 39-frame clip.
1911        let l = Layout::fl2va(tl, lt, lh, lw, at, &[(0, 39), (38, 39)], &[]);
1912        assert_eq!(l.cond_rows(), 2 * l.frame_rows);
1913        assert_eq!(l.seq_len, base.seq_len + 2 * l.frame_rows);
1914        let kinds: Vec<_> = l.segments.iter().map(|s| s.kind).collect();
1915        assert_eq!(
1916            kinds,
1917            vec![Kind::Text, Kind::Cond, Kind::Cond, Kind::Audio, Kind::Video]
1918        );
1919        // A keyframe never advances the cursor: audio and video start
1920        // where they would have without it.
1921        let (a0, v0) = (l.segment(Kind::Audio), l.segment(Kind::Video));
1922        let (ba, bv) = (base.segment(Kind::Audio), base.segment(Kind::Video));
1923        assert_eq!(l.pos[a0.start][0], base.pos[ba.start][0]);
1924        assert_eq!(l.pos[v0.start][0], base.pos[bv.start][0]);
1925        // The first frame's rows sit at the text's end; the last one's a
1926        // whole clip further on, minus one span.
1927        let c: Vec<_> = l.segments.iter().filter(|s| s.kind == Kind::Cond).collect();
1928        assert_eq!(l.pos[c[0].start][0], tl as f64);
1929        let spans: f64 = (0..lt)
1930            .map(|k| FRAME_RESCALE * FRAME_PER_TOKEN[k % 5])
1931            .sum();
1932        assert!((l.pos[c[1].start][0] - (tl as f64 + spans - FRAME_RESCALE)).abs() < 1e-12);
1933        // Both share the TARGET spatial grid.
1934        assert_eq!(l.pos[c[0].start][1], l.pos[v0.start][1]);
1935        assert_eq!(l.pos[c[0].start][2], l.pos[v0.start][2]);
1936    }
1937
1938    #[test]
1939    fn layout_places_the_target_streams_last() {
1940        let l = Layout::t2va(8, 3, 8, 12, 5);
1941        assert_eq!(l.frame_rows, 4 * 6);
1942        assert_eq!(l.seq_len, 8 + 2 * 5 + 3 * 24);
1943        assert_eq!(l.segments[0].kind, Kind::Text);
1944        assert_eq!(l.segments[1].kind, Kind::Audio);
1945        assert_eq!(l.segments[2].kind, Kind::Video);
1946        // The video t axis advances by the 1,4,4,4,4 span pattern.
1947        let v = l.segment(Kind::Video);
1948        let t0 = l.pos[v.start][0];
1949        let t1 = l.pos[v.start + l.frame_rows][0];
1950        assert!((t1 - t0 - FRAME_RESCALE).abs() < 1e-12);
1951    }
1952}