1use crate::dit::Proj;
38use crate::ltxlora::{LoraBank, LoraBranch};
39
40pub 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
60pub 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
99const FRAME_PER_TOKEN: [f64; 5] = [1.0, 4.0, 4.0, 4.0, 4.0];
102const FRAME_RESCALE: f64 = 5.0 / 3.0;
103
104const TAG_VIDEO: usize = 0;
106const TAG_TEXT: usize = 1;
107const TAG_AUDIO: usize = 2;
108const MODALITIES: usize = 3;
109const EXPAND: usize = 6;
111const VISUAL_COND_TIMESTEP: f64 = 0.999;
113const AUDIO_COND_TIMESTEP: f64 = 1.0;
115
116pub 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
124pub 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#[derive(Clone, Copy, PartialEq, Debug)]
134pub enum Kind {
135 Text,
136 Cond,
138 RefImg,
140 RefAudio,
142 Audio,
143 Video,
144}
145
146#[derive(Clone, Copy, Debug)]
148pub struct Segment {
149 pub start: usize,
150 pub stop: usize,
151 pub kind: Kind,
152}
153
154pub struct Layout {
157 pub seq_len: usize,
158 pub segments: Vec<Segment>,
159 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 pub frame_rows: usize,
168 pub text_tags: Vec<u8>,
171}
172
173#[derive(Clone, Debug)]
175pub enum Ref {
176 Image { lat_h: usize, lat_w: usize },
178 Audio { t: usize },
180 Video {
183 latent_t: usize,
184 lat_h: usize,
185 lat_w: usize,
186 audio_t: usize,
187 },
188}
189
190fn 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
202fn 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 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 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 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 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 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 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 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 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 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
493struct Adaln {
496 w: Proj,
498 b: Vec<f32>,
499 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 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 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, out: Proj, q_norm: Vec<f32>,
559 k_norm: Vec<f32>,
560 fc1: Proj, fc2: Proj, adaln: Option<Adaln>,
563 lora: BlockLora,
568}
569
570#[derive(Default)]
575struct BlockLora {
576 qkv: Option<LoraBranch>,
577 out: Option<LoraBranch>,
578 fc1: Option<LoraBranch>,
579 fc2: Option<LoraBranch>,
580}
581
582fn 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 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 pub cond_aug: f64,
670 pub cond_aug_audio: f64,
673 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 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 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 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 pub fn lora_bound(&self) -> usize {
808 self.lora_bound.len()
809 }
810
811 pub fn lora_binds(&self, key: &str) -> bool {
813 self.lora_bound.contains(key)
814 }
815
816 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 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 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 #[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 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 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 if std::env::var("CMF_MMH3_ATTN").as_deref() != Ok("cpu")
984 && crate::gpu::enabled_here()
985 && n >= 256
986 {
987 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 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 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 Self::prof(2, t_repack);
1052 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 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 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 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 let t_qkv = std::time::Instant::now();
1146 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 if qk_gpu
1160 && !BlockLora::live(&blk.lora.qkv)
1161 && std::env::var("CMF_MMH3_FUSEQKV").as_deref() != Ok("0")
1162 {
1163 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 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 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 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 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 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 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 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 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 rms_norm_into(&src[p * hs..(p + 1) * hs], w, eps, unsafe {
1384 ptr.row(p * hs, hs)
1385 });
1386 }
1387 });
1388 }
1389
1390 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 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 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 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 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 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 let t_cond_a = t_a.max(self.cond_aug_audio);
1490 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 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 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 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 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 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 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 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 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 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 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 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
1753fn 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
1763pub 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
1788pub 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
1810pub 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
1824pub 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
1837struct SendPtr(*mut f32);
1840unsafe impl Send for SendPtr {}
1841unsafe impl Sync for SendPtr {}
1842impl SendPtr {
1843 #[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 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 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 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 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 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}