From fba967ed9c6e8f59f3ca77e8e9ba8afd1eaa3cdd Mon Sep 17 00:00:00 2001 From: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:01:06 -0700 Subject: [PATCH] [diffusion] Support diffusion decoder parallel tiling for LTX-2.5 (#36026) Co-authored-by: Claude Opus 5 Co-authored-by: Mick Qian --- docs/cookbook/diffusion/LTX/LTX2.5.mdx | 61 +++ .../configs/pipeline_configs/base.py | 21 +- .../runtime/entrypoints/openai/video_api.py | 1 + .../decoders/ltx_2_5_diffusion_decoder.py | 269 +++++++++-- .../runtime/pipelines/ltx_2_pipeline.py | 30 +- .../ltx_2/decoding_av.py | 11 +- .../multimodal_gen/test/server/gpu_cases.py | 25 ++ .../test/server/perf_baselines/h100.json | 8 + .../sglang/multimodal_gen/test/test_utils.py | 2 +- .../test_diffusion_decoder_tile_sharding.py | 422 ++++++++++++++++++ 10 files changed, 795 insertions(+), 55 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_diffusion_decoder_tile_sharding.py diff --git a/docs/cookbook/diffusion/LTX/LTX2.5.mdx b/docs/cookbook/diffusion/LTX/LTX2.5.mdx index 94def6262..6a75c25b4 100644 --- a/docs/cookbook/diffusion/LTX/LTX2.5.mdx +++ b/docs/cookbook/diffusion/LTX/LTX2.5.mdx @@ -128,6 +128,7 @@ Selecting weights: | 2 GPUs, long sequences | `--num-gpus 2 --ulysses-degree 2` | Sequence parallel; the memory/long-sequence tool. | | 2 GPUs, large DiT | `--num-gpus 2 --tp-size 2` | Tensor parallel across attention heads. | | 2 GPUs, dev weights | `--num-gpus 2 --enable-cfg-parallel` | Splits the guided and unguided branches across GPUs. Measured 1.77x on denoising (15.1s to 8.5s, 960×544 / 57 frames / 30 steps). | +| 2 GPUs, diffusion decoder | `--num-gpus 2 --ulysses-degree 2` | The decoder's tiles are split across the ranks by default. See [section 4.6.1](#4-6-1-memory-and-multi-gpu). | **CFG parallelism does not apply on the default (distilled) path.** That DiT @@ -340,3 +341,63 @@ uv pip install natten==0.21.6+torch2110cu130 -f https://whl.natten.org/ Nothing else changes if you skip it: the decoder still produces the same video, just slower. + +#### 4.6.1 Memory and multi-GPU + +Two flags govern how the decode is executed. Both default to on, so the numbers +below are what you already get — they matter when you want to turn one off. + +`--diffusion-decoder-tiling` runs the decoder's two expensive stages over +overlapping tiles instead of the whole volume. It is a **memory** control, not a +speed one: it costs wall clock and buys headroom. + +`--diffusion-decoder-parallel-tiling` splits those tiles across the +decode-parallel ranks -- the TP, SP, PP and CFG ranks of one replica, since the +decoder is replicated over all of them. Without it every one of those ranks +decodes every tile and keeps its own identical copy. It only applies on the +tiled path, so it does nothing when tiling is off, and nothing at a single +rank. + +Decoding stage on 2xH200 at 960×544, Ulysses degree 2. "Peak" is the whole +process, not the decode alone: + +| Frames | Tiling | Parallel tiling | Decode | Peak | +| --- | --- | --- | --- | --- | +| 121 | off | *(n/a)* | **3.38s** | 103.2 GB | +| 121 | on | off | 5.68s | 78.2 GB | +| 121 | on | on | 4.08s | 79.3 GB | +| 49 | off | *(n/a)* | **1.93s** | 84.0 GB | +| 49 | on | off | 2.79s | 78.2 GB | +| 49 | on | on | 2.12s | 78.2 GB | + +Reading that: + +- **Untiled is the fastest option** whenever it fits. Tiling exists for the + ~25 GB it saves at 121 frames, which is the difference between fitting on an + 80 GB card and not. +- **Parallel tiling recovers most of tiling's cost** — 1.39x at 121 frames, + 1.32x at 49 — but does not beat an untiled decode. It also adds about 1 GB + for the gather buffers. +- The gain grows with tile count, so it is larger at higher resolution: at + 1920×1088 / 49 frames the decode goes from 12.32s to 7.86s, **1.57x**. + +Output is bitwise identical however many ranks the tiles are split over: every +rank draws the whole grid's noise in the same order, and only the decode is +shared out. Turning *tiling* on or off does change the result slightly near +tile borders, so pick one and stay with it if you need reproducible frames: + +```bash +# Fastest, if the untiled decode fits in VRAM +sglang serve \ + --model-path Lightricks/LTX-2.5-Diffusers \ + --pipeline-class-name LTX2Pipeline \ + --load-diffusion-decoder \ + --diffusion-decoder-tiling false + +# Memory-bound: keep tiling, and split the tiles over both GPUs +sglang serve \ + --model-path Lightricks/LTX-2.5-Diffusers \ + --pipeline-class-name LTX2Pipeline \ + --load-diffusion-decoder \ + --num-gpus 2 --ulysses-degree 2 +``` diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index f6180c0a5..762465ee4 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index b9665e440..2fe259932 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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"), diff --git a/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py b/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py index 1b336f6dd..ac8d86f9f 100644 --- a/python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 988ec73c3..6a750db23 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py index e5f58189f..ab6fad3e0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index e32b4e93e..5da7490cc 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index d7b2b0cd0..f3c62b12d 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 33db17fed..0e0102486 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/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. diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_decoder_tile_sharding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_decoder_tile_sharding.py new file mode 100644 index 000000000..cd2425008 --- /dev/null +++ b/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()