diff --git a/xfuser/model_executor/layers/ltx2/__init__.py b/xfuser/model_executor/layers/ltx2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xfuser/model_executor/layers/ltx2/diffusion_decoder.py b/xfuser/model_executor/layers/ltx2/diffusion_decoder.py new file mode 100644 index 00000000..bf8ab442 --- /dev/null +++ b/xfuser/model_executor/layers/ltx2/diffusion_decoder.py @@ -0,0 +1,268 @@ +"""Tile-parallel diffusion VAE decode for LTX-2.5 across 2/4/8 GPUs. + +The stock ``LTX2VideoDiffusionDecoderModel.tiled_decode`` runs a triple-nested loop over +``(temporal × height × width)`` tiles, each independently calling ``forward_stage_4`` ++ ``denoise`` (8 neighbourhood-attention blocks each). + +``xFuserLTX2VideoDiffusionDecoderWrapper`` distributes those tiles across the SP group +via round-robin ownership, then does one small all_reduce for shape metadata followed by +per-tile broadcasts to reassemble the full video on every rank. Communication is a +single collective per tile. + +Tile geometry for the default 1024×1536×121 config: 12 tiles, adequate for 2/4/8 GPUs. +""" +from __future__ import annotations + +import math + +import torch +import torch.distributed + +from diffusers import LTX2VideoDiffusionDecoderModel +from diffusers.models.autoencoders.ltx2_diffusion_decoder import _tile_intervals +from diffusers.utils.torch_utils import randn_tensor + +from xfuser.core.distributed import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, +) + + +class xFuserLTX2VideoDiffusionDecoderWrapper(LTX2VideoDiffusionDecoderModel): + """``LTX2VideoDiffusionDecoderModel`` with tile-parallel decode across the SP group. + Drop-in replacement: use ``from_pretrained`` on this class exactly as on the base + class. The runner activates tile-parallel decode by: + 1. Setting ``decoder._parallel_decode = True`` (done when ``--use_parallel_vae`` is on). + 2. Calling ``decoder.enable_tiling()`` so that ``decode()`` dispatches to ``tiled_decode``. + When ``_parallel_decode`` is False, or ``sp_world_size == 1``, or the tile count is + less than 2, ``tiled_decode`` delegates to the stock implementation unchanged. + """ + + # Set to True by the runner when --use_parallel_vae is on. + _parallel_decode: bool = False + + def tiled_decode( + self, + z: torch.Tensor, + generator: torch.Generator | None = None, + num_inference_steps: int | None = None, + ) -> torch.Tensor: + """Decode with tiles distributed across SP ranks. + Each rank owns a round-robin subset of tiles. Owned tiles are computed locally; + all ranks gather all tiles via a shape-metadata all_reduce + per-tile broadcast, + then run the stock blend/assembly loop to produce the full video. + Noise determinism: + - *Shipping 1-step x0 path*: each tile is seeded from ``base_seed XOR f(t,h,w)``, + making noise rank-invariant without requiring non-owners to know tile shapes. + Output differs from single-GPU stock (different per-tile seeds) but is an + equally valid sample. + - *Multi-step path*: the full noise canvas is drawn identically on all ranks from + a shared generator seeded by ``base_seed``, then each owner slices its region. + """ + sp_world_size = get_sequence_parallel_world_size() + if not self._parallel_decode or sp_world_size <= 1: + return super().tiled_decode(z, generator=generator, + num_inference_steps=num_inference_steps) + + sp_group = get_sp_group() + sp_rank = get_sequence_parallel_rank() + + # Tile geometry + decoder = self.decoder + num_inference_steps = num_inference_steps or decoder.default_num_inference_steps + batch_size = z.shape[0] + patch_size = decoder.patch_size + + upsample_stride = decoder.upsamples[-1].stride + scale_t = upsample_stride[0] + scale_h = upsample_stride[1] * patch_size + scale_w = upsample_stride[2] * patch_size + + tile_t = self.tile_sample_min_num_frames // scale_t + stride_t = self.tile_sample_stride_num_frames // scale_t + tile_h = self.tile_sample_min_height // scale_h + stride_h = self.tile_sample_stride_height // scale_h + tile_w = self.tile_sample_min_width // scale_w + stride_w = self.tile_sample_stride_width // scale_w + + min_sizes = [ + max(k4, -(-k5 // s)) + for k4, k5, s in zip( + self.config.decoder_stage_kernels[-1], + self.config.decoder_stage5_kernel, + upsample_stride, + ) + ] + + # Cheap deterministic stages run on the full volume on every rank. + # This avoids communicating the intermediate features tensor and keeps the + # tile-grid computation identical across ranks. + features = decoder.forward_stages_1_to_3(z) + ghost_frames = decoder.trailing_pad_latent_frames * math.prod( + up.stride[0] for up in decoder.upsamples[:-1] + ) + num_frames = features.shape[1] - ghost_frames + height, width = features.shape[2], features.shape[3] + + temporal_tiles = _tile_intervals(num_frames, tile_t, stride_t, min_sizes[0]) + height_tiles = _tile_intervals(height, tile_h, stride_h, min_sizes[1]) + width_tiles = _tile_intervals(width, tile_w, stride_w, min_sizes[2]) + + blend_frames = (tile_t - stride_t) * scale_t + blend_height = (tile_h - stride_h) * scale_h + blend_width = (tile_w - stride_w) * scale_w + + n_tiles = len(temporal_tiles) * len(height_tiles) * len(width_tiles) + + # Fall back to stock when there are too few tiles to distribute. + if n_tiles < 2: + return super().tiled_decode(z, generator=generator, + num_inference_steps=num_inference_steps) + + # Noise setup + single_step_x0 = (num_inference_steps == 1 and decoder.model_output_type == "x0") + + # Draw one integer from the shared generator to derive a base seed. + # Generator state is guaranteed identical across SP ranks at decode entry, so + # all ranks draw the same value, advancing the generator identically. + if generator is not None: + base_seed_t = torch.randint(0, 2**31, (1,), generator=generator, device=z.device) + base_seed = int(base_seed_t.item()) + else: + # generator=None: use a fixed seed (in practice the runner always provides one) + base_seed = 0 + + x_t_full: torch.Tensor | None = None + if not single_step_x0: + pixel_frames = num_frames * scale_t - (1 if scale_t == 2 else 0) + shared_gen = torch.Generator(device=z.device).manual_seed(base_seed) + x_t_full = randn_tensor( + (batch_size, decoder.out_channels, + pixel_frames, height * scale_h, width * scale_w), + generator=shared_gen, + device=z.device, + dtype=z.dtype, + ) + + # Compute owned tiles + # shape_table[i] = (T_px, H_px, W_px) for tile i; zero for non-owned tiles. + # After all_reduce(SUM) every rank knows all tile pixel shapes. + shape_table = torch.zeros(n_tiles, 3, dtype=torch.int64, device=z.device) + tile_outputs: dict[int, torch.Tensor] = {} + + tile_idx = 0 + for t_idx, (t0, t1) in enumerate(temporal_tiles): + is_origin = (t0 == 0) + is_trailing = (t1 == num_frames) + # Trailing temporal tile carries ghost frames into forward_stage_4. + feature_t1 = features.shape[1] if is_trailing else t1 + + for h_idx, (h0, h1) in enumerate(height_tiles): + for w_idx, (w0, w1) in enumerate(width_tiles): + owner = tile_idx % sp_world_size + + if owner == sp_rank: + context = decoder.forward_stage_4( + features[:, t0:feature_t1, h0:h1, w0:w1], + drop_leading_frame=is_origin, + crop_trailing_ghost=is_trailing, + ) + tile_pixel_shape = ( + batch_size, + decoder.out_channels, + context.shape[1], + context.shape[2] * patch_size, + context.shape[3] * patch_size, + ) + + if single_step_x0: + # Per-tile deterministic seeding: owner-only draw, no + # cross-rank RNG synchronisation needed. + # Small integer hash keeps seeds well-separated. + tile_seed = (base_seed ^ (t_idx * 0x3_D6F1 + + h_idx * 0x1_E3 + + w_idx * 0x7)) & 0x7FFF_FFFF + tile_gen = torch.Generator(device=z.device).manual_seed(tile_seed) + x_t = randn_tensor(tile_pixel_shape, generator=tile_gen, + device=z.device, dtype=z.dtype) + else: + # Slice from the shared noise canvas (drawn identically on all + # ranks above). + 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_pixel_shape[2], + h0 * scale_h : h0 * scale_h + tile_pixel_shape[3], + w0 * scale_w : w0 * scale_w + tile_pixel_shape[4], + ] + + out = decoder.denoise(context, x_t, num_inference_steps) + tile_outputs[tile_idx] = out + shape_table[tile_idx, 0] = out.shape[2] # T_px + shape_table[tile_idx, 1] = out.shape[3] # H_px + shape_table[tile_idx, 2] = out.shape[4] # W_px + + tile_idx += 1 + + # Gather: shape metadata + per-tile broadcast + # all_reduce(SUM): safe because each element is non-zero on exactly one rank. + torch.distributed.all_reduce( + shape_table, + op=torch.distributed.ReduceOp.SUM, + group=sp_group.device_group, + ) + + gathered: dict[int, torch.Tensor] = {} + for i in range(n_tiles): + owner = i % sp_world_size + T_px = int(shape_table[i, 0]) + H_px = int(shape_table[i, 1]) + W_px = int(shape_table[i, 2]) + if sp_rank == owner: + buf = tile_outputs[i] + else: + buf = torch.empty( + batch_size, decoder.out_channels, T_px, H_px, W_px, + dtype=z.dtype, device=z.device, + ) + # src is the SP-group-local rank (0 … sp_world_size-1). + sp_group.broadcast(buf, src=owner) + gathered[i] = buf + + # Assembly — identical to stock tiled_decode + frame_groups: list[torch.Tensor] = [] + tile_idx = 0 + for _t_idx, (_t0, _t1) in enumerate(temporal_tiles): + rows: list[list[torch.Tensor]] = [] + for _h_idx, _ in enumerate(height_tiles): + row: list[torch.Tensor] = [] + for _w_idx, _ in enumerate(width_tiles): + row.append(gathered[tile_idx]) + tile_idx += 1 + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + keep_height = stride_h * scale_h if i < len(rows) - 1 else tile.shape[3] + keep_width = stride_w * scale_w if j < len(row) - 1 else tile.shape[4] + result_row.append(tile[:, :, :, :keep_height, :keep_width]) + result_rows.append(torch.cat(result_row, dim=4)) + frame_groups.append(torch.cat(result_rows, dim=3)) + + result = [] + for k, group in enumerate(frame_groups): + if k > 0: + group = self.blend_t(frame_groups[k - 1], group, blend_frames) + if k < len(frame_groups) - 1: + keep_frames = stride_t * scale_t - (1 if k == 0 and scale_t == 2 else 0) + group = group[:, :, :keep_frames] + result.append(group) + return torch.cat(result, dim=2) \ No newline at end of file diff --git a/xfuser/model_executor/layers/ltx2_na3d_eager_attn.py b/xfuser/model_executor/layers/ltx2/na3d_eager_attn.py similarity index 100% rename from xfuser/model_executor/layers/ltx2_na3d_eager_attn.py rename to xfuser/model_executor/layers/ltx2/na3d_eager_attn.py diff --git a/xfuser/model_executor/layers/norms.py b/xfuser/model_executor/layers/norms.py new file mode 100644 index 00000000..43975fec --- /dev/null +++ b/xfuser/model_executor/layers/norms.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import torch + +try: + import aiter as _aiter + + # Wrapping aiter.rms_norm as a non-mutating custom op fixes graph capture + @torch.library.custom_op("xfuser::aiter_rms_norm", mutates_args=()) + def _aiter_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: + return _aiter.rms_norm(x, w, eps) + + @_aiter_rms_norm.register_fake + def _(x, w, eps): + return torch.empty_like(x) + + @_aiter_rms_norm.register_autograd + def _(ctx, grad_output): + # Inference only, backward is never called; return identity gradient. + return grad_output, None, None + + _HAS_AITER = True +except ImportError: + _HAS_AITER = False + + +class _AITERRMSNorm(torch.nn.Module): + """Drop-in for diffusers RMSNorm using AITER's CK-Tile kernel. + Diffusers RMSNorm up-casts to float32 internally; AITER's kernel works + natively in bfloat16 and is faster. Handles elementwise_affine=False + (no learned weight) via a registered ones buffer so AITER always receives a + weight tensor. + """ + + def __init__(self, weight: torch.nn.Parameter | None, eps: float, dim: int) -> None: + super().__init__() + if weight is not None: + self.weight = weight + self._use_ones = False + else: + self.register_buffer("_ones_weight", torch.ones(dim)) + self._use_ones = True + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shape = x.shape + w = self._ones_weight if self._use_ones else self.weight + out = _aiter_rms_norm(x.view(-1, shape[-1]), w, self.eps) + return out.view(shape) + + +def _replace_rms_norms_with_aiter(model: torch.nn.Module) -> None: + """Walk ``model`` and replace diffusers RMSNorm modules with :class:`_AITERRMSNorm`. + Only replaces ``diffusers.models.normalization.RMSNorm`` (Python impl, + float32 up-cast), leaving ``torch.nn.RMSNorm`` (C++ impl) untouched. + Must be called BEFORE ``model.to(device)`` so buffers are placed correctly. + """ + if not _HAS_AITER: + return + try: + from diffusers.models.normalization import RMSNorm as _DiffusersRMSNorm + except ImportError: + return + + replacements: list[tuple[torch.nn.Module, str, _AITERRMSNorm]] = [] + for _parent_name, parent_module in model.named_modules(): + for child_name, child_module in parent_module.named_children(): + if type(child_module) is _DiffusersRMSNorm and hasattr(child_module, "eps"): + if child_module.weight is not None: + dim = child_module.weight.shape[0] + elif hasattr(child_module, "dim"): + dim = ( + int(child_module.dim[0]) + if hasattr(child_module.dim, "__len__") + else int(child_module.dim) + ) + else: + continue + replacements.append( + ( + parent_module, + child_name, + _AITERRMSNorm(child_module.weight, child_module.eps, dim), + ) + ) + for parent, name, replacement in replacements: + setattr(parent, name, replacement) \ No newline at end of file diff --git a/xfuser/model_executor/models/runner_models/ltx.py b/xfuser/model_executor/models/runner_models/ltx.py index 3b367fe2..92f63a57 100644 --- a/xfuser/model_executor/models/runner_models/ltx.py +++ b/xfuser/model_executor/models/runner_models/ltx.py @@ -445,6 +445,7 @@ class _xFuserLTX25VideoModelBase(xFuserModel): ulysses_degree=True, ring_degree=True, enable_tiling=True, + use_parallel_vae=True, use_fp8_gemms=True, use_fp4_gemms=True, ) @@ -497,12 +498,15 @@ def _load_model(self) -> DiffusionPipeline: # Diffusion decoder — replaces convolutional VAE decode for both distilled # and full model pipelines. - from diffusers import LTX2VideoDiffusionDecoderModel from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode import ( LTX2VideoDiffusionDecodePipeline, ) - diff_decoder = LTX2VideoDiffusionDecoderModel.from_pretrained( + from xfuser.model_executor.layers.ltx2.diffusion_decoder import ( + xFuserLTX2VideoDiffusionDecoderWrapper, + ) + + diff_decoder = xFuserLTX2VideoDiffusionDecoderWrapper.from_pretrained( self.settings.model_name, subfolder="diffusion_decoder", torch_dtype=torch.bfloat16, @@ -518,13 +522,13 @@ def _load_model(self) -> DiffusionPipeline: # NATTEN is not available # Fall back to tiled PyTorch SDPA # Works on CUDA, ROCm and CPU. Ported from LTX-2 EagerSdpaAttention. - from xfuser.model_executor.layers.ltx2_na3d_eager_attn import ( + from xfuser.model_executor.layers.ltx2.na3d_eager_attn import ( LTX2VideoVaeEagerSdpaAttnProcessor, ) diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor()) log( - "Diffusion decoder: NATTEN unavailable; using Triton na3d attention fallback." + "Diffusion decoder: NATTEN unavailable; using tiled PyTorch SDPA fallback." ) self.decode_pipe = LTX2VideoDiffusionDecodePipeline( diffusion_decoder=diff_decoder, @@ -535,9 +539,16 @@ def _load_model(self) -> DiffusionPipeline: def _enable_options(self) -> None: super()._enable_options() - if self.config.enable_tiling: + # Tiling on the diffusion decoder is enabled when either --enable_tiling is set + # (for single-GPU memory savings) or --use_parallel_vae is set (required for + # tiled_decode dispatch, which is what activates tile-parallel decode). + if self.config.enable_tiling or self.config.use_parallel_vae: self.pipe.vae.enable_tiling() self.decode_pipe.diffusion_decoder.enable_tiling() + # Gate tile-parallel decode behind --use_parallel_vae. The wrapper's tiled_decode + # checks this flag before distributing tiles across SP ranks. + if self.config.use_parallel_vae: + self.decode_pipe.diffusion_decoder._parallel_decode = True def _preprocess_args_images(self, input_args: dict) -> dict: input_args = super()._preprocess_args_images(input_args) diff --git a/xfuser/model_executor/models/transformers/transformer_ltx2.py b/xfuser/model_executor/models/transformers/transformer_ltx2.py index fabc13a4..8cc21890 100644 --- a/xfuser/model_executor/models/transformers/transformer_ltx2.py +++ b/xfuser/model_executor/models/transformers/transformer_ltx2.py @@ -18,6 +18,7 @@ AttentionMaskWithMeta, make_attn_mask_with_meta, ) +from xfuser.model_executor.layers.norms import _replace_rms_norms_with_aiter from xfuser.model_executor.layers.usp import USP, attention @@ -352,6 +353,11 @@ def __init__( self._enc_mask_cache: dict = {} self._audio_enc_mask_cache: dict = {} + # If AITER is available, replace diffusers RMSNorm (slow float32 cast) + # with AITER RMSNorm. + # Called BEFORE device move so buffers are placed correctly by model.to(). + _replace_rms_norms_with_aiter(self) + if perturbed_attn: attn_processor_cls = xFuserLTX2PerturbedAttnProcessor else: