From 75a263af0e34a092e9a1e0fdef9cf2079948986e Mon Sep 17 00:00:00 2001 From: James Brink Date: Sun, 2 Aug 2026 18:40:16 -0700 Subject: [PATCH 1/3] wip(ltx2): latent tiling primitives for stage-2 refinement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure arithmetic only: interval splitting, the separable trapezoidal blend window, and tile enumeration. No tensor or device dependency, so the part that is easy to get subtly wrong is exhaustively testable on CPU. Fixtures are values captured from running upstream's own `split_by_count` in its virtualenv, so this is pinned to the reference rather than to a reading of it. Includes a guard upstream lacks: `split_by_count` silently breaks the partition of unity once the stride drops below the overlap, because three trapezoids then stack and `blend` has no normalization pass. Measured sums reach 1.16, which would brighten the seam rather than fail. Solving `tile_size >= 2 * overlap` gives `overlap <= dim / (n + 1)`; two-tile layouts cannot triple-overlap so upstream's shipped 2/8-6-6 is untouched. Not yet wired into `render_real_two_stage_av` — the module is dead code until then, so this is not PR-ready. --- crates/mold-inference/src/ltx2/mod.rs | 1 + crates/mold-inference/src/ltx2/tiling.rs | 518 +++++++++++++++++++++++ 2 files changed, 519 insertions(+) create mode 100644 crates/mold-inference/src/ltx2/tiling.rs diff --git a/crates/mold-inference/src/ltx2/mod.rs b/crates/mold-inference/src/ltx2/mod.rs index 2e9d5c8e4..5cafd34c5 100644 --- a/crates/mold-inference/src/ltx2/mod.rs +++ b/crates/mold-inference/src/ltx2/mod.rs @@ -16,6 +16,7 @@ mod runtime; mod sampler; pub(crate) mod single_file; mod text; +mod tiling; pub use chain::{extract_tail_latents, tail_latent_frame_count}; pub use pipeline::Ltx2Engine; diff --git a/crates/mold-inference/src/ltx2/tiling.rs b/crates/mold-inference/src/ltx2/tiling.rs new file mode 100644 index 000000000..971289faa --- /dev/null +++ b/crates/mold-inference/src/ltx2/tiling.rs @@ -0,0 +1,518 @@ +//! Latent-space tiling for LTX-2 stage-2 refinement. +//! +//! Upstream reaches resolutions the transformer never saw in training by +//! running stage 2 over overlapping latent tiles, each denoised at a shape the +//! model handles well, then blending them back with a separable trapezoidal +//! window. See `packages/ltx-core/src/ltx_core/tiling.py` and +//! `modality_tiling.py` in , driven by +//! `packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py` — the only +//! single-GPU reference for the technique. +//! +//! This module is the pure arithmetic: interval splitting, the blend window, +//! and tile enumeration. It has no tensor or device dependency so it can be +//! tested exhaustively on CPU, which matters because a silent desync between +//! the blend weights and the token order produces plausible-looking output +//! rather than an error. + +use anyhow::{Result, bail}; + +/// Half-open `[start, end)` slice of one latent axis, plus the ramp widths +/// that fade it into its neighbours. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DimensionInterval { + pub start: usize, + pub end: usize, + pub left_ramp: usize, + pub right_ramp: usize, +} + +impl DimensionInterval { + pub fn len(self) -> usize { + self.end - self.start + } +} + +/// Per-axis tile count and overlap, in **latent** cells. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DimensionTiling { + pub num_tiles: usize, + pub overlap: usize, +} + +impl DimensionTiling { + pub const fn new(num_tiles: usize, overlap: usize) -> Self { + Self { num_tiles, overlap } + } + + /// A single un-ramped tile — the axis is not split. + pub const fn none() -> Self { + Self { + num_tiles: 1, + overlap: 0, + } + } +} + +/// Tile layout across the three latent axes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct TileCountConfig { + pub frames: DimensionTiling, + pub height: DimensionTiling, + pub width: DimensionTiling, +} + +impl TileCountConfig { + /// Upstream's shipped stage-2 layout: 2 tiles per axis, overlaps 8/6/6 + /// latent cells (`hdr_ic_lora.py:89-93`). + pub const fn upstream_stage2() -> Self { + Self { + frames: DimensionTiling::new(2, 8), + height: DimensionTiling::new(2, 6), + width: DimensionTiling::new(2, 6), + } + } + + pub const fn untiled() -> Self { + Self { + frames: DimensionTiling::none(), + height: DimensionTiling::none(), + width: DimensionTiling::none(), + } + } + + pub fn is_untiled(self) -> bool { + self.frames.num_tiles <= 1 && self.height.num_tiles <= 1 && self.width.num_tiles <= 1 + } +} + +/// One tile: which latent cells it covers per axis, and its blend window. +#[derive(Debug, Clone)] +pub(crate) struct Tile { + pub frames: DimensionInterval, + pub height: DimensionInterval, + pub width: DimensionInterval, + /// Separable per-axis trapezoids, in tile-local order `(f, h, w)`. + pub masks: [Vec; 3], +} + +impl Tile { + /// Flattened `(f, h, w)` row-major blend window — the product of the three + /// per-axis trapezoids (`tiling.py:332-357`). Length matches this tile's + /// token count, in the same order `patchify` emits. + pub fn blend_window(&self) -> Vec { + let [mf, mh, mw] = &self.masks; + let mut out = Vec::with_capacity(mf.len() * mh.len() * mw.len()); + for f in mf { + for h in mh { + for w in mw { + out.push(f * h * w); + } + } + } + out + } + + /// Indices of this tile's tokens within the full latent's flattened + /// `(f, h, w)` token sequence. Strides come from the FULL latent, not the + /// tile (`modality_tiling.py:207-211`). + pub fn token_indices(&self, full_height: usize, full_width: usize) -> Vec { + let mut out = Vec::with_capacity(self.token_count()); + for f in self.frames.start..self.frames.end { + for h in self.height.start..self.height.end { + for w in self.width.start..self.width.end { + out.push(f * full_height * full_width + h * full_width + w); + } + } + } + out + } + + pub fn token_count(&self) -> usize { + self.frames.len() * self.height.len() * self.width.len() + } + + /// Pixel dimensions this tile is denoised at. The inverse of + /// `VideoLatentShape::from_pixel_shape`: the causal VAE's first latent + /// frame maps to a single pixel frame, so this is `(f - 1) * 8 + 1` and + /// emphatically not `f * 8` (`hdr_ic_lora.py:548-550`). + pub fn pixel_shape(&self) -> (usize, usize, usize) { + ( + self.width.len() * LATENT_SPATIAL_STRIDE, + self.height.len() * LATENT_SPATIAL_STRIDE, + (self.frames.len() - 1) * LATENT_TEMPORAL_STRIDE + 1, + ) + } +} + +const LATENT_SPATIAL_STRIDE: usize = 32; +const LATENT_TEMPORAL_STRIDE: usize = 8; + +/// Trapezoidal window for one axis (`tiling.py:11-47`). +/// +/// The ramp values are `i / (ramp + 1)` for `i = 1..=ramp` — deliberately +/// never reaching 0 or 1, which is exactly what makes two adjacent ramps sum +/// to one across an overlap. +pub(crate) fn trapezoidal_mask(length: usize, left_ramp: usize, right_ramp: usize) -> Vec { + let left = left_ramp.min(length); + let right = right_ramp.min(length); + let mut mask = vec![1.0f32; length]; + for (i, slot) in mask.iter_mut().take(left).enumerate() { + *slot *= (i + 1) as f32 / (left + 1) as f32; + } + if right > 0 { + let base = length - right; + for (j, slot) in mask.iter_mut().skip(base).enumerate() { + // Upstream multiplies, so a left and right ramp that overlap + // compound rather than clamp. + *slot *= 1.0 - (j + 1) as f32 / (right + 1) as f32; + } + } + for value in &mut mask { + *value = value.clamp(0.0, 1.0); + } + mask +} + +/// Split an axis into `size`-long intervals overlapping by `overlap` +/// (`tiling.py:133-170`). +fn split_by_size(dimension_size: usize, size: usize, overlap: usize) -> Result> { + if size == 0 { + bail!("tile size must be > 0"); + } + if overlap >= size { + bail!("overlap must satisfy 0 <= overlap < size, got overlap={overlap}, size={size}"); + } + if dimension_size <= size { + return Ok(vec![DimensionInterval { + start: 0, + end: dimension_size, + left_ramp: 0, + right_ramp: 0, + }]); + } + let stride = size - overlap; + let amount = (dimension_size + size - 2 * overlap - 1) / stride; + let mut out = Vec::with_capacity(amount); + out.push(DimensionInterval { + start: 0, + end: size, + left_ramp: 0, + right_ramp: overlap, + }); + for i in 1..amount.saturating_sub(1) { + out.push(DimensionInterval { + start: i * stride, + end: i * stride + size, + left_ramp: overlap, + right_ramp: overlap, + }); + } + if amount > 1 { + out.push(DimensionInterval { + start: (amount - 1) * stride, + end: dimension_size, + left_ramp: overlap, + right_ramp: 0, + }); + } + Ok(out) +} + +/// Split an axis into exactly `num_tiles` intervals (`tiling.py:246-291`). +pub(crate) fn split_by_count( + dimension_size: usize, + num_tiles: usize, + overlap: usize, +) -> Result> { + if num_tiles == 0 { + bail!("num_tiles must be >= 1"); + } + if num_tiles > dimension_size { + bail!("cannot split {dimension_size} cells into {num_tiles} tiles"); + } + if num_tiles == 1 { + return Ok(vec![DimensionInterval { + start: 0, + end: dimension_size, + left_ramp: 0, + right_ramp: 0, + }]); + } + let total = dimension_size + overlap * (num_tiles - 1); + let tile_size = total / num_tiles; + let remainder = total % num_tiles; + let base = split_by_size(dimension_size - remainder, tile_size, overlap)?; + Ok(base + .into_iter() + .enumerate() + .map(|(i, interval)| { + let shift = i.min(remainder); + let grow = usize::from(i < remainder); + DimensionInterval { + start: interval.start + shift, + end: interval.end + shift + grow, + ..interval + } + }) + .collect()) +} + +/// Shrink a requested layout to something this latent can actually be split +/// into (`hdr_ic_lora.py:100-142`), and additionally refuse a layout whose +/// windows would not sum to one. +pub(crate) fn clamp_dimension_tiling(cfg: DimensionTiling, dimension_size: usize) -> DimensionTiling { + if cfg.num_tiles <= 1 { + return DimensionTiling::none(); + } + if dimension_size < cfg.num_tiles { + return DimensionTiling::none(); + } + // `split_by_count` needs `overlap < tile_size`, which reduces to this. + let max_overlap = dimension_size - cfg.num_tiles; + let mut overlap = cfg.overlap.min(max_overlap); + + // Upstream stops here, which is unsound for more than two tiles: once the + // stride drops below the overlap, tiles i and i+2 also overlap, three + // trapezoids stack, and the windows sum to as much as 1.16 — a silently + // brightened seam, because `blend` has no normalization pass. + // + // No triple overlap requires `tile_size >= 2 * overlap`. Substituting + // `tile_size = (dim + overlap * (n - 1)) / n` and solving gives + // `overlap <= dim / (n + 1)`. Two tiles can never triple-overlap at all, + // so upstream's shipped 2/8-6-6 layout is untouched. + if cfg.num_tiles > 2 { + overlap = overlap.min(dimension_size / (cfg.num_tiles + 1)); + } + DimensionTiling { + num_tiles: cfg.num_tiles, + overlap, + } +} + +/// Enumerate tiles in `(frames, height, width)` row-major order. +/// +/// The order is load-bearing: it defines each tile's index, which upstream +/// uses to seed that tile's noise (`hdr_ic_lora.py:547`). +pub(crate) fn create_tiles( + latent_frames: usize, + latent_height: usize, + latent_width: usize, + cfg: TileCountConfig, +) -> Result> { + let frames = clamp_dimension_tiling(cfg.frames, latent_frames); + let height = clamp_dimension_tiling(cfg.height, latent_height); + let width = clamp_dimension_tiling(cfg.width, latent_width); + + let f_intervals = split_by_count(latent_frames, frames.num_tiles, frames.overlap)?; + let h_intervals = split_by_count(latent_height, height.num_tiles, height.overlap)?; + let w_intervals = split_by_count(latent_width, width.num_tiles, width.overlap)?; + + let mut tiles = Vec::with_capacity(f_intervals.len() * h_intervals.len() * w_intervals.len()); + for f in &f_intervals { + for h in &h_intervals { + for w in &w_intervals { + tiles.push(Tile { + frames: *f, + height: *h, + width: *w, + masks: [ + trapezoidal_mask(f.len(), f.left_ramp, f.right_ramp), + trapezoidal_mask(h.len(), h.left_ramp, h.right_ramp), + trapezoidal_mask(w.len(), w.left_ramp, w.right_ramp), + ], + }); + } + } + } + Ok(tiles) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Expected values were produced by running upstream's own + /// `split_by_count` in its virtualenv, so this pins the port to the + /// reference rather than to my reading of it. + #[test] + fn split_by_count_matches_upstream() { + let cases: &[(usize, usize, usize, &[(usize, usize, usize, usize)])] = &[ + (21, 2, 8, &[(0, 15, 0, 8), (7, 21, 8, 0)]), + (34, 2, 6, &[(0, 20, 0, 6), (14, 34, 6, 0)]), + (60, 2, 6, &[(0, 33, 0, 6), (27, 60, 6, 0)]), + (10, 2, 8, &[(0, 9, 0, 8), (1, 10, 8, 0)]), + ( + 20, + 4, + 2, + &[(0, 7, 0, 2), (5, 12, 2, 2), (10, 16, 2, 2), (14, 20, 2, 0)], + ), + ]; + for &(dim, n, overlap, expected) in cases { + let got: Vec<_> = split_by_count(dim, n, overlap) + .unwrap() + .into_iter() + .map(|i| (i.start, i.end, i.left_ramp, i.right_ramp)) + .collect(); + assert_eq!(got, expected, "split_by_count({dim}, {n}, {overlap})"); + } + } + + #[test] + fn split_by_count_rejects_impossible_layouts() { + // tile_size == overlap == 8 + assert!(split_by_count(9, 2, 8).is_err()); + assert!(split_by_count(4, 5, 0).is_err()); + assert_eq!(split_by_count(12, 1, 0).unwrap().len(), 1); + } + + #[test] + fn trapezoid_ramps_never_reach_zero_or_one() { + let mask = trapezoidal_mask(15, 0, 8); + assert_eq!(&mask[..7], &[1.0; 7]); + for (i, value) in mask[7..].iter().enumerate() { + let expected = 1.0 - (i + 1) as f32 / 9.0; + assert!((value - expected).abs() < 1e-6, "slot {i}: {value}"); + } + assert!(mask.iter().all(|v| *v > 0.0 && *v <= 1.0)); + } + + /// The whole technique rests on this: `blend` has no normalization pass, + /// so if the windows do not sum to one the seam is silently wrong. + fn assert_partition_of_unity(dim: usize, n: usize, overlap: usize) { + let intervals = split_by_count(dim, n, overlap).unwrap(); + let mut acc = vec![0.0f32; dim]; + for interval in &intervals { + let mask = trapezoidal_mask(interval.len(), interval.left_ramp, interval.right_ramp); + for (offset, weight) in mask.iter().enumerate() { + acc[interval.start + offset] += weight; + } + } + for (i, total) in acc.iter().enumerate() { + assert!( + (total - 1.0).abs() < 1e-6, + "cell {i} of split({dim}, {n}, {overlap}) sums to {total}, not 1" + ); + } + } + + #[test] + fn blend_windows_sum_to_one_across_tiles() { + for &(dim, n, overlap) in &[ + (21, 2, 8), + (34, 2, 6), + (60, 2, 6), + (21, 3, 4), + (20, 4, 2), + (9, 2, 7), + (12, 2, 0), + ] { + assert_partition_of_unity(dim, n, overlap); + } + } + + /// Upstream clamps only enough to avoid a `ValueError`, so a three-tile + /// layout with a wide overlap triple-stacks and sums to ~1.07. Our clamp + /// narrows the overlap instead of shipping a brightened seam. + #[test] + fn wide_overlaps_are_narrowed_rather_than_triple_stacked() { + for &(dim, n, requested) in &[(21, 3, 8), (21, 4, 8), (34, 3, 12)] { + let clamped = clamp_dimension_tiling(DimensionTiling::new(n, requested), dim); + assert!(clamped.overlap <= requested); + assert_partition_of_unity(dim, clamped.num_tiles, clamped.overlap); + } + // Two tiles can never triple-overlap, so the shipped layout is + // untouched. + assert_eq!( + clamp_dimension_tiling(DimensionTiling::new(2, 8), 21), + DimensionTiling::new(2, 8) + ); + } + + #[test] + fn clamping_degrades_to_a_single_tile_when_the_axis_is_too_small() { + assert_eq!( + clamp_dimension_tiling(DimensionTiling::new(4, 2), 3), + DimensionTiling::none() + ); + assert_eq!( + clamp_dimension_tiling(DimensionTiling::new(2, 99), 10), + DimensionTiling::new(2, 8) + ); + } + + /// Tile order defines each tile's noise seed upstream, so it is part of + /// the contract, not an implementation detail. + /// Fixture uses 2 tiles / overlap 2 per axis, matching the layout the + /// upstream numbers below were captured with. + fn two_by_two() -> TileCountConfig { + TileCountConfig { + frames: DimensionTiling::new(2, 2), + height: DimensionTiling::new(2, 2), + width: DimensionTiling::new(2, 2), + } + } + + #[test] + fn tiles_enumerate_frames_then_height_then_width() { + let tiles = create_tiles(9, 8, 10, two_by_two()).unwrap(); + let coords: Vec<_> = tiles + .iter() + .map(|t| { + ( + (t.frames.start, t.frames.end), + (t.height.start, t.height.end), + (t.width.start, t.width.end), + ) + }) + .collect(); + assert_eq!(coords.len(), 8); + assert_eq!(coords[0], ((0, 6), (0, 5), (0, 6))); + assert_eq!(coords[1], ((0, 6), (0, 5), (4, 10))); + assert_eq!(coords[2], ((0, 6), (3, 8), (0, 6))); + assert_eq!(coords[3], ((0, 6), (3, 8), (4, 10))); + assert_eq!(coords[4], ((4, 9), (0, 5), (0, 6))); + } + + /// A tile's window must line up cell-for-cell with the tokens it covers, + /// in the same `(f, h, w)` order `patchify` emits. A transpose here would + /// desync the two silently. + #[test] + fn blend_windows_reconstruct_the_full_field() { + let (frames, height, width) = (9usize, 8usize, 10usize); + let tiles = create_tiles(frames, height, width, two_by_two()).unwrap(); + let mut acc = vec![0.0f32; frames * height * width]; + for tile in &tiles { + let window = tile.blend_window(); + let indices = tile.token_indices(height, width); + assert_eq!(window.len(), indices.len()); + assert_eq!(window.len(), tile.token_count()); + for (token, weight) in indices.iter().zip(window) { + acc[*token] += weight; + } + } + for (token, total) in acc.iter().enumerate() { + assert!( + (total - 1.0).abs() < 1e-5, + "token {token} accumulated {total}, not 1" + ); + } + } + + #[test] + fn tile_pixel_shape_inverts_the_latent_grid() { + let tiles = create_tiles(9, 8, 10, two_by_two()).unwrap(); + // First tile is 6 latent frames x 5 rows x 6 cols. + let (w, h, f) = tiles[0].pixel_shape(); + // 6 latent frames -> (6 - 1) * 8 + 1 = 41 pixel frames, not 48. + assert_eq!((w, h, f), (6 * 32, 5 * 32, 41)); + } + + #[test] + fn an_untiled_layout_yields_one_full_cover_tile() { + let tiles = create_tiles(9, 8, 10, TileCountConfig::untiled()).unwrap(); + assert_eq!(tiles.len(), 1); + assert_eq!(tiles[0].token_count(), 9 * 8 * 10); + assert!(tiles[0].blend_window().iter().all(|w| *w == 1.0)); + } +} From a38c1095d5775a90b8ee009ded521ca82f68c5a1 Mon Sep 17 00:00:00 2001 From: James Brink Date: Sun, 2 Aug 2026 19:47:23 -0700 Subject: [PATCH 2/3] feat(ltx2): plan a stage-2 tile layout from the trained RoPE span MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tile arithmetic needs a policy: when to tile at all, and into how many tiles. Both answers come from the same fact that motivates tiling — the checkpoints' RoPE was trained over a 2048px span, which is 64 latent cells after the VAE's /32 compression. A shape already inside that span is left untiled, because tiling costs a full denoise pass per tile and must not be paid for nothing. That covers every resolution mold ships today, 1920x1088 included. An oversized axis is split into the smallest number of tiles that brings every tile back inside the span, using upstream's shipped overlaps. Time stays whole: the duration budget already caps frames, so the temporal axis is not what runs past the trained span at high resolution. The 4K test asserts the property that matters — every tile of a 120x68 latent plan fits the trained span, and the blend windows still reconstruct the field to 1e-5. --- crates/mold-inference/src/ltx2/tiling.rs | 100 ++++++++++++++++++++++- 1 file changed, 97 insertions(+), 3 deletions(-) diff --git a/crates/mold-inference/src/ltx2/tiling.rs b/crates/mold-inference/src/ltx2/tiling.rs index 971289faa..d01c0be84 100644 --- a/crates/mold-inference/src/ltx2/tiling.rs +++ b/crates/mold-inference/src/ltx2/tiling.rs @@ -14,7 +14,7 @@ //! the blend weights and the token order produces plausible-looking output //! rather than an error. -use anyhow::{Result, bail}; +use anyhow::{bail, Result}; /// Half-open `[start, end)` slice of one latent axis, plus the ramp widths /// that fade it into its neighbours. @@ -175,7 +175,11 @@ pub(crate) fn trapezoidal_mask(length: usize, left_ramp: usize, right_ramp: usiz /// Split an axis into `size`-long intervals overlapping by `overlap` /// (`tiling.py:133-170`). -fn split_by_size(dimension_size: usize, size: usize, overlap: usize) -> Result> { +fn split_by_size( + dimension_size: usize, + size: usize, + overlap: usize, +) -> Result> { if size == 0 { bail!("tile size must be > 0"); } @@ -260,7 +264,10 @@ pub(crate) fn split_by_count( /// Shrink a requested layout to something this latent can actually be split /// into (`hdr_ic_lora.py:100-142`), and additionally refuse a layout whose /// windows would not sum to one. -pub(crate) fn clamp_dimension_tiling(cfg: DimensionTiling, dimension_size: usize) -> DimensionTiling { +pub(crate) fn clamp_dimension_tiling( + cfg: DimensionTiling, + dimension_size: usize, +) -> DimensionTiling { if cfg.num_tiles <= 1 { return DimensionTiling::none(); } @@ -327,6 +334,47 @@ pub(crate) fn create_tiles( Ok(tiles) } +/// Per-axis span, in latent cells, that the checkpoints' RoPE was trained on. +/// +/// `positional_embedding_max_pos = [20, 2048, 2048]` in pixels, and the video +/// VAE compresses space by 32, so the spatial axes are trained over 64 latent +/// cells. Past that, positions land outside the trained range. +pub(crate) const TRAINED_SPATIAL_LATENT_SPAN: usize = 2_048 / LATENT_SPATIAL_STRIDE; + +/// Choose a tile layout for a stage-2 latent. +/// +/// Returns [`TileCountConfig::untiled`] when the shape is already inside the +/// span the model was trained on — tiling costs a full denoise pass per tile, +/// so it must not be paid unless it buys something. Otherwise split each +/// oversized axis into just enough tiles to bring every tile back inside the +/// span, using upstream's shipped overlaps. +pub(crate) fn plan_stage2_tiling( + latent_frames: usize, + latent_height: usize, + latent_width: usize, +) -> TileCountConfig { + let axis = |size: usize, overlap: usize| -> DimensionTiling { + if size <= TRAINED_SPATIAL_LATENT_SPAN { + return DimensionTiling::none(); + } + // Ceiling division: the smallest tile count whose tiles fit the span. + let tiles = size.div_ceil(TRAINED_SPATIAL_LATENT_SPAN).max(2); + clamp_dimension_tiling(DimensionTiling::new(tiles, overlap), size) + }; + + let upstream = TileCountConfig::upstream_stage2(); + let planned = TileCountConfig { + // Time is not the axis that blows past the trained span at high + // resolution — the duration budget already caps it — so frames stay + // whole and only space is tiled. + frames: DimensionTiling::none(), + height: axis(latent_height, upstream.height.overlap), + width: axis(latent_width, upstream.width.overlap), + }; + let _ = latent_frames; + planned +} + #[cfg(test)] mod tests { use super::*; @@ -508,6 +556,52 @@ mod tests { assert_eq!((w, h, f), (6 * 32, 5 * 32, 41)); } + /// Tiling costs a full denoise pass per tile, so a shape the model was + /// trained on must not pay for it. + #[test] + fn shapes_inside_the_trained_span_are_not_tiled() { + // 1216x704 -> 38x22 latent cells: the LTX-2 default, comfortably inside. + assert!(plan_stage2_tiling(13, 22, 38).is_untiled()); + // 1920x1088 -> 60x34: still inside, which is why #668 could ship it + // without tiling. + assert!(plan_stage2_tiling(13, 34, 60).is_untiled()); + // Exactly at the span. + assert!(plan_stage2_tiling(13, 64, 64).is_untiled()); + } + + /// Every tile of a 4K plan must land back inside the trained span, or + /// tiling has not bought anything. + #[test] + fn oversized_axes_are_split_until_every_tile_fits_the_trained_span() { + // 3840x2160 -> 120x68 latent cells. + let cfg = plan_stage2_tiling(13, 68, 120); + assert!(!cfg.is_untiled()); + assert_eq!(cfg.frames, DimensionTiling::none(), "time stays whole"); + + let tiles = create_tiles(13, 68, 120, cfg).unwrap(); + for tile in &tiles { + assert!( + tile.height.len() <= TRAINED_SPATIAL_LATENT_SPAN, + "tile is {} cells tall, past the trained span", + tile.height.len() + ); + assert!( + tile.width.len() <= TRAINED_SPATIAL_LATENT_SPAN, + "tile is {} cells wide, past the trained span", + tile.width.len() + ); + } + + // And the windows still reconstruct the field exactly. + let mut acc = vec![0.0f32; 13 * 68 * 120]; + for tile in &tiles { + for (token, weight) in tile.token_indices(68, 120).iter().zip(tile.blend_window()) { + acc[*token] += weight; + } + } + assert!(acc.iter().all(|total| (total - 1.0).abs() < 1e-5)); + } + #[test] fn an_untiled_layout_yields_one_full_cover_tile() { let tiles = create_tiles(9, 8, 10, TileCountConfig::untiled()).unwrap(); From 6291a3a587c80dc76e96e8bcf38fccb33aa5d770 Mon Sep 17 00:00:00 2001 From: James Brink Date: Sun, 2 Aug 2026 19:49:45 -0700 Subject: [PATCH 3/3] feat(ltx2): report when stage 2 runs past the trained RoPE span MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A render above the span the checkpoints were trained on still produces a picture, which is exactly why it needs saying out loud — the failure mode is degraded structure, not an error. Stage 2 now reports the shape and the tile layout that would bring it back inside, on the existing LTX-2 VRAM telemetry target. Also documents why the tile execution arithmetic is staged ahead of its caller: it is the part that fails silently when wrong, so it is verified against upstream now rather than written under pressure later. --- crates/mold-inference/src/ltx2/runtime.rs | 27 +++++++++++++++++++++++ crates/mold-inference/src/ltx2/tiling.rs | 17 +++++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/crates/mold-inference/src/ltx2/runtime.rs b/crates/mold-inference/src/ltx2/runtime.rs index a84ca8f05..15e95a1e4 100644 --- a/crates/mold-inference/src/ltx2/runtime.rs +++ b/crates/mold-inference/src/ltx2/runtime.rs @@ -2344,6 +2344,7 @@ fn render_real_distilled_av( fps: plan.frame_rate as f32, }; let stage2_video_latent_shape = video_latent_shape_from_tensor(&stage2_clean_video_latents)?; + report_stage2_trained_span(stage2_video_latent_shape); let stage2_pixel_shape = pixel_shape_for_video_latents(stage2_video_latent_shape, plan.frame_rate); let stage2_video_conditioning = maybe_load_stage_video_conditioning( @@ -2748,6 +2749,7 @@ fn render_real_two_stage_av( fps: plan.frame_rate as f32, }; let stage2_video_latent_shape = video_latent_shape_from_tensor(&stage2_clean_video_latents)?; + report_stage2_trained_span(stage2_video_latent_shape); let stage2_pixel_shape = pixel_shape_for_video_latents(stage2_video_latent_shape, plan.frame_rate); let stage2_video_conditioning = maybe_load_stage_video_conditioning( @@ -5723,6 +5725,31 @@ fn ltx_debug_timings_enabled() -> bool { const LTX2_VRAM_TARGET: &str = "mold::ltx2::vram"; +/// Report whether stage 2 is running a shape past the span the checkpoints' +/// RoPE was trained on, and what tile layout would bring it back inside. +/// +/// A render beyond that span still produces a picture, which is exactly why +/// it needs saying out loud: the failure mode is degraded structure, not an +/// error. Reported once per stage-2 pass so the operator can correlate it +/// with the output. +fn report_stage2_trained_span(shape: VideoLatentShape) { + let plan = crate::ltx2::tiling::plan_stage2_tiling(shape.frames, shape.height, shape.width); + if plan.is_untiled() { + return; + } + tracing::info!( + target: LTX2_VRAM_TARGET, + "[ltx2-vram] stage 2 latent {}x{} exceeds the {}-cell span these checkpoints were \ + trained on; tiled refinement would use {}x{} tiles (not yet enabled — see \ + https://github.com/utensils/mold/issues/673). Expect degraded structure at this size.", + shape.width, + shape.height, + crate::ltx2::tiling::TRAINED_SPATIAL_LATENT_SPAN, + plan.width.num_tiles, + plan.height.num_tiles, + ); +} + thread_local! { /// The residency plan the most recent transformer build on this thread /// used. Stashed rather than threaded through the call graph because the diff --git a/crates/mold-inference/src/ltx2/tiling.rs b/crates/mold-inference/src/ltx2/tiling.rs index d01c0be84..69d2445b9 100644 --- a/crates/mold-inference/src/ltx2/tiling.rs +++ b/crates/mold-inference/src/ltx2/tiling.rs @@ -1,3 +1,13 @@ +// `plan_stage2_tiling` is live — stage-2 telemetry reports when a shape +// exceeds the span the checkpoints were trained on. The tile *execution* +// arithmetic (interval splitting, blend windows, tile enumeration) is verified +// against upstream but has no production caller until the per-tile denoise +// lands in https://github.com/utensils/mold/issues/673. It is staged here +// rather than written later because it is the part that fails silently when +// it is wrong: a desynced blend window produces a plausible-looking seam, not +// an error. +#![allow(dead_code)] + //! Latent-space tiling for LTX-2 stage-2 refinement. //! //! Upstream reaches resolutions the transformer never saw in training by @@ -384,7 +394,12 @@ mod tests { /// reference rather than to my reading of it. #[test] fn split_by_count_matches_upstream() { - let cases: &[(usize, usize, usize, &[(usize, usize, usize, usize)])] = &[ + /// `(start, end, left_ramp, right_ramp)` as upstream reports it. + type ExpectedInterval = (usize, usize, usize, usize); + /// `(dimension, tiles, overlap, expected intervals)`. + type SplitCase = (usize, usize, usize, &'static [ExpectedInterval]); + + let cases: &[SplitCase] = &[ (21, 2, 8, &[(0, 15, 0, 8), (7, 21, 8, 0)]), (34, 2, 6, &[(0, 20, 0, 6), (14, 34, 6, 0)]), (60, 2, 6, &[(0, 33, 0, 6), (27, 60, 6, 0)]),