[diffusion] Support diffusion decoder parallel tiling for LTX-2.5 (#36026)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -236,11 +236,15 @@ class PipelineConfig:
     # dtype resident so lossless requests never consume pre-rounded weights.
     vae_decode_precision_high: str | None = None
     vae_tiling: bool = True
+    vae_slicing: bool = False
+    vae_sp: bool = True
+
+    # Diffusion Decoder configuration
     # Bounds the attention grid the diffusion decoder's stages see, which is
     # what makes a full-length decode tractable.
     diffusion_decoder_tiling: bool = True
-    vae_slicing: bool = False
-    vae_sp: bool = True
+    # Splits those tiles across the decode-parallel ranks.
+    diffusion_decoder_parallel_tiling: bool = True
 
     # Image encoder configuration
     image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
@@ -865,6 +869,19 @@ class PipelineConfig:
             default=PipelineConfig.diffusion_decoder_tiling,
             help="Enable tiling for the LTX-2.5 diffusion decoder",
         )
+        parser.add_argument(
+            f"--{prefix_with_dot}diffusion-decoder-parallel-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}diffusion_decoder_parallel_tiling",
+            default=PipelineConfig.diffusion_decoder_parallel_tiling,
+            help=(
+                "Split the LTX-2.5 diffusion decoder's tiles across the "
+                "decode-parallel ranks (TP/SP/PP/CFG within a replica). "
+                "Requires --diffusion-decoder-tiling, since the tiles it "
+                "splits only exist on that path; inert otherwise, and at a "
+                "single rank"
+            ),
+        )
         parser.add_argument(
             f"--{prefix_with_dot}vae-slicing",
             action=StoreBoolean,

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -286,6 +286,7 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         "max_sequence_length": request.max_sequence_length,
         "flow_shift": request.flow_shift,
         "enable_teacache": request.enable_teacache,
+        "use_diffusion_decoder": _extra_value(request, "use_diffusion_decoder"),
         "enable_cache_dit": _extra_value(request, "enable_cache_dit"),
         "cache_dit_params": _extra_value(request, "cache_dit_params"),
         "cfg_gate_step": _extra_value(request, "cfg_gate_step"),

python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py

@@ -403,7 +403,8 @@ class LTX2VideoVaeNeighborhoodAttention(nn.Module):
         key = self.norm_k(key)
         query = query * self.scale
         query, key = self.rope.forward_pair(query, key)
-        return query, key, value
+
+        return query.to(value.dtype), key.to(value.dtype), value
 
     def build_block_mask(self, hidden_states: torch.Tensor):
         """The window mask for this grid, or `None` when NATTEN handles it.
@@ -764,6 +765,30 @@ class LTX2VideoDiffusionDecoder3d(nn.Module):
             ]
         return hidden_states
 
+    def stage_4_output_extent(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        drop_leading_frame: bool = True,
+        crop_trailing_ghost: bool = True,
+    ) -> tuple[int, int, int]:
+        """`forward_stage_4`'s `(T, H, W)` without running it.
+
+        The blocks preserve the grid and the upsample scales it, so the extent
+        is a pure function of the input extent. Tiled decoding needs the size
+        of a tile's noise before deciding whether to spend the stage on it, and
+        `tiled_decode` asserts this against the real context it later builds.
+        """
+        stride_t, stride_h, stride_w = self.upsamples[-1].stride
+        num_frames *= stride_t
+        if stride_t == 2 and drop_leading_frame:
+            num_frames -= 1
+        num_pad = self.trailing_pad_latent_frames
+        if crop_trailing_ghost and num_pad > 0:
+            num_frames -= num_pad * self.temporal_compression_ratio
+        return num_frames, height * stride_h, width * stride_w
+
     def forward_diffusion_step(
         self, latent_context: torch.Tensor, x_t: torch.Tensor, timestep: torch.Tensor
     ) -> torch.Tensor:
@@ -847,6 +872,57 @@ class LTX2VideoDiffusionDecoder3d(nn.Module):
         return self.denoise(latent_context, x_t, num_inference_steps)
 
 
+def _all_gather_tiles(
+    local_tiles: list[torch.Tensor],
+    local_indices: list[int],
+    total: int,
+    group,
+    device: torch.device,
+) -> list[torch.Tensor]:
+    """Collect every rank's decoded tiles, restoring global tile order.
+
+    Tiles at the grid edges differ in shape, so the payloads are flattened and
+    padded to a common length and the shapes travel alongside them. A rank can
+    hold no tile at all -- there may be fewer tiles than ranks -- and it still
+    has to enter the collective with the dtype the others are sending, so the
+    dtype travels with the shapes rather than being assumed.
+    """
+    import torch.distributed as dist
+
+    from sglang.multimodal_gen.runtime.distributed.utils import all_gather_single
+
+    world_size = dist.get_world_size(group=group)
+    local_meta = [(i, tuple(t.shape)) for i, t in zip(local_indices, local_tiles)]
+    meta: list = [None] * world_size
+    dist.all_gather_object(
+        meta,
+        (local_tiles[0].dtype if local_tiles else None, local_meta),
+        group=group,
+    )
+    # Every rank that decoded a tile ran the same decoder over the same dtype,
+    # so any one of them names the dtype for the ranks that decoded none.
+    dtype = next(d for d, _ in meta if d is not None)
+
+    max_size = max(
+        sum(math.prod(shape) for _, shape in per_rank) for _, per_rank in meta
+    )
+    padded = torch.zeros(max_size, device=device, dtype=dtype)
+    if local_tiles:
+        payload = torch.cat([t.reshape(-1) for t in local_tiles])
+        padded[: payload.numel()] = payload
+    gathered = torch.empty(world_size * max_size, device=device, dtype=dtype)
+    all_gather_single(gathered, padded, group=group)
+
+    tiles: list = [None] * total
+    for rank, (_, per_rank) in enumerate(meta):
+        offset = rank * max_size
+        for tile_index, shape in per_rank:
+            count = math.prod(shape)
+            tiles[tile_index] = gathered[offset : offset + count].reshape(shape)
+            offset += count
+    return tiles
+
+
 def _tile_intervals(
     length: int, tile_size: int, stride: int, min_size: int
 ) -> list[tuple[int, int]]:
@@ -898,6 +974,11 @@ class LTX2VideoDiffusionDecoderModel(nn.Module, LayerwiseOffloadableModuleMixin)
         # output only moves near tile borders. Set by the decoding stage from
         # `--diffusion-decoder-tiling`; tile sizes match upstream.
         self.use_tiling = False
+        # Shards the tiles across the decode-parallel ranks. Only meaningful
+        # with tiling on, and only when there is more than one rank; the
+        # decoding stage sets it from `--diffusion-decoder-parallel-tiling`,
+        # and clears it when the stage does not run on every one of them.
+        self.use_parallel_tiling = False
         self.tile_sample_min_height = 768
         self.tile_sample_min_width = 768
         self.tile_sample_min_num_frames = 32
@@ -905,6 +986,33 @@ class LTX2VideoDiffusionDecoderModel(nn.Module, LayerwiseOffloadableModuleMixin)
         self.tile_sample_stride_width = 512
         self.tile_sample_stride_num_frames = 16
 
+    def _tile_shard(self):
+        """`(rank, world_size, group)` to split tiles over, or `(0, 1, None)`.
+
+        The decoder is replicated across the whole decode-parallel group -- TP,
+        SP, PP and CFG within one DP replica -- so those are exactly the ranks
+        that would otherwise each decode every tile.
+        """
+        if not self.use_parallel_tiling:
+            return 0, 1, None
+        from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+            get_decode_parallel_group_coordinator,
+            get_decode_parallel_rank,
+            get_decode_parallel_world_size,
+            model_parallel_is_initialized,
+        )
+
+        if not model_parallel_is_initialized():
+            return 0, 1, None
+        world_size = get_decode_parallel_world_size()
+        if world_size <= 1:
+            return 0, 1, None
+        return (
+            get_decode_parallel_rank(),
+            world_size,
+            get_decode_parallel_group_coordinator().device_group,
+        )
+
     @staticmethod
     def _blend(a: torch.Tensor, b: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
         """Linear cross-fade of `a`'s tail into `b`'s head along `dim`."""
@@ -1009,49 +1117,124 @@ class LTX2VideoDiffusionDecoderModel(nn.Module, LayerwiseOffloadableModuleMixin)
                 dtype=hidden_states.dtype,
             )
 
-        frame_groups = []
-        for t0, t1 in temporal_tiles:
-            is_origin = t0 == 0
+        # Flatten the grid first: the tiles are independent, so which rank runs
+        # which is just a slice of this list.
+        coords = [
+            (t0, t1, h0, h1, w0, w1)
+            for t0, t1 in temporal_tiles
+            for h0, h1 in height_tiles
+            for w0, w1 in width_tiles
+        ]
+        rank, world_size, group = self._tile_shard()
+        logger.debug(
+            "Diffusion decoder tiling: %d tiles (%dx%dx%d), rank %d of %d",
+            len(coords),
+            len(temporal_tiles),
+            len(height_tiles),
+            len(width_tiles),
+            rank,
+            world_size,
+        )
+
+        def _tile_extent(index: int) -> tuple[int, int, int, int, int]:
+            """Tile `index`'s pixel shape, without paying for its context."""
+            t0, t1, h0, h1, w0, w1 = coords[index]
             is_trailing = t1 == num_frames
-            feature_t1 = features.shape[1] if is_trailing else t1
-            rows = []
-            for h0, h1 in height_tiles:
-                row = []
-                for w0, w1 in width_tiles:
-                    context = decoder.forward_stage_4(
-                        features[:, t0:feature_t1, h0:h1, w0:w1],
-                        drop_leading_frame=is_origin,
-                        crop_trailing_ghost=is_trailing,
-                    )
-                    tile_shape = (
-                        batch_size,
-                        decoder.out_channels,
-                        context.shape[1],
-                        context.shape[2] * patch_size,
-                        context.shape[3] * patch_size,
-                    )
-                    if single_step_x0:
-                        x_t = torch.randn(
-                            tile_shape,
-                            generator=generator,
-                            device=hidden_states.device,
-                            dtype=hidden_states.dtype,
-                        )
-                    else:
-                        # A non-origin tile keeps its duplicate leading frame, so
-                        # it starts one pixel frame earlier than t0 * scale_t.
-                        pixel_t0 = t0 * scale_t - (
-                            1 if not is_origin and scale_t == 2 else 0
-                        )
-                        x_t = x_t_full[
-                            :,
-                            :,
-                            pixel_t0 : pixel_t0 + tile_shape[2],
-                            h0 * scale_h : h0 * scale_h + tile_shape[3],
-                            w0 * scale_w : w0 * scale_w + tile_shape[4],
-                        ]
-                    row.append(decoder.denoise(context, x_t, num_inference_steps))
-                rows.append(row)
+            context_t, context_h, context_w = decoder.stage_4_output_extent(
+                (features.shape[1] if is_trailing else t1) - t0,
+                h1 - h0,
+                w1 - w0,
+                drop_leading_frame=t0 == 0,
+                crop_trailing_ghost=is_trailing,
+            )
+            return (
+                batch_size,
+                decoder.out_channels,
+                context_t,
+                context_h * patch_size,
+                context_w * patch_size,
+            )
+
+        def _tile_x_t(index: int) -> torch.Tensor:
+            """Tile `index`'s starting noise.
+
+            Every rank walks the whole grid in order so the generator sees the
+            same sequence of draws it saw when a single rank decoded every
+            tile: sharding the tiles must not move the output. A draw skipped
+            here is far cheaper than the stage it feeds, so only the decode is
+            worth splitting.
+            """
+            tile_shape = _tile_extent(index)
+            if single_step_x0:
+                return torch.randn(
+                    tile_shape,
+                    generator=generator,
+                    device=hidden_states.device,
+                    dtype=hidden_states.dtype,
+                )
+            t0, _, h0, _, w0, _ = coords[index]
+            # A non-origin tile keeps its duplicate leading frame, so it starts
+            # one pixel frame earlier than t0 * scale_t.
+            pixel_t0 = t0 * scale_t - (1 if t0 != 0 and scale_t == 2 else 0)
+            return x_t_full[
+                :,
+                :,
+                pixel_t0 : pixel_t0 + tile_shape[2],
+                h0 * scale_h : h0 * scale_h + tile_shape[3],
+                w0 * scale_w : w0 * scale_w + tile_shape[4],
+            ]
+
+        def _decode_tile(index: int, x_t: torch.Tensor) -> torch.Tensor:
+            t0, t1, h0, h1, w0, w1 = coords[index]
+            is_trailing = t1 == num_frames
+            context = decoder.forward_stage_4(
+                features[
+                    :, t0 : features.shape[1] if is_trailing else t1, h0:h1, w0:w1
+                ],
+                drop_leading_frame=t0 == 0,
+                crop_trailing_ghost=is_trailing,
+            )
+            # `_tile_extent` predicted this to size the noise; a mismatch means
+            # the two have drifted apart and every tile's noise is now wrong.
+            assert tuple(x_t.shape[2:]) == (
+                context.shape[1],
+                context.shape[2] * patch_size,
+                context.shape[3] * patch_size,
+            ), "tile noise shape does not match the context it conditions on"
+            return decoder.denoise(context, x_t, num_inference_steps)
+
+        local_indices = []
+        local_tiles = []
+        for index in range(len(coords)):
+            x_t = _tile_x_t(index)
+            if index % world_size != rank:
+                del x_t
+                continue
+            local_indices.append(index)
+            local_tiles.append(_decode_tile(index, x_t))
+
+        if world_size > 1:
+            flat_tiles = _all_gather_tiles(
+                local_tiles,
+                local_indices,
+                len(coords),
+                group,
+                hidden_states.device,
+            )
+        else:
+            flat_tiles = local_tiles
+
+        # Back into the (t, h, w) nesting the blending below expects.
+        per_row = len(width_tiles)
+        per_group = len(height_tiles) * per_row
+        frame_groups = []
+        for g in range(len(temporal_tiles)):
+            rows = [
+                flat_tiles[
+                    g * per_group + r * per_row : g * per_group + (r + 1) * per_row
+                ]
+                for r in range(len(height_tiles))
+            ]
 
             result_rows = []
             for i, row in enumerate(rows):

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -388,14 +388,30 @@ class _BaseLTX2Pipeline(LoRAPipeline):
 
     @staticmethod
     def _declares_component(model_path: str, component_name: str) -> bool:
+        """Whether `model_index.json` names a real component.
+
+        Falls through to the hub when `model_path` is a repo id rather than a
+        directory. Reading it as a directory alone makes every optional
+        component look absent -- silently for `duration_head`, and as a
+        spurious "checkpoint does not declare it" for the diffusion decoder.
+        """
         index_path = os.path.join(str(model_path), "model_index.json")
-        if not os.path.exists(index_path):
-            return False
-        try:
-            with open(index_path) as f:
-                model_index = json.load(f)
-        except (OSError, ValueError):
-            return False
+        if os.path.exists(index_path):
+            try:
+                with open(index_path) as f:
+                    model_index = json.load(f)
+            except (OSError, ValueError):
+                return False
+        else:
+            from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+                maybe_download_model_index,
+            )
+
+            try:
+                model_index = maybe_download_model_index(str(model_path))
+            except Exception as e:
+                logger.debug("No model_index.json for %s: %s", model_path, e)
+                return False
         entry = model_index.get(component_name)
         # model_index.json records absent optional components as [null, null].
         return bool(entry) and entry[0] is not None

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py

@@ -5,6 +5,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
     ComponentUse,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    StageParallelismType,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -76,9 +79,13 @@ class LTX2AVDecodingStage(DecodingStage):
         It is a diffusion model in its own right, so it needs a generator; the
         request's seed keeps a decode reproducible.
         """
-        # Untiled, every stage attends over the whole volume -- minutes at a
-        # full-length 121-frame grid.
+        # Untiled, every stage attends over the whole volume
         decoder.use_tiling = bool(server_args.pipeline_config.diffusion_decoder_tiling)
+        # Splitting the tiles is a collective over the decode-parallel group
+        decoder.use_parallel_tiling = (
+            bool(server_args.pipeline_config.diffusion_decoder_parallel_tiling)
+            and self.parallelism_type == StageParallelismType.REPLICATED
+        )
         generator = torch.Generator(device=latents.device).manual_seed(int(batch.seed))
         return decoder(latents, generator=generator)
 

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -1049,6 +1049,31 @@ TWO_GPU_CASES = [
         DiffusionSamplingParams(prompt=T2V_PROMPT, extras={"seed": 42}),
         run_component_accuracy_check=False,
     ),
+    # LTX-2.5's diffusion decoder
+    DiffusionTestCase(
+        "ltx_2_5_diffusion_decoder_2gpus",
+        DiffusionServerArgs(
+            model_path="Lightricks/LTX-2.5-Diffusers",
+            modality="video",
+            ulysses_degree=2,
+            # Offload both the DiT and text encoder between stages to leave
+            # decoder headroom on 80 GB GPUs.
+            extras=[
+                "--load-diffusion-decoder",
+                "--component-residency "
+                "transformer=component-offload,text_encoder=component-offload",
+            ],
+        ),
+        DiffusionSamplingParams(
+            prompt=T2V_PROMPT,
+            output_size="768x448",
+            num_frames=49,
+            expect_audio_output=True,
+            extras={"seed": 42, "use_diffusion_decoder": True},
+        ),
+        run_perf_check=False,
+        run_component_accuracy_check=False,
+    ),
     # I2V LoRA test case
     DiffusionTestCase(
         "wan2_1_i2v_14b_lora_2gpu",

python/sglang/multimodal_gen/test/server/perf_baselines/h100.json

@@ -2763,6 +2763,14 @@
             "runtime_peak_allocated_mb": 56951.0,
             "estimated_full_test_time_s": 186.0
         },
+        "ltx_2_5_diffusion_decoder_2gpus": {
+            "stages_ms": {},
+            "denoise_step_ms": {},
+            "expected_e2e_ms": 0.0,
+            "expected_avg_denoise_ms": 0.0,
+            "expected_median_denoise_ms": 0.0,
+            "estimated_full_test_time_s": 545.1
+        },
         "ltx_2_3_two_stage_ti2v_2gpus": {
             "stages_ms": {
                 "InputValidationStage": 3.21,

python/sglang/multimodal_gen/test/test_utils.py

@@ -40,7 +40,7 @@ logger = init_logger(__name__)
 # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows
 # publish.
 SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion"
-SGL_TEST_FILES_CI_DATA_REVISION = "285ffa7fa7a8afcfb90e2344e3f9d6ca12a5e2ea"
+SGL_TEST_FILES_CI_DATA_REVISION = "4ce5eeb9606e378478b2d0964d83e960af4e88cf"
 
 # The NPU pin is kept as a separate branch so ascend GT can be bumped independently
 # when it's regenerated on its own cadence.

python/sglang/multimodal_gen/test/unit/test_diffusion_decoder_tile_sharding.py

@@ -0,0 +1,422 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Splitting the diffusion decoder's tiles across ranks must not move the output.
+
+The tiles of a diffusion decode are independent, so they can be shared out over
+the ranks that would otherwise each decode all of them. That is only a
+performance optimization if the result is bit-for-bit what the single-rank,
+tile-by-tile decode produced -- which means the noise has to keep coming off
+one generator in one global tile order, however few tiles a rank ends up
+decoding.
+
+The distributed cases run over gloo on CPU, so they need no GPU.
+"""
+
+import math
+import os
+import socket
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from sglang.multimodal_gen.runtime.models.decoders.ltx_2_5_diffusion_decoder import (
+    LTX2VideoDiffusionDecoder3d,
+    LTX2VideoDiffusionDecoderModel,
+    LTX2VideoVaePixelShuffleUpsampler,
+    _all_gather_tiles,
+    _tile_intervals,
+)
+
+PARALLEL_STATE = "sglang.multimodal_gen.runtime.distributed.parallel_state"
+
+
+class _Upsampler:
+    """Just the stride; the stub reproduces the shuffle itself."""
+
+    def __init__(self, stride):
+        self.stride = stride
+
+
+class _StubDecoder:
+    """A decoder with the real stage geometry and arithmetic values.
+
+    `tiled_decode` only needs the shapes to line up and the values to depend on
+    both the tile's content and its noise -- that is enough for a misplaced
+    tile or a misdrawn noise to show up as a different output.
+    """
+
+    # Under test: the shape prediction that lets a rank size a tile's noise
+    # without paying for the tile's context.
+    stage_4_output_extent = LTX2VideoDiffusionDecoder3d.stage_4_output_extent
+
+    def __init__(self, num_inference_steps=1):
+        self.upsamples = [_Upsampler((2, 1, 1)), _Upsampler((2, 2, 2))]
+        self.temporal_compression_ratio = math.prod(u.stride[0] for u in self.upsamples)
+        self.trailing_pad_latent_frames = 2
+        self.patch_size = 2
+        self.out_channels = 3
+        self.context_channels = 4
+        self.default_num_inference_steps = num_inference_steps
+        self.model_output_type = "x0"
+        self.drawn_noise = []
+
+    @property
+    def ghost(self) -> int:
+        return self.trailing_pad_latent_frames * math.prod(
+            u.stride[0] for u in self.upsamples[:-1]
+        )
+
+    def forward_stages_1_to_3(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        features = hidden_states.mean(dim=1)
+        trailing = features[:, -1:].repeat(1, self.ghost, 1, 1)
+        features = torch.cat([features, trailing], dim=1)
+        channels = torch.arange(
+            self.context_channels, dtype=features.dtype, device=features.device
+        )
+        return features.unsqueeze(-1) + channels
+
+    def forward_stage_4(
+        self,
+        hidden_states: torch.Tensor,
+        drop_leading_frame: bool = True,
+        crop_trailing_ghost: bool = True,
+    ) -> torch.Tensor:
+        stride_t, stride_h, stride_w = self.upsamples[-1].stride
+        out = (
+            hidden_states.repeat_interleave(stride_t, dim=1)
+            .repeat_interleave(stride_h, dim=2)
+            .repeat_interleave(stride_w, dim=3)
+        )
+        if stride_t == 2 and drop_leading_frame:
+            out = out[:, 1:]
+        num_pad = self.trailing_pad_latent_frames
+        if crop_trailing_ghost and num_pad > 0:
+            out = out[:, : -num_pad * self.temporal_compression_ratio]
+        return out
+
+    def denoise(self, context, x_t, num_inference_steps):
+        self.drawn_noise.append(x_t.clone())
+        pixels = context[..., : self.out_channels].permute(0, 4, 1, 2, 3)
+        pixels = pixels.repeat_interleave(self.patch_size, dim=3).repeat_interleave(
+            self.patch_size, dim=4
+        )
+        return pixels + x_t * float(num_inference_steps)
+
+
+def _build_model(num_inference_steps: int = 1) -> LTX2VideoDiffusionDecoderModel:
+    model = LTX2VideoDiffusionDecoderModel.__new__(LTX2VideoDiffusionDecoderModel)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(
+        arch_config=SimpleNamespace(
+            decoder_stage_kernels=[(3, 3, 3)],
+            decoder_stage5_kernel=(3, 3, 3),
+        )
+    )
+    model.decoder = _StubDecoder(num_inference_steps)
+    model.use_tiling = True
+    model.use_parallel_tiling = False
+    # Divided by (scale_t, scale_h, scale_w) = (2, 4, 4) inside `tiled_decode`,
+    # so: 4-cell tiles on a 2-cell stride, on every axis.
+    model.tile_sample_min_num_frames = 8
+    model.tile_sample_stride_num_frames = 4
+    model.tile_sample_min_height = 16
+    model.tile_sample_stride_height = 8
+    model.tile_sample_min_width = 16
+    model.tile_sample_stride_width = 8
+    return model
+
+
+def _latent(shape, dtype=torch.float32) -> torch.Tensor:
+    numel = math.prod(shape)
+    return torch.arange(numel, dtype=torch.float32).reshape(shape).div_(numel).to(dtype)
+
+
+# A 2 x 2 x 3 = 12-tile grid, and the width axis is ragged: 7 cells over a
+# 2-cell stride leaves a last tile of 3 where the others are 4.
+MANY_TILES = (1, 4, 6, 6, 7)
+# One tile on every axis -- fewer tiles than ranks, so some rank decodes none.
+ONE_TILE = (1, 4, 4, 4, 4)
+
+
+def _decode(model, latent, seed, shard=(0, 1, None)):
+    model._tile_shard = lambda: shard
+    model.decoder.drawn_noise = []
+    generator = torch.Generator().manual_seed(seed)
+    return model.tiled_decode(latent, generator=generator)
+
+
+def _serial_reference(model, latent, seed):
+    """The pre-parallel decode: one generator, consumed tile by tile.
+
+    Deliberately a second implementation of the tile walk rather than a call
+    into the one under test -- it is the thing the optimization must not have
+    changed.
+    """
+    decoder = model.decoder
+    stride_t, stride_h, stride_w = decoder.upsamples[-1].stride
+    patch = decoder.patch_size
+    scale_t, scale_h, scale_w = stride_t, stride_h * patch, stride_w * patch
+    features = decoder.forward_stages_1_to_3(latent)
+    num_frames = features.shape[1] - decoder.ghost
+    axes = [
+        _tile_intervals(num_frames, 8 // scale_t, 4 // scale_t, 3),
+        _tile_intervals(features.shape[2], 16 // scale_h, 8 // scale_h, 3),
+        _tile_intervals(features.shape[3], 16 // scale_w, 8 // scale_w, 3),
+    ]
+    generator = torch.Generator().manual_seed(seed)
+    noise = []
+    for t0, t1 in axes[0]:
+        for h0, h1 in axes[1]:
+            for w0, w1 in axes[2]:
+                is_trailing = t1 == num_frames
+                context = decoder.forward_stage_4(
+                    features[
+                        :,
+                        t0 : features.shape[1] if is_trailing else t1,
+                        h0:h1,
+                        w0:w1,
+                    ],
+                    drop_leading_frame=t0 == 0,
+                    crop_trailing_ghost=is_trailing,
+                )
+                noise.append(
+                    torch.randn(
+                        (
+                            latent.shape[0],
+                            decoder.out_channels,
+                            context.shape[1],
+                            context.shape[2] * patch,
+                            context.shape[3] * patch,
+                        ),
+                        generator=generator,
+                        dtype=latent.dtype,
+                    )
+                )
+    return noise
+
+
+class TestStage4OutputExtent(unittest.TestCase):
+    """The predicted extent has to be what stage 4 actually produces."""
+
+    def _decoder(self, stride):
+        decoder = _StubDecoder()
+        decoder.upsamples = [_Upsampler(stride)]
+        decoder.temporal_compression_ratio = stride[0]
+        return decoder
+
+    def test_matches_the_real_upsampler(self):
+        for stride in [(2, 2, 2), (1, 2, 2), (2, 1, 1)]:
+            for drop, crop in [(True, True), (True, False), (False, True)]:
+                with self.subTest(stride=stride, drop=drop, crop=crop):
+                    decoder = self._decoder(stride)
+                    upsampler = LTX2VideoVaePixelShuffleUpsampler(8, stride)
+                    hidden = torch.zeros(1, 6, 5, 4, 8)
+                    with torch.no_grad():
+                        out = upsampler(hidden, drop_leading_frame=drop)
+                    num_pad = decoder.trailing_pad_latent_frames
+                    if crop and num_pad > 0:
+                        out = out[:, : -num_pad * decoder.temporal_compression_ratio]
+                    self.assertEqual(
+                        decoder.stage_4_output_extent(
+                            6, 5, 4, drop_leading_frame=drop, crop_trailing_ghost=crop
+                        ),
+                        tuple(out.shape[1:4]),
+                    )
+
+
+class TestSerialDecodeIsUnchanged(unittest.TestCase):
+    """A single rank must still see the noise the pre-parallel decode saw."""
+
+    def test_noise_matches_the_pre_parallel_stream(self):
+        for shape in (MANY_TILES, ONE_TILE):
+            with self.subTest(tiles=shape):
+                model = _build_model()
+                latent = _latent(shape)
+                _decode(model, latent, seed=7)
+                expected = _serial_reference(model, latent, seed=7)
+                self.assertEqual(len(model.decoder.drawn_noise), len(expected))
+                for got, want in zip(model.decoder.drawn_noise, expected):
+                    self.assertTrue(torch.equal(got, want))
+
+    def test_the_request_seed_still_changes_the_result(self):
+        model = _build_model()
+        latent = _latent(MANY_TILES)
+        self.assertFalse(
+            torch.equal(_decode(model, latent, seed=7), _decode(model, latent, seed=8))
+        )
+
+
+class TestTileShard(unittest.TestCase):
+    """Which ranks the tiles are split over."""
+
+    def test_off_when_parallel_tiling_is_disabled(self):
+        model = _build_model()
+        model.use_parallel_tiling = False
+        self.assertEqual(model._tile_shard(), (0, 1, None))
+
+    def test_off_before_the_parallel_state_exists(self):
+        model = _build_model()
+        model.use_parallel_tiling = True
+        with patch(
+            f"{PARALLEL_STATE}.model_parallel_is_initialized", return_value=False
+        ):
+            self.assertEqual(model._tile_shard(), (0, 1, None))
+
+    def test_uses_the_decode_parallel_group_not_the_sp_group(self):
+        # The decoder is replicated over TP/SP/PP/CFG, so sharding only over SP
+        # would leave the TP and CFG ranks redecoding the whole volume.
+        model = _build_model()
+        model.use_parallel_tiling = True
+        group = object()
+        coordinator = SimpleNamespace(device_group=group)
+        with patch(
+            f"{PARALLEL_STATE}.model_parallel_is_initialized", return_value=True
+        ):
+            with patch(
+                f"{PARALLEL_STATE}.get_decode_parallel_world_size", return_value=4
+            ):
+                with patch(
+                    f"{PARALLEL_STATE}.get_decode_parallel_rank", return_value=2
+                ):
+                    with patch(
+                        f"{PARALLEL_STATE}.get_decode_parallel_group_coordinator",
+                        return_value=coordinator,
+                    ):
+                        self.assertEqual(model._tile_shard(), (2, 4, group))
+
+    def test_off_at_a_single_decode_rank(self):
+        model = _build_model()
+        model.use_parallel_tiling = True
+        with patch(
+            f"{PARALLEL_STATE}.model_parallel_is_initialized", return_value=True
+        ):
+            with patch(
+                f"{PARALLEL_STATE}.get_decode_parallel_world_size", return_value=1
+            ):
+                self.assertEqual(model._tile_shard(), (0, 1, None))
+
+
+def _free_port() -> int:
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _check_gather(group, rank, world_size, dtype, tile_shapes):
+    """Every rank must rebuild the same tile list, whatever it held itself."""
+    tiles = [
+        torch.full(shape, float(i + 1), dtype=dtype)
+        for i, shape in enumerate(tile_shapes)
+    ]
+    local_indices = list(range(rank, len(tiles), world_size))
+    gathered = _all_gather_tiles(
+        [tiles[i] for i in local_indices],
+        local_indices,
+        len(tiles),
+        group,
+        torch.device("cpu"),
+    )
+    assert len(gathered) == len(tiles)
+    for want, got in zip(tiles, gathered):
+        assert got is not None, "a tile went missing from the gather"
+        assert got.dtype == dtype, f"{got.dtype} != {dtype}"
+        assert torch.equal(got, want)
+
+
+def _check_sharded_decode(group, rank, world_size, dtype, latent_shape, steps):
+    model = _build_model(steps)
+    latent = _latent(latent_shape, dtype)
+    serial = _decode(model, latent, seed=11)
+    serial_noise = list(model.decoder.drawn_noise)
+    sharded = _decode(model, latent, seed=11, shard=(rank, world_size, group))
+    assert torch.equal(serial, sharded), "sharding the tiles moved the output"
+    # ...and this rank did less work to get there, on the tiles it owns.
+    owned = list(range(rank, len(serial_noise), world_size))
+    assert len(model.decoder.drawn_noise) == len(owned)
+    for index, got in zip(owned, model.decoder.drawn_noise):
+        assert torch.equal(got, serial_noise[index])
+
+
+def _distributed_worker(rank, world_size, port):
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+    try:
+        # A group narrower than the world, so an even split gets covered too.
+        pair = dist.new_group([0, 1])
+        groups = [(dist.group.WORLD, rank, world_size)]
+        if rank < 2:
+            groups.append((pair, rank, 2))
+
+        for group, group_rank, group_size in groups:
+            for dtype in (torch.float32, torch.float16, torch.bfloat16):
+                # Ragged tiles, and more tiles than ranks.
+                _check_gather(
+                    group,
+                    group_rank,
+                    group_size,
+                    dtype,
+                    [(2, 3, 4), (2, 3, 4), (2, 3, 2), (1, 3, 4)],
+                )
+                # Fewer tiles than ranks: the tail ranks send nothing, and have
+                # only the metadata to tell them what dtype to send it as.
+                _check_gather(group, group_rank, group_size, dtype, [(2, 3, 4)])
+
+                for latent_shape in (MANY_TILES, ONE_TILE):
+                    for steps in (1, 2):
+                        _check_sharded_decode(
+                            group,
+                            group_rank,
+                            group_size,
+                            dtype,
+                            latent_shape,
+                            steps,
+                        )
+    finally:
+        dist.destroy_process_group()
+
+
+class TestShardedDecode(unittest.TestCase):
+    """The real gather, over a real process group."""
+
+    def test_three_ranks_agree_with_the_serial_decode(self):
+        mp.spawn(
+            _distributed_worker,
+            args=(3, _free_port()),
+            nprocs=3,
+            join=True,
+        )
+
+
+class TestTileIntervals(unittest.TestCase):
+    def test_intervals_cover_the_axis(self):
+        intervals = _tile_intervals(30, 12, 8, 4)
+        self.assertEqual(intervals[0][0], 0)
+        self.assertEqual(intervals[-1][1], 30)
+
+    def test_a_short_remnant_is_merged_into_the_previous_tile(self):
+        # 25 with stride 8 would leave a 1-long trailing tile, which is below
+        # the neighborhood kernel and cannot be decoded on its own.
+        for start, end in _tile_intervals(25, 12, 8, 4):
+            self.assertGreaterEqual(end - start, 4)
+
+    def test_an_axis_shorter_than_one_tile_stays_whole(self):
+        self.assertEqual(_tile_intervals(5, 12, 8, 4), [(0, 5)])
+
+    def test_a_round_robin_split_covers_every_tile_exactly_once(self):
+        # How `tiled_decode` assigns tiles: rank r takes r, r+W, r+2W, ...
+        for total in (1, 7, 14, 15):
+            for world_size in (1, 2, 3, 4):
+                assigned = [
+                    i
+                    for rank in range(world_size)
+                    for i in range(rank, total, world_size)
+                ]
+                self.assertEqual(sorted(assigned), list(range(total)))
+
+
+if __name__ == "__main__":
+    unittest.main()