sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

# Copied and adapted from LTX-2 and WanVideo implementations.
#
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import Any, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
    can_use_ltx2_qknorm_split_rope_cuda,
    ltx2_qknorm_split_rope_cuda,
)
from sglang.jit_kernel.diffusion.residual_gate_add import (
    can_use_residual_gate_add_cuda,
    residual_gate_add_cuda,
)
from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config
from sglang.multimodal_gen.runtime.distributed import (
    get_sp_parallel_rank,
    get_sp_world_size,
    get_tp_rank,
    get_tp_world_size,
    model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.distributed.communication_op import (
    sequence_model_parallel_all_gather,
    tensor_model_parallel_all_reduce,
)
from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
    AttentionBackendEnum,
    current_platform,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

_is_npu = current_platform.is_npu()

logger = init_logger(__name__)

ADALN_NUM_BASE_PARAMS = 6
ADALN_NUM_CROSS_ATTN_PARAMS = 3
_LTX2_RESIDUAL_GATE_CUDA_DISABLED = False
_LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = False


def _ltx2_residual_gate_add(
    residual: torch.Tensor,
    update: torch.Tensor,
    gate: torch.Tensor,
) -> torch.Tensor:
    global _LTX2_RESIDUAL_GATE_CUDA_DISABLED

    if not _LTX2_RESIDUAL_GATE_CUDA_DISABLED and can_use_residual_gate_add_cuda(
        residual, update, gate
    ):
        try:
            return residual_gate_add_cuda(residual, update, gate)
        except Exception as exc:
            if torch.compiler.is_compiling():
                raise
            logger.warning_once(f"Disabling LTX2 residual-gate CUDA fast path: {exc}")
            _LTX2_RESIDUAL_GATE_CUDA_DISABLED = True

    return residual + update * gate


def _ltx2_try_fused_qknorm_split_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: nn.Module,
    k_norm: nn.Module,
    q_cos: torch.Tensor,
    q_sin: torch.Tensor,
    k_cos: torch.Tensor,
    k_sin: torch.Tensor,
    *,
    eps: float,
    num_heads: int,
    head_dim: int,
) -> tuple[torch.Tensor, torch.Tensor] | None:
    global _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED

    if (
        _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED
        or get_tp_world_size() != 1
        or not isinstance(q_norm, nn.RMSNorm)
        or not isinstance(k_norm, nn.RMSNorm)
        or float(q_norm.eps) != float(eps)
        or float(k_norm.eps) != float(eps)
        or not can_use_ltx2_qknorm_split_rope_cuda(
            q,
            q_cos,
            q_sin,
            q_norm.weight,
            k,
            k_cos,
            k_sin,
            k_norm.weight,
            num_heads=num_heads,
            head_dim=head_dim,
        )
    ):
        return None

    try:
        return ltx2_qknorm_split_rope_cuda(
            q,
            q_cos,
            q_sin,
            q_norm.weight,
            k,
            k_cos,
            k_sin,
            k_norm.weight,
            eps=eps,
            num_heads=num_heads,
            head_dim=head_dim,
        )
    except Exception as exc:
        if torch.compiler.is_compiling():
            raise
        logger.warning_once(f"Disabling LTX2 QKNorm split-RoPE CUDA fast path: {exc}")
        _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = True
        return None


_LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = False


def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
    return ADALN_NUM_BASE_PARAMS + (
        ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0
    )


def _ltx2_disable_fused_ada_values(exc: Exception) -> None:
    global _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED
    _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = True
    logger.warning_once(f"Disabling LTX2 fused Ada values fast path: {exc}")


def _ltx2_try_fused_ada_values9(
    scale_shift_table: torch.Tensor,
    batch_size: int,
    timestep: torch.Tensor,
) -> tuple[torch.Tensor, ...] | None:
    if (
        _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED
        or get_tp_world_size() != 1
        or not timestep.is_cuda
        or timestep.dtype != torch.bfloat16
        or timestep.ndim != 3
        or int(timestep.shape[0]) != int(batch_size)
        or not timestep.is_contiguous()
        or not scale_shift_table.is_cuda
        or scale_shift_table.dtype not in (torch.bfloat16, torch.float32)
        or scale_shift_table.ndim != 2
        or int(scale_shift_table.shape[0]) != 9
        or scale_shift_table.stride(-1) != 1
    ):
        return None

    hidden = int(scale_shift_table.shape[1])
    if hidden % 256 != 0 or hidden > 8192 or timestep.shape[-1] != 9 * hidden:
        return None

    try:
        from sglang.jit_kernel.diffusion.triton.ltx2_ada_values import (
            ltx2_ada_values9,
        )

        return ltx2_ada_values9(scale_shift_table, timestep)
    except Exception as exc:
        _ltx2_disable_fused_ada_values(exc)
        return None


def _ltx2_is_perturbed(
    perturbation_config: dict[str, object],
    key: str,
    block_idx: int,
) -> bool:
    value = perturbation_config.get(key)
    if value is None:
        return False
    if key.endswith("_blocks"):
        return block_idx in value
    return bool(value)


def _ltx2_build_batched_perturbation_states(
    perturbation_configs: tuple[dict[str, object], ...],
    key: str,
    block_indices: tuple[int, ...],
    values: torch.Tensor,
) -> dict[int, tuple[torch.Tensor | None, bool]]:
    mask_cache: dict[tuple[int, ...], torch.Tensor] = {}
    states: dict[int, tuple[torch.Tensor | None, bool]] = {}
    for block_idx in block_indices:
        keep_values = []
        any_perturbed = False
        all_perturbed = True
        for config in perturbation_configs:
            perturbed = _ltx2_is_perturbed(config, key, block_idx)
            any_perturbed = any_perturbed or perturbed
            all_perturbed = all_perturbed and perturbed
            keep_values.append(0 if perturbed else 1)

        if not any_perturbed:
            states[block_idx] = (None, False)
        elif all_perturbed:
            states[block_idx] = (None, True)
        else:
            cache_key = tuple(keep_values)
            mask = mask_cache.get(cache_key)
            if mask is None:
                mask = torch.tensor(
                    keep_values, device=values.device, dtype=values.dtype
                ).view(len(keep_values), *([1] * (values.ndim - 1)))
                mask_cache[cache_key] = mask
            states[block_idx] = (mask, False)
    return states


def apply_interleaved_rotary_emb(
    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor:
    cos, sin = freqs
    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)
    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
    return x * cos + x_rotated * sin


def apply_split_rotary_emb(
    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor:
    cos, sin = freqs
    if (
        x.ndim == 3
        and cos.ndim == 4
        and sin.ndim == 4
        and x.dtype == torch.bfloat16
        and cos.dtype == torch.bfloat16
        and sin.dtype == torch.bfloat16
        and x.is_cuda
        and x.is_contiguous()
        and cos.is_cuda
        and sin.is_cuda
    ):
        from sglang.jit_kernel.diffusion.triton.ltx2_rotary import (
            apply_ltx2_split_rotary_emb,
        )

        return apply_ltx2_split_rotary_emb(x, cos, sin)

    x_dtype = x.dtype
    needs_reshape = False
    if x.ndim != 4 and cos.ndim == 4:
        b = x.shape[0]
        _, h, t, _ = cos.shape
        x = x.reshape(b, t, h, -1).swapaxes(1, 2)
        needs_reshape = True

    last = x.shape[-1]
    if last % 2 != 0:
        raise ValueError(
            f"Expected x.shape[-1] to be even for split rotary, got {last}."
        )
    r = last // 2

    split_x = x.reshape(*x.shape[:-1], 2, r)
    first_x = split_x[..., :1, :]
    second_x = split_x[..., 1:, :]

    cos_u = cos.unsqueeze(-2)
    sin_u = sin.unsqueeze(-2)

    out = split_x * cos_u
    first_out = out[..., :1, :]
    second_out = out[..., 1:, :]
    first_out.addcmul_(-sin_u, second_x)
    second_out.addcmul_(sin_u, first_x)

    out = out.reshape(*out.shape[:-2], last)
    if needs_reshape:
        out = out.swapaxes(1, 2).reshape(b, t, -1)
    return out.to(dtype=x_dtype)


# ==============================================================================
# Layers and Embeddings
# ==============================================================================


class LTX2AudioVideoRotaryPosEmbed(nn.Module):
    def __init__(
        self,
        dim: int,
        patch_size: int = 1,
        patch_size_t: int = 1,
        base_num_frames: int = 20,
        base_height: int = 2048,
        base_width: int = 2048,
        sampling_rate: int = 16000,
        hop_length: int = 160,
        scale_factors: Tuple[int, ...] = (8, 32, 32),
        theta: float = 10000.0,
        causal_offset: int = 1,
        modality: str = "video",
        double_precision: bool = True,
        rope_type: str = "interleaved",
        num_attention_heads: int = 32,
    ) -> None:
        super().__init__()
        self.dim = int(dim)
        self.patch_size = int(patch_size)
        self.patch_size_t = int(patch_size_t)

        if rope_type not in ["interleaved", "split"]:
            raise ValueError(
                f"{rope_type=} not supported. Choose between 'interleaved' and 'split'."
            )
        self.rope_type = rope_type

        self.base_num_frames = int(base_num_frames)
        self.num_attention_heads = int(num_attention_heads)

        self.base_height = int(base_height)
        self.base_width = int(base_width)

        self.sampling_rate = int(sampling_rate)
        self.hop_length = int(hop_length)
        self.audio_latents_per_second = (
            float(self.sampling_rate) / float(self.hop_length) / float(scale_factors[0])
        )

        self.scale_factors = tuple(int(x) for x in scale_factors)
        self.theta = float(theta)
        self.causal_offset = int(causal_offset)

        self.modality = modality
        if self.modality not in ["video", "audio"]:
            raise ValueError(
                f"Modality {modality} is not supported. Supported modalities are `video` and `audio`."
            )
        self.double_precision = bool(double_precision)

    def prepare_video_coords(
        self,
        batch_size: int,
        num_frames: int,
        height: int,
        width: int,
        device: torch.device,
        fps: float = 24.0,
        *,
        start_frame: int = 0,
    ) -> torch.Tensor:
        grid_f = torch.arange(
            start=int(start_frame),
            end=int(num_frames) + int(start_frame),
            step=self.patch_size_t,
            dtype=torch.float32,
            device=device,
        )
        grid_h = torch.arange(
            start=0,
            end=height,
            step=self.patch_size,
            dtype=torch.float32,
            device=device,
        )
        grid_w = torch.arange(
            start=0,
            end=width,
            step=self.patch_size,
            dtype=torch.float32,
            device=device,
        )
        grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij")
        grid = torch.stack(grid, dim=0)

        patch_size = (self.patch_size_t, self.patch_size, self.patch_size)
        patch_size_delta = torch.tensor(
            patch_size, dtype=grid.dtype, device=grid.device
        )
        patch_ends = grid + patch_size_delta.view(3, 1, 1, 1)

        latent_coords = torch.stack([grid, patch_ends], dim=-1)
        latent_coords = latent_coords.flatten(1, 3)
        latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1)

        scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device)
        broadcast_shape = [1] * latent_coords.ndim
        broadcast_shape[1] = -1
        pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape)
        pixel_coords[:, 0, ...] = (
            pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0]
        ).clamp(min=0)
        pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
        return pixel_coords

    def prepare_audio_coords(
        self,
        batch_size: int,
        num_frames: int,
        device: torch.device,
        *,
        start_frame: int = 0,
    ) -> torch.Tensor:
        grid_f = torch.arange(
            start=int(start_frame),
            end=int(num_frames) + int(start_frame),
            step=self.patch_size_t,
            dtype=torch.float32,
            device=device,
        )

        audio_scale_factor = self.scale_factors[0]
        grid_start_mel = grid_f * audio_scale_factor
        grid_start_mel = (
            grid_start_mel + self.causal_offset - audio_scale_factor
        ).clip(min=0)
        grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate

        grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor
        grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip(
            min=0
        )
        grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate

        audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1)
        audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1)
        audio_coords = audio_coords.unsqueeze(1)
        return audio_coords

    def prepare_coords(self, *args, **kwargs):
        if self.modality == "video":
            return self.prepare_video_coords(*args, **kwargs)
        return self.prepare_audio_coords(*args, **kwargs)

    def forward(
        self,
        coords: torch.Tensor,
        device: Optional[Union[str, torch.device]] = None,
        out_dtype: Optional[torch.dtype] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        device = device or coords.device
        out_dtype = out_dtype or coords.dtype
        num_pos_dims = coords.shape[1]

        if coords.ndim == 4:
            coords_start, coords_end = coords.chunk(2, dim=-1)
            coords = (coords_start + coords_end) / 2.0
            coords = coords.squeeze(-1)

        if self.modality == "video":
            max_positions = (self.base_num_frames, self.base_height, self.base_width)
        else:
            max_positions = (self.base_num_frames,)

        grid = torch.stack(
            [coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1
        ).to(device)

        num_rope_elems = num_pos_dims * 2
        # LTX-2.3 HQ is sensitive to RoPE rounding; keep frequency generation on
        # the target device instead of caching a CPU/NumPy tensor.
        freqs_dtype = torch.float64 if self.double_precision else torch.float32
        pow_indices = torch.pow(
            self.theta,
            torch.linspace(
                start=0.0,
                end=1.0,
                steps=self.dim // num_rope_elems,
                dtype=freqs_dtype,
                device=device,
            ),
        )
        freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)

        freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs
        freqs = freqs.transpose(-1, -2).flatten(2)

        if self.rope_type == "interleaved":
            cos_freqs = freqs.cos().repeat_interleave(2, dim=-1)
            sin_freqs = freqs.sin().repeat_interleave(2, dim=-1)

            if self.dim % num_rope_elems != 0:
                cos_padding = torch.ones_like(
                    cos_freqs[:, :, : self.dim % num_rope_elems]
                )
                sin_padding = torch.zeros_like(
                    cos_freqs[:, :, : self.dim % num_rope_elems]
                )
                cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1)
                sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1)
        else:
            expected_freqs = self.dim // 2
            current_freqs = freqs.shape[-1]
            pad_size = expected_freqs - current_freqs
            cos_freq = freqs.cos()
            sin_freq = freqs.sin()

            if pad_size != 0:
                cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
                sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
                cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
                sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)

            b = cos_freq.shape[0]
            t = cos_freq.shape[1]
            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1)
            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1)
            cos_freqs = torch.swapaxes(cos_freq, 1, 2)
            sin_freqs = torch.swapaxes(sin_freq, 1, 2)

        return cos_freqs.to(dtype=out_dtype), sin_freqs.to(dtype=out_dtype)


class LTX2TextProjection(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_size: int,
        out_features: int | None = None,
        act_fn: str = "gelu_tanh",
    ) -> None:
        super().__init__()
        if out_features is None:
            out_features = hidden_size

        self.linear_1 = ColumnParallelLinear(
            in_features, hidden_size, bias=True, gather_output=True
        )
        if act_fn == "gelu_tanh":
            self.act_1 = nn.GELU(approximate="tanh")
        elif act_fn == "silu":
            self.act_1 = nn.SiLU()
        else:
            raise ValueError(f"Unknown activation function: {act_fn}")

        self.linear_2 = ColumnParallelLinear(
            hidden_size, out_features, bias=True, gather_output=True
        )

    def forward(self, caption: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.linear_1(caption)
        hidden_states = self.act_1(hidden_states)
        hidden_states, _ = self.linear_2(hidden_states)
        return hidden_states


class LTX2TimestepEmbedder(nn.Module):
    def __init__(self, embedding_dim: int, in_channels: int = 256) -> None:
        super().__init__()
        self.linear_1 = ColumnParallelLinear(
            in_channels, embedding_dim, bias=True, gather_output=True
        )
        self.linear_2 = ColumnParallelLinear(
            embedding_dim, embedding_dim, bias=True, gather_output=True
        )

    def forward(self, t_emb: torch.Tensor) -> torch.Tensor:
        x, _ = self.linear_1(t_emb)
        x = F.silu(x)
        x, _ = self.linear_2(x)
        return x


class LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module):
    def __init__(self, embedding_dim: int) -> None:
        super().__init__()
        self.timestep_embedder = LTX2TimestepEmbedder(embedding_dim, in_channels=256)

    def forward(
        self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None
    ) -> torch.Tensor:
        t = timestep.reshape(-1).to(dtype=torch.float32)
        t_emb = timestep_embedding(t, dim=256, max_period=10000, dtype=torch.float32)
        if hidden_dtype is not None:
            t_emb = t_emb.to(dtype=hidden_dtype)
        return self.timestep_embedder(t_emb)


class LTX2AdaLayerNormSingle(nn.Module):
    def __init__(self, embedding_dim: int, embedding_coefficient: int = 6) -> None:
        super().__init__()
        self.emb = LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim)
        self.silu = nn.SiLU()
        self.linear = ColumnParallelLinear(
            embedding_dim,
            embedding_coefficient * embedding_dim,
            bias=True,
            gather_output=True,
        )

    def forward(
        self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype).to(
            dtype=self.linear.weight.dtype
        )
        out, _ = self.linear(self.silu(embedded_timestep))
        return out, embedded_timestep


class LTX2TPRMSNormAcrossHeads(nn.Module):
    def __init__(
        self, full_hidden_size: int, local_hidden_size: int, eps: float
    ) -> None:
        super().__init__()
        self.full_hidden_size = full_hidden_size
        self.local_hidden_size = local_hidden_size
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(local_hidden_size))

        tp_rank = get_tp_rank()

        def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
            shard = loaded_weight.narrow(
                0, tp_rank * local_hidden_size, local_hidden_size
            )
            param.data.copy_(shard.to(dtype=param.dtype, device=param.device))

        setattr(self.weight, "weight_loader", _weight_loader)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Keep track of the original dtype. We do the statistics in fp32 for
        # numerical stability, but cast the output back to the input dtype to
        orig_dtype = x.dtype
        if get_tp_world_size() == 1:
            var = x.float().pow(2).mean(dim=-1, keepdim=True)
        else:
            local_sumsq = x.float().pow(2).sum(dim=-1, keepdim=True)
            global_sumsq = tensor_model_parallel_all_reduce(local_sumsq)
            var = global_sumsq / float(self.full_hidden_size)

        inv_rms_fp32 = torch.rsqrt(var + self.eps)
        y = (x.float() * inv_rms_fp32).to(dtype=orig_dtype)
        return y * self.weight.to(dtype=orig_dtype)


class LTX2Attention(nn.Module):
    def __init__(
        self,
        query_dim: int,
        context_dim: int | None = None,
        heads: int = 8,
        dim_head: int = 64,
        norm_eps: float = 1e-6,
        qk_norm: bool = True,
        use_local_attention: bool = False,
        apply_gated_attention: bool = False,
        enable_packed_qkv_input_a2a: bool = False,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()

        self.query_dim = int(query_dim)
        self.context_dim = int(query_dim if context_dim is None else context_dim)
        self.heads = int(heads)
        self.dim_head = int(dim_head)
        self.inner_dim = self.heads * self.dim_head
        self.norm_eps = float(norm_eps)
        self.qk_norm = bool(qk_norm)
        self.use_local_attention = bool(use_local_attention)
        self.apply_gated_attention = bool(apply_gated_attention)
        self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a)
        self.prefix = prefix

        tp_size = get_tp_world_size()
        if tp_size <= 0:
            raise ValueError(f"Invalid {tp_size=}. Expected tp_size >= 1.")
        if self.heads % tp_size != 0:
            raise ValueError(
                f"LTX2Attention requires heads divisible by tp_size, got "
                f"{self.heads=} {tp_size=}."
            )
        if self.inner_dim % tp_size != 0:
            # This should follow from heads % tp_size, but keep explicit for clarity.
            raise ValueError(
                f"LTX2Attention requires inner_dim divisible by tp_size, got "
                f"{self.inner_dim=} {tp_size=}."
            )
        self.local_heads = self.heads // tp_size

        self.to_q = ColumnParallelLinear(
            self.query_dim,
            self.inner_dim,
            bias=True,
            gather_output=False,
            quant_config=quant_config,
        )
        self.to_k = ColumnParallelLinear(
            self.context_dim,
            self.inner_dim,
            bias=True,
            gather_output=False,
            quant_config=quant_config,
        )
        self.to_v = ColumnParallelLinear(
            self.context_dim,
            self.inner_dim,
            bias=True,
            gather_output=False,
            quant_config=quant_config,
        )
        self.to_gate_logits: ColumnParallelLinear | None = None
        if self.apply_gated_attention:
            self.to_gate_logits = ColumnParallelLinear(
                self.query_dim,
                self.heads,
                bias=True,
                gather_output=False,
                quant_config=quant_config,
            )

        self.q_norm: nn.Module | None = None
        self.k_norm: nn.Module | None = None
        if self.qk_norm:
            if tp_size == 1:
                self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)
                self.k_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)
            else:
                self.q_norm = LTX2TPRMSNormAcrossHeads(
                    full_hidden_size=self.inner_dim,
                    local_hidden_size=self.inner_dim // tp_size,
                    eps=self.norm_eps,
                )
                self.k_norm = LTX2TPRMSNormAcrossHeads(
                    full_hidden_size=self.inner_dim,
                    local_hidden_size=self.inner_dim // tp_size,
                    eps=self.norm_eps,
                )

        self.to_out = nn.Sequential(
            RowParallelLinear(
                self.inner_dim,
                self.query_dim,
                bias=True,
                input_is_parallel=True,
                quant_config=quant_config,
            ),
            nn.Identity(),
        )

        if self.use_local_attention:
            self.attn = LocalAttention(
                num_heads=self.local_heads,
                head_size=self.dim_head,
                num_kv_heads=self.local_heads,
                softmax_scale=None,
                causal=False,
                supported_attention_backends=supported_attention_backends,
                prefix=f"{prefix}.attn",
                enable_packed_qkv_input_a2a=self.enable_packed_qkv_input_a2a,
                # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
                allow_cudnn_sdp=True,
            )
        else:
            self.attn = USPAttention(
                num_heads=self.local_heads,
                head_size=self.dim_head,
                num_kv_heads=self.local_heads,
                dropout_rate=0,
                softmax_scale=None,
                causal=False,
                supported_attention_backends=supported_attention_backends,
                prefix=f"{prefix}.attn",
                # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
                allow_cudnn_sdp=True,
            )

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor | None = None,
        mask: torch.Tensor | None = None,
        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
        perturbation_mask: torch.Tensor | None = None,
        all_perturbed: bool = False,
        skip_sequence_parallel_override: bool = False,
        gather_context_kv_for_sp: bool = False,
        context_replicated_prefix_len: int = 0,
    ) -> torch.Tensor:
        gate_input = x
        context_ = x if context is None else context
        v, _ = self.to_v(context_)
        use_attention = not all_perturbed

        if use_attention:
            q, _ = self.to_q(x)
            k, _ = self.to_k(context_)

            fused_qk = None
            if pe is not None:
                cos, sin = pe
                k_cos, k_sin = pe if k_pe is None else k_pe
                tp_size = get_tp_world_size()
                if tp_size > 1:
                    tp_rank = get_tp_rank()
                    cos, sin = self._slice_rope_for_tp(
                        cos, sin, tp_rank=tp_rank, tp_size=tp_size
                    )
                    k_cos, k_sin = self._slice_rope_for_tp(
                        k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size
                    )
                if self.qk_norm and cos.dim() != 3:
                    assert self.q_norm is not None and self.k_norm is not None
                    fused_qk = _ltx2_try_fused_qknorm_split_rope(
                        q,
                        k,
                        self.q_norm,
                        self.k_norm,
                        cos,
                        sin,
                        k_cos,
                        k_sin,
                        eps=self.norm_eps,
                        num_heads=self.local_heads,
                        head_dim=self.dim_head,
                    )

            if fused_qk is not None:
                q, k = fused_qk
            else:
                if self.qk_norm:
                    assert self.q_norm is not None and self.k_norm is not None
                    q = self.q_norm(q)
                    k = self.k_norm(k)

                if pe is not None and cos.dim() == 3:
                    q = apply_interleaved_rotary_emb(q, (cos, sin))
                    k = apply_interleaved_rotary_emb(k, (k_cos, k_sin))
                elif pe is not None:
                    q = apply_split_rotary_emb(q, (cos, sin))
                    k = apply_split_rotary_emb(k, (k_cos, k_sin))

        v = v.view(*v.shape[:-1], self.local_heads, self.dim_head)
        if use_attention:
            q = q.view(*q.shape[:-1], self.local_heads, self.dim_head)
            k = k.view(*k.shape[:-1], self.local_heads, self.dim_head)

            if gather_context_kv_for_sp:
                # Replicated prefix (e.g. JoyEcho memory) is identical on every rank; only gather the sharded suffix.
                if context_replicated_prefix_len > 0:
                    prefix = int(context_replicated_prefix_len)
                    k_prefix, k_suffix = k[:, :prefix], k[:, prefix:]
                    v_prefix, v_suffix = v[:, :prefix], v[:, prefix:]
                    k_full = torch.cat(
                        [
                            k_prefix,
                            sequence_model_parallel_all_gather(
                                k_suffix.contiguous(), dim=1
                            ),
                        ],
                        dim=1,
                    )
                    v_full = torch.cat(
                        [
                            v_prefix,
                            sequence_model_parallel_all_gather(
                                v_suffix.contiguous(), dim=1
                            ),
                        ],
                        dim=1,
                    )
                    gathered_mask = mask
                else:
                    k_full = sequence_model_parallel_all_gather(k.contiguous(), dim=1)
                    v_full = sequence_model_parallel_all_gather(v.contiguous(), dim=1)
                    gathered_mask = None
                    if mask is not None:
                        gathered_mask = sequence_model_parallel_all_gather(
                            mask.contiguous(), dim=1
                        )
                if self.use_local_attention:
                    out = self.attn(q, k_full, v_full, attn_mask=gathered_mask)
                else:
                    out = self.attn(
                        q,
                        k_full,
                        v_full,
                        attn_mask=gathered_mask,
                        skip_sequence_parallel_override=True,
                    )
            elif self.use_local_attention:
                out = self.attn(q, k, v, attn_mask=mask)
            else:
                out = self.attn(
                    q,
                    k,
                    v,
                    attn_mask=mask,
                    skip_sequence_parallel_override=skip_sequence_parallel_override,
                )

            if perturbation_mask is not None:
                if perturbation_mask.ndim == out.ndim - 1:
                    perturbation_mask = perturbation_mask.unsqueeze(-1)
                out = out * perturbation_mask + v * (1 - perturbation_mask)

        if not use_attention:
            out = v

        if self.to_gate_logits is not None:
            gate_logits, _ = self.to_gate_logits(gate_input)
            b, t = out.shape[:2]
            out = out.view(b, t, self.local_heads, self.dim_head)
            out = out * (2.0 * torch.sigmoid(gate_logits).unsqueeze(-1))
            out = out.view(b, t, self.local_heads * self.dim_head)

        out_flat = out.flatten(2)
        out_proj, _ = self.to_out[0](out_flat)

        return out_proj

    def _slice_rope_for_tp(
        self,
        cos: torch.Tensor,
        sin: torch.Tensor,
        *,
        tp_rank: int,
        tp_size: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Slice RoPE tensors to the local TP shard.

        - split-rope: cos/sin are shaped [B, H, T, R] (head-major), slice by heads.
        - interleaved-rope: cos/sin are shaped [B, T, D], where D matches the projected
          feature dimension and is sharded by TP.
        """
        if cos.ndim == 4:
            # [B, H, T, R]
            start = tp_rank * self.local_heads
            end = start + self.local_heads
            return cos[:, start:end, :, :], sin[:, start:end, :, :]
        elif cos.ndim == 3:
            # [B, T, D]
            d = cos.shape[-1]
            if d % tp_size != 0:
                raise ValueError(
                    f"RoPE dim must be divisible by tp_size, got {d=} {tp_size=}."
                )
            local_d = d // tp_size
            start = tp_rank * local_d
            end = start + local_d
            return cos[:, :, start:end], sin[:, :, start:end]
        raise ValueError(f"Unexpected RoPE tensor rank: {cos.ndim}. Expected 3 or 4.")


class LTX2FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        dim_out: int | None = None,
        mult: int = 4,
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__()
        if dim_out is None:
            dim_out = dim
        inner_dim = int(dim * mult)

        self.proj_in = ColumnParallelLinear(
            dim, inner_dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.act = nn.GELU(approximate="tanh")
        self.proj_out = RowParallelLinear(
            inner_dim,
            dim_out,
            bias=True,
            input_is_parallel=True,
            quant_config=quant_config,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.proj_in(x)
        x = self.act(x)
        x, _ = self.proj_out(x)
        return x


class LTX2TransformerBlock(nn.Module):
    def __init__(
        self,
        idx: int,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        cross_attention_dim: int,
        audio_dim: int,
        audio_num_attention_heads: int,
        audio_attention_head_dim: int,
        audio_cross_attention_dim: int,
        qk_norm: bool = True,
        norm_eps: float = 1e-6,
        apply_gated_attention: bool = False,
        cross_attention_adaln: bool = False,
        use_local_av_cross_attention: bool = False,
        force_sdpa_v2a_cross_attention: bool = False,
        enable_packed_qkv_input_a2a: bool = False,
        supported_attention_backends: set[AttentionBackendEnum] | None = None,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.idx = idx
        self.norm_eps = norm_eps
        self.rms_norm = RMSNormNoWeight()
        # LTX2.3
        self.cross_attention_adaln = cross_attention_adaln
        self.use_local_av_cross_attention = use_local_av_cross_attention

        # 1. Self-Attention (video and audio)
        self.attn1 = LTX2Attention(
            query_dim=dim,
            heads=num_attention_heads,
            dim_head=attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            apply_gated_attention=apply_gated_attention,
            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn1",
            quant_config=quant_config,
        )
        self.audio_attn1 = LTX2Attention(
            query_dim=audio_dim,
            heads=audio_num_attention_heads,
            dim_head=audio_attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            apply_gated_attention=apply_gated_attention,
            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_attn1",
            quant_config=quant_config,
        )

        # 2. Prompt Cross-Attention
        # Prompt KV is replicated across SP ranks, so prompt cross-attn should
        # stay local and preserve the explicit KV mask semantics from official.
        self.attn2 = LTX2Attention(
            query_dim=dim,
            context_dim=cross_attention_dim,
            heads=num_attention_heads,
            dim_head=attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            use_local_attention=True,
            apply_gated_attention=apply_gated_attention,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn2",
            quant_config=quant_config,
        )
        self.audio_attn2 = LTX2Attention(
            query_dim=audio_dim,
            context_dim=audio_cross_attention_dim,
            heads=audio_num_attention_heads,
            dim_head=audio_attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            use_local_attention=True,
            apply_gated_attention=apply_gated_attention,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_attn2",
            quant_config=quant_config,
        )

        # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
        self.audio_to_video_attn = LTX2Attention(
            query_dim=dim,
            context_dim=audio_dim,
            heads=audio_num_attention_heads,
            dim_head=audio_attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            use_local_attention=use_local_av_cross_attention,
            apply_gated_attention=apply_gated_attention,
            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.audio_to_video_attn",
            quant_config=quant_config,
        )
        self.video_to_audio_attn = LTX2Attention(
            query_dim=audio_dim,
            context_dim=dim,
            heads=audio_num_attention_heads,
            dim_head=audio_attention_head_dim,
            norm_eps=norm_eps,
            qk_norm=qk_norm,
            use_local_attention=use_local_av_cross_attention,
            apply_gated_attention=apply_gated_attention,
            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
            supported_attention_backends=(
                {AttentionBackendEnum.TORCH_SDPA}
                if force_sdpa_v2a_cross_attention
                else supported_attention_backends
            ),
            prefix=f"{prefix}.video_to_audio_attn",
            quant_config=quant_config,
        )

        # 4. Feedforward layers
        self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config)
        self.audio_ff = LTX2FeedForward(
            audio_dim, dim_out=audio_dim, quant_config=quant_config
        )

        # 5. Modulation Parameters
        num_ada_params = adaln_embedding_coefficient(cross_attention_adaln)
        self.scale_shift_table = nn.Parameter(
            torch.randn(num_ada_params, dim) / dim**0.5
        )
        self.audio_scale_shift_table = nn.Parameter(
            torch.randn(num_ada_params, audio_dim) / audio_dim**0.5
        )
        self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim))
        self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(
            torch.randn(5, audio_dim)
        )
        if self.cross_attention_adaln:
            # LTX2.3
            self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
            self.audio_prompt_scale_shift_table = nn.Parameter(
                torch.randn(2, audio_dim)
            )

    def get_ada_values(
        self,
        scale_shift_table: torch.Tensor,
        batch_size: int,
        timestep: torch.Tensor,
        indices: slice,
    ) -> tuple[torch.Tensor, ...]:
        num_ada_params = int(scale_shift_table.shape[0])
        ada_values = (
            scale_shift_table[indices]
            .unsqueeze(0)
            .unsqueeze(0)
            .to(device=timestep.device, dtype=timestep.dtype)
            + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[
                :, :, indices, :
            ]
        ).unbind(dim=2)
        return [t.squeeze(2) for t in ada_values]

    def forward(
        self,
        hidden_states: torch.Tensor,
        audio_hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        audio_encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        temb_audio: torch.Tensor,
        temb_prompt: torch.Tensor | None,
        temb_audio_prompt: torch.Tensor | None,
        temb_ca_scale_shift: torch.Tensor,
        temb_ca_audio_scale_shift: torch.Tensor,
        temb_ca_gate: torch.Tensor,
        temb_ca_audio_gate: torch.Tensor,
        video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        ca_video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        ca_audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        audio_encoder_attention_mask: Optional[torch.Tensor] = None,
        video_self_attention_mask: Optional[torch.Tensor] = None,
        audio_self_attention_mask: Optional[torch.Tensor] = None,
        a2v_cross_attention_mask: Optional[torch.Tensor] = None,
        v2a_cross_attention_mask: Optional[torch.Tensor] = None,
        skip_video_self_attn: bool = False,
        skip_audio_self_attn: bool = False,
        skip_a2v_cross_attn: bool = False,
        skip_v2a_cross_attn: bool = False,
        video_self_attn_perturbation_mask: Optional[torch.Tensor] = None,
        audio_self_attn_perturbation_mask: Optional[torch.Tensor] = None,
        a2v_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
        v2a_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
        audio_replicated_for_sp: bool = False,
        video_memory_prefix_len: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size = hidden_states.size(0)
        video_ada_values = _ltx2_try_fused_ada_values9(
            self.scale_shift_table, batch_size, temb
        )
        audio_ada_values = _ltx2_try_fused_ada_values9(
            self.audio_scale_shift_table, batch_size, temb_audio
        )

        # 1. Video and Audio Self-Attention
        if video_ada_values is None:
            vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
                self.scale_shift_table, batch_size, temb, slice(0, 3)
            )
        else:
            vshift_msa, vscale_msa, vgate_msa = video_ada_values[0:3]
        norm_hidden_states = (
            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa
        )
        attn_hidden_states = self.attn1(
            norm_hidden_states,
            mask=video_self_attention_mask,
            pe=video_rotary_emb,
            perturbation_mask=video_self_attn_perturbation_mask,
            all_perturbed=skip_video_self_attn,
            gather_context_kv_for_sp=audio_replicated_for_sp,
            context_replicated_prefix_len=video_memory_prefix_len,
        )
        hidden_states = _ltx2_residual_gate_add(
            hidden_states, attn_hidden_states, vgate_msa
        )

        if audio_ada_values is None:
            ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
                self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3)
            )
        else:
            ashift_msa, ascale_msa, agate_msa = audio_ada_values[0:3]
        norm_audio_hidden_states = (
            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa)
            + ashift_msa
        )
        attn_audio_hidden_states = self.audio_attn1(
            norm_audio_hidden_states,
            mask=audio_self_attention_mask,
            pe=audio_rotary_emb,
            perturbation_mask=audio_self_attn_perturbation_mask,
            all_perturbed=skip_audio_self_attn,
            skip_sequence_parallel_override=audio_replicated_for_sp,
        )
        audio_hidden_states = _ltx2_residual_gate_add(
            audio_hidden_states, attn_audio_hidden_states, agate_msa
        )
        # 2. Prompt Cross-Attention
        if self.cross_attention_adaln:
            # LTX2.3
            if temb_prompt is None or temb_audio_prompt is None:
                raise ValueError(
                    "cross_attention_adaln requires prompt modulation tensors."
                )
            if video_ada_values is None:
                vshift_q, vscale_q, vgate_q = self.get_ada_values(
                    self.scale_shift_table, batch_size, temb, slice(6, 9)
                )
            else:
                vshift_q, vscale_q, vgate_q = video_ada_values[6:9]
            v_prompt_shift, v_prompt_scale = self.get_ada_values(
                self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None)
            )
            norm_hidden_states = (
                self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q
            )
            mod_encoder_hidden_states = (
                encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift
            )
            attn_hidden_states = self.attn2(
                norm_hidden_states,
                context=mod_encoder_hidden_states,
                mask=encoder_attention_mask,
            )
            hidden_states = _ltx2_residual_gate_add(
                hidden_states, attn_hidden_states, vgate_q
            )

            if audio_ada_values is None:
                ashift_q, ascale_q, agate_q = self.get_ada_values(
                    self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9)
                )
            else:
                ashift_q, ascale_q, agate_q = audio_ada_values[6:9]
            a_prompt_shift, a_prompt_scale = self.get_ada_values(
                self.audio_prompt_scale_shift_table,
                batch_size,
                temb_audio_prompt,
                slice(None),
            )
            norm_audio_hidden_states = (
                self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q)
                + ashift_q
            )
            mod_audio_encoder_hidden_states = (
                audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift
            )
            attn_audio_hidden_states = self.audio_attn2(
                norm_audio_hidden_states,
                context=mod_audio_encoder_hidden_states,
                mask=audio_encoder_attention_mask,
            )
            audio_hidden_states = _ltx2_residual_gate_add(
                audio_hidden_states, attn_audio_hidden_states, agate_q
            )
        else:
            norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps)
            attn_hidden_states = self.attn2(
                norm_hidden_states,
                context=encoder_hidden_states,
                mask=encoder_attention_mask,
            )
            hidden_states = hidden_states + attn_hidden_states

            norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps)
            attn_audio_hidden_states = self.audio_attn2(
                norm_audio_hidden_states,
                context=audio_encoder_hidden_states,
                mask=audio_encoder_attention_mask,
            )
            audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
        # 3. Audio-to-Video and Video-to-Audio Cross-Attention
        norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps)
        norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps)

        # Compute combined ada params
        video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[
            :4, :
        ]
        video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]

        video_ca_scale_shift_table = (
            video_per_layer_ca_scale_shift[None, None, :, :].to(
                dtype=temb_ca_scale_shift.dtype, device=temb_ca_scale_shift.device
            )
            + temb_ca_scale_shift.reshape(
                batch_size, temb_ca_scale_shift.shape[1], 4, -1
            )
        ).unbind(dim=2)
        video_ca_gate = (
            video_per_layer_ca_gate[None, None, :, :].to(
                dtype=temb_ca_gate.dtype, device=temb_ca_gate.device
            )
            + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1)
        ).unbind(dim=2)

        (
            video_a2v_ca_scale,
            video_a2v_ca_shift,
            video_v2a_ca_scale,
            video_v2a_ca_shift,
        ) = [t.squeeze(2) for t in video_ca_scale_shift_table]
        a2v_gate = video_ca_gate[0].squeeze(2)

        audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[
            :4, :
        ]
        audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]

        audio_ca_scale_shift_table = (
            audio_per_layer_ca_scale_shift[None, None, :, :].to(
                dtype=temb_ca_audio_scale_shift.dtype,
                device=temb_ca_audio_scale_shift.device,
            )
            + temb_ca_audio_scale_shift.reshape(
                batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1
            )
        ).unbind(dim=2)
        audio_ca_gate = (
            audio_per_layer_ca_gate[None, None, :, :].to(
                dtype=temb_ca_audio_gate.dtype, device=temb_ca_audio_gate.device
            )
            + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1)
        ).unbind(dim=2)

        (
            audio_a2v_ca_scale,
            audio_a2v_ca_shift,
            audio_v2a_ca_scale,
            audio_v2a_ca_shift,
        ) = [t.squeeze(2) for t in audio_ca_scale_shift_table]
        v2a_gate = audio_ca_gate[0].squeeze(2)

        # A2V
        mod_norm_hidden_states = (
            norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift
        )
        mod_norm_audio_hidden_states = (
            norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift
        )

        if not skip_a2v_cross_attn:
            a2v_attn_hidden_states = self.audio_to_video_attn(
                mod_norm_hidden_states,
                context=mod_norm_audio_hidden_states,
                pe=ca_video_rotary_emb,
                k_pe=ca_audio_rotary_emb,
                mask=a2v_cross_attention_mask,
                skip_sequence_parallel_override=audio_replicated_for_sp,
            )
            if a2v_cross_attn_perturbation_mask is not None:
                a2v_attn_hidden_states = (
                    a2v_attn_hidden_states * a2v_cross_attn_perturbation_mask
                )
            hidden_states = _ltx2_residual_gate_add(
                hidden_states, a2v_attn_hidden_states, a2v_gate
            )

        # V2A
        mod_norm_hidden_states = (
            norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift
        )
        mod_norm_audio_hidden_states = (
            norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift
        )

        if not skip_v2a_cross_attn:
            v2a_attn_hidden_states = self.video_to_audio_attn(
                mod_norm_audio_hidden_states,
                context=mod_norm_hidden_states,
                pe=ca_audio_rotary_emb,
                k_pe=ca_video_rotary_emb,
                mask=v2a_cross_attention_mask,
                gather_context_kv_for_sp=audio_replicated_for_sp,
                context_replicated_prefix_len=video_memory_prefix_len,
            )
            if v2a_cross_attn_perturbation_mask is not None:
                v2a_attn_hidden_states = (
                    v2a_attn_hidden_states * v2a_cross_attn_perturbation_mask
                )
            audio_hidden_states = _ltx2_residual_gate_add(
                audio_hidden_states, v2a_attn_hidden_states, v2a_gate
            )
        # 4. Feedforward
        if video_ada_values is None:
            vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
                self.scale_shift_table, batch_size, temb, slice(3, 6)
            )
        else:
            vshift_mlp, vscale_mlp, vgate_mlp = video_ada_values[3:6]
        norm_hidden_states = (
            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
        )
        ff_output = self.ff(norm_hidden_states)
        hidden_states = _ltx2_residual_gate_add(hidden_states, ff_output, vgate_mlp)

        if audio_ada_values is None:
            ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
                self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6)
            )
        else:
            ashift_mlp, ascale_mlp, agate_mlp = audio_ada_values[3:6]
        norm_audio_hidden_states = (
            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp)
            + ashift_mlp
        )
        audio_ff_output = self.audio_ff(norm_audio_hidden_states)
        audio_hidden_states = _ltx2_residual_gate_add(
            audio_hidden_states, audio_ff_output, agate_mlp
        )
        return hidden_states, audio_hidden_states


class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
    _fsdp_shard_conditions = LTX2ArchConfig()._fsdp_shard_conditions
    _compile_conditions = LTX2ArchConfig()._compile_conditions
    _supported_attention_backends = LTX2ArchConfig()._supported_attention_backends
    param_names_mapping = LTX2ArchConfig().param_names_mapping
    reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping
    lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping

    @staticmethod
    def _collapse_prompt_timestep(timestep: torch.Tensor) -> torch.Tensor:
        if timestep.ndim <= 1:
            return timestep
        return timestep.amax(dim=tuple(range(1, timestep.ndim)))

    def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor:
        ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2"))
        if ltx_variant == "ltx_2_3" and bool(
            getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False)
        ):
            return timestep * float(self.timestep_scale_multiplier)
        return timestep

    def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None:
        """Validate TP-related dimension constraints (fail-fast)."""
        if tp_size < 1:
            raise ValueError(f"Invalid tp_size={tp_size}. Expected tp_size >= 1.")

        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                "video hidden_size must be divisible by num_attention_heads, got "
                f"{self.hidden_size=} {self.num_attention_heads=}."
            )
        if self.audio_hidden_size % self.audio_num_attention_heads != 0:
            raise ValueError(
                "audio_hidden_size must be divisible by audio_num_attention_heads, got "
                f"{self.audio_hidden_size=} {self.audio_num_attention_heads=}."
            )

        if tp_size == 1:
            return

        if self.num_attention_heads % tp_size != 0:
            raise ValueError(
                "num_attention_heads must be divisible by tp_size, got "
                f"{self.num_attention_heads=} {tp_size=}."
            )
        if self.audio_num_attention_heads % tp_size != 0:
            raise ValueError(
                "audio_num_attention_heads must be divisible by tp_size, got "
                f"{self.audio_num_attention_heads=} {tp_size=}."
            )
        if self.hidden_size % tp_size != 0:
            raise ValueError(
                "hidden_size must be divisible by tp_size for TP-sharded projections, got "
                f"{self.hidden_size=} {tp_size=}."
            )
        if self.audio_hidden_size % tp_size != 0:
            raise ValueError(
                "audio_hidden_size must be divisible by tp_size for TP-sharded projections, got "
                f"{self.audio_hidden_size=} {tp_size=}."
            )
        if int(arch.out_channels) % tp_size != 0:
            raise ValueError(
                "out_channels must be divisible by tp_size for TP-sharded output projection, got "
                f"{arch.out_channels=} {tp_size=}."
            )
        if int(arch.audio_out_channels) % tp_size != 0:
            raise ValueError(
                "audio_out_channels must be divisible by tp_size for TP-sharded output projection, got "
                f"{arch.audio_out_channels=} {tp_size=}."
            )

    def __init__(
        self,
        config: LTX2Config,
        hf_config: dict[str, Any],
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        arch = config.arch_config
        self.hidden_size = arch.hidden_size
        self.num_attention_heads = arch.num_attention_heads
        self.audio_hidden_size = arch.audio_hidden_size
        self.audio_num_attention_heads = arch.audio_num_attention_heads
        self.norm_eps = arch.norm_eps

        tp_size = get_tp_world_size()
        self._validate_tp_config(arch=arch, tp_size=tp_size)

        # 1. Patchification input projections
        # Matches LTX2Config().param_names_mapping
        self.patchify_proj = ColumnParallelLinear(
            arch.in_channels,
            self.hidden_size,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
        )
        self.audio_patchify_proj = ColumnParallelLinear(
            arch.audio_in_channels,
            self.audio_hidden_size,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
        )

        # 2. Prompt embeddings
        self.caption_projection: LTX2TextProjection | None = None
        self.audio_caption_projection: LTX2TextProjection | None = None
        if not arch.caption_proj_before_connector:
            self.caption_projection = LTX2TextProjection(
                in_features=arch.caption_channels, hidden_size=self.hidden_size
            )
            self.audio_caption_projection = LTX2TextProjection(
                in_features=arch.caption_channels, hidden_size=self.audio_hidden_size
            )

        # 3. Timestep Modulation Params and Embedding
        self.adaln_single = LTX2AdaLayerNormSingle(
            self.hidden_size,
            embedding_coefficient=adaln_embedding_coefficient(
                arch.cross_attention_adaln
            ),
        )
        self.audio_adaln_single = LTX2AdaLayerNormSingle(
            self.audio_hidden_size,
            embedding_coefficient=adaln_embedding_coefficient(
                arch.cross_attention_adaln
            ),
        )
        self.prompt_adaln_single: LTX2AdaLayerNormSingle | None = None
        self.audio_prompt_adaln_single: LTX2AdaLayerNormSingle | None = None
        if arch.cross_attention_adaln:
            self.prompt_adaln_single = LTX2AdaLayerNormSingle(
                self.hidden_size, embedding_coefficient=2
            )
            self.audio_prompt_adaln_single = LTX2AdaLayerNormSingle(
                self.audio_hidden_size, embedding_coefficient=2
            )

        # Global Cross Attention Modulation Parameters
        self.av_ca_video_scale_shift_adaln_single = LTX2AdaLayerNormSingle(
            self.hidden_size, embedding_coefficient=4
        )
        self.av_ca_a2v_gate_adaln_single = LTX2AdaLayerNormSingle(
            self.hidden_size, embedding_coefficient=1
        )
        self.av_ca_audio_scale_shift_adaln_single = LTX2AdaLayerNormSingle(
            self.audio_hidden_size, embedding_coefficient=4
        )
        self.av_ca_v2a_gate_adaln_single = LTX2AdaLayerNormSingle(
            self.audio_hidden_size, embedding_coefficient=1
        )

        # Output Layer Scale/Shift Modulation parameters
        self.scale_shift_table = nn.Parameter(
            torch.randn(2, self.hidden_size) / self.hidden_size**0.5
        )
        self.audio_scale_shift_table = nn.Parameter(
            torch.randn(2, self.audio_hidden_size) / self.audio_hidden_size**0.5
        )

        hf_patch_size = int(hf_config.get("patch_size", 1))
        hf_patch_size_t = int(hf_config.get("patch_size_t", 1))
        self.patch_size = (hf_patch_size_t, hf_patch_size, hf_patch_size)

        hf_audio_patch_size = int(hf_config.get("audio_patch_size", 1))
        hf_audio_patch_size_t = int(hf_config.get("audio_patch_size_t", 1))

        rope_type = (
            arch.rope_type.value
            if hasattr(arch.rope_type, "value")
            else str(arch.rope_type)
        )
        frequencies_precision = hf_config.get("frequencies_precision")
        if frequencies_precision is None:
            frequencies_precision = getattr(arch, "frequencies_precision", None)

        # diffusers/LTX configs use `frequencies_precision` for this RoPE switch
        rope_double_precision = (
            str(frequencies_precision) == "float64"
            if frequencies_precision is not None
            else bool(
                hf_config.get("rope_double_precision", arch.double_precision_rope)
            )
        )
        self.quantize_video_rope_coords_to_hidden_dtype = bool(
            hf_config.get("quantize_video_rope_coords_to_hidden_dtype", False)
        )
        causal_offset = int(hf_config.get("causal_offset", 1))

        pos_embed_max_pos = int(arch.positional_embedding_max_pos[0])
        base_height = int(arch.positional_embedding_max_pos[1])
        base_width = int(arch.positional_embedding_max_pos[2])

        audio_pos_embed_max_pos = int(arch.audio_positional_embedding_max_pos[0])

        self.video_scale_factors = (8, 32, 32)
        self.audio_scale_factors = (4,)

        self.rope = LTX2AudioVideoRotaryPosEmbed(
            dim=self.hidden_size,
            patch_size=hf_patch_size,
            patch_size_t=hf_patch_size_t,
            base_num_frames=pos_embed_max_pos,
            base_height=base_height,
            base_width=base_width,
            scale_factors=self.video_scale_factors,
            theta=float(arch.positional_embedding_theta),
            causal_offset=causal_offset,
            modality="video",
            double_precision=rope_double_precision,
            rope_type=rope_type,
            num_attention_heads=self.num_attention_heads,
        )
        self.audio_rope = LTX2AudioVideoRotaryPosEmbed(
            dim=self.audio_hidden_size,
            patch_size=hf_audio_patch_size,
            patch_size_t=hf_audio_patch_size_t,
            base_num_frames=audio_pos_embed_max_pos,
            sampling_rate=16000,
            hop_length=160,
            scale_factors=self.audio_scale_factors,
            theta=float(arch.positional_embedding_theta),
            causal_offset=causal_offset,
            modality="audio",
            double_precision=rope_double_precision,
            rope_type=rope_type,
            num_attention_heads=self.audio_num_attention_heads,
        )

        cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos)
        self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed(
            dim=int(arch.audio_cross_attention_dim),
            patch_size=hf_patch_size,
            patch_size_t=hf_patch_size_t,
            base_num_frames=cross_attn_pos_embed_max_pos,
            base_height=base_height,
            base_width=base_width,
            theta=float(arch.positional_embedding_theta),
            causal_offset=causal_offset,
            modality="video",
            double_precision=rope_double_precision,
            rope_type=rope_type,
            num_attention_heads=self.num_attention_heads,
        )
        self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed(
            dim=int(arch.audio_cross_attention_dim),
            patch_size=hf_audio_patch_size,
            patch_size_t=hf_audio_patch_size_t,
            base_num_frames=cross_attn_pos_embed_max_pos,
            sampling_rate=16000,
            hop_length=160,
            scale_factors=self.audio_scale_factors,
            theta=float(arch.positional_embedding_theta),
            causal_offset=causal_offset,
            modality="audio",
            double_precision=rope_double_precision,
            rope_type=rope_type,
            num_attention_heads=self.audio_num_attention_heads,
        )

        self.cross_pe_max_pos = cross_attn_pos_embed_max_pos

        # 5. Transformer Blocks
        self.transformer_blocks = nn.ModuleList(
            [
                LTX2TransformerBlock(
                    idx=idx,
                    dim=self.hidden_size,
                    num_attention_heads=self.num_attention_heads,
                    attention_head_dim=self.hidden_size // self.num_attention_heads,
                    cross_attention_dim=arch.cross_attention_dim,
                    audio_dim=self.audio_hidden_size,
                    audio_num_attention_heads=self.audio_num_attention_heads,
                    audio_attention_head_dim=self.audio_hidden_size
                    // self.audio_num_attention_heads,
                    audio_cross_attention_dim=arch.audio_cross_attention_dim,
                    norm_eps=self.norm_eps,
                    qk_norm=True,  # Always True in LTX2
                    apply_gated_attention=arch.apply_gated_attention,
                    cross_attention_adaln=arch.cross_attention_adaln,
                    use_local_av_cross_attention=bool(
                        getattr(arch, "use_local_av_cross_attention", False)
                    ),
                    force_sdpa_v2a_cross_attention=bool(
                        getattr(arch, "force_sdpa_v2a_cross_attention", False)
                    ),
                    enable_packed_qkv_input_a2a=arch.enable_packed_qkv_input_a2a,
                    supported_attention_backends=self._supported_attention_backends,
                    prefix=config.prefix,
                    quant_config=quant_config,
                )
                for idx in range(arch.num_layers)
            ]
        )

        # 6. Output layers
        self.norm_out = nn.LayerNorm(
            self.hidden_size, eps=self.norm_eps, elementwise_affine=False
        )
        self.proj_out = ColumnParallelLinear(
            self.hidden_size,
            arch.out_channels,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
        )

        self.audio_norm_out = nn.LayerNorm(
            self.audio_hidden_size, eps=self.norm_eps, elementwise_affine=False
        )
        self.audio_proj_out = ColumnParallelLinear(
            self.audio_hidden_size,
            arch.audio_out_channels,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
        )

        self.out_channels_raw = arch.out_channels // (
            self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
        )
        self.audio_out_channels = arch.audio_out_channels
        self.timestep_scale_multiplier = arch.timestep_scale_multiplier
        self.av_ca_timestep_scale_multiplier = arch.av_ca_timestep_scale_multiplier

        self.layer_names = ["transformer_blocks"]

    def _maybe_quantize_video_rope_coords(
        self,
        video_coords: torch.Tensor,
        hidden_device: torch.device,
        hidden_dtype: torch.dtype,
    ) -> torch.Tensor:
        if self.quantize_video_rope_coords_to_hidden_dtype:
            return video_coords.to(device=hidden_device, dtype=hidden_dtype)
        return video_coords.to(device=hidden_device)

    def _get_av_ca_gate_timestep_factor(self) -> float:
        ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2"))
        if ltx_variant == "ltx_2_3":
            return self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier
        return float(self.av_ca_timestep_scale_multiplier)

    def _get_av_ca_timesteps(
        self,
        timestep: torch.Tensor,
        audio_timestep: torch.Tensor,
        prompt_timestep: torch.Tensor | None,
        audio_prompt_timestep: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2"))
        if ltx_variant != "ltx_2_3":
            return timestep, audio_timestep

        video_timestep = (
            self._collapse_prompt_timestep(timestep)
            if prompt_timestep is None
            else prompt_timestep
        )
        audio_timestep_for_ca = (
            self._collapse_prompt_timestep(audio_timestep)
            if audio_prompt_timestep is None
            else audio_prompt_timestep
        )
        return video_timestep, audio_timestep_for_ca

    def forward(
        self,
        hidden_states: torch.Tensor,
        audio_hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        audio_encoder_hidden_states: torch.Tensor,
        timestep: torch.LongTensor,
        audio_timestep: Optional[torch.LongTensor] = None,
        prompt_timestep: Optional[torch.Tensor] = None,
        audio_prompt_timestep: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        audio_encoder_attention_mask: Optional[torch.Tensor] = None,
        num_frames: Optional[int] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        fps: float = 24.0,
        audio_num_frames: Optional[int] = None,
        video_coords: Optional[torch.Tensor] = None,
        audio_coords: Optional[torch.Tensor] = None,
        video_self_attention_mask: Optional[torch.Tensor] = None,
        audio_self_attention_mask: Optional[torch.Tensor] = None,
        a2v_cross_attention_mask: Optional[torch.Tensor] = None,
        v2a_cross_attention_mask: Optional[torch.Tensor] = None,
        skip_video_self_attn_blocks: Optional[tuple[int, ...]] = None,
        skip_audio_self_attn_blocks: Optional[tuple[int, ...]] = None,
        disable_a2v_cross_attn: bool = False,
        disable_v2a_cross_attn: bool = False,
        audio_replicated_for_sp: bool = False,
        video_memory_prefix_len: int = 0,
        late_layer_ratio: float = 1.0,
        late_audio_self_attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        batch_size = hidden_states.size(0)
        audio_timestep = audio_timestep if audio_timestep is not None else timestep

        if num_frames is None or height is None or width is None:
            raise ValueError(
                "num_frames/height/width must be provided for RoPE coordinate generation."
            )
        if audio_num_frames is None:
            raise ValueError(
                "audio_num_frames must be provided for RoPE coordinate generation."
            )
        perturbation_configs = kwargs.get("perturbation_configs")
        if perturbation_configs is not None and len(perturbation_configs) != batch_size:
            raise ValueError(
                "perturbation_configs length must match batch size, got "
                f"{len(perturbation_configs)=} {batch_size=}."
            )

        if video_coords is None:
            # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local
            # time shard but RoPE positions must be offset to global time.
            #
            # We assume equal time sharding across SP ranks.
            if model_parallel_is_initialized():
                sp_world_size = get_sp_world_size()
                sp_rank = get_sp_parallel_rank()
            else:
                sp_world_size = 1
                sp_rank = 0

            video_shift = int(sp_rank) * int(num_frames) if sp_world_size > 1 else 0
            video_coords = self.rope.prepare_video_coords(
                batch_size=batch_size,
                num_frames=num_frames,
                height=height,
                width=width,
                device=hidden_states.device,
                fps=fps,
                start_frame=video_shift,
            )
        if audio_coords is None:
            audio_coords = self.audio_rope.prepare_audio_coords(
                batch_size=batch_size,
                num_frames=audio_num_frames,
                device=audio_hidden_states.device,
            )

        video_coords = self._maybe_quantize_video_rope_coords(
            video_coords, hidden_states.device, hidden_states.dtype
        )
        audio_coords = audio_coords.to(device=audio_hidden_states.device)
        video_rotary_emb = self.rope(
            video_coords,
            device=hidden_states.device,
            out_dtype=hidden_states.dtype,
        )
        audio_rotary_emb = self.audio_rope(
            audio_coords,
            device=audio_hidden_states.device,
            out_dtype=audio_hidden_states.dtype,
        )
        ca_video_rotary_emb = self.cross_attn_rope(
            video_coords[:, 0:1, :],
            device=hidden_states.device,
            out_dtype=hidden_states.dtype,
        )
        ca_audio_rotary_emb = self.cross_attn_audio_rope(
            audio_coords[:, 0:1, :],
            device=audio_hidden_states.device,
            out_dtype=audio_hidden_states.dtype,
        )

        # 2. Patchify input projections
        hidden_states, _ = self.patchify_proj(hidden_states)
        audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states)
        # 3. Prepare timestep embeddings
        # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters
        timestep_for_adaln = self._scale_timestep_for_adaln(timestep)
        audio_timestep_for_adaln = self._scale_timestep_for_adaln(audio_timestep)
        temb, embedded_timestep = self.adaln_single(
            timestep_for_adaln.flatten(),
            hidden_dtype=hidden_states.dtype,
        )
        temb = temb.view(batch_size, -1, temb.size(-1))
        embedded_timestep = embedded_timestep.view(
            batch_size, -1, embedded_timestep.size(-1)
        )

        temb_audio, audio_embedded_timestep = self.audio_adaln_single(
            audio_timestep_for_adaln.flatten(),
            hidden_dtype=audio_hidden_states.dtype,
        )
        temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
        audio_embedded_timestep = audio_embedded_timestep.view(
            batch_size, -1, audio_embedded_timestep.size(-1)
        )
        temb_prompt = None
        temb_audio_prompt = None
        if self.prompt_adaln_single is not None:
            prompt_timestep = (
                self._collapse_prompt_timestep(timestep)
                if prompt_timestep is None
                else prompt_timestep
            )
            prompt_timestep_for_adaln = self._scale_timestep_for_adaln(prompt_timestep)
            temb_prompt, _ = self.prompt_adaln_single(
                prompt_timestep_for_adaln.flatten(), hidden_dtype=hidden_states.dtype
            )
            temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1))
        if self.audio_prompt_adaln_single is not None:
            audio_prompt_timestep = (
                self._collapse_prompt_timestep(audio_timestep)
                if audio_prompt_timestep is None
                else audio_prompt_timestep
            )
            audio_prompt_timestep_for_adaln = self._scale_timestep_for_adaln(
                audio_prompt_timestep
            )
            temb_audio_prompt, _ = self.audio_prompt_adaln_single(
                audio_prompt_timestep_for_adaln.flatten(),
                hidden_dtype=audio_hidden_states.dtype,
            )
            temb_audio_prompt = temb_audio_prompt.view(
                batch_size, -1, temb_audio_prompt.size(-1)
            )

        # 3.2. Prepare global modality cross attention modulation parameters
        hidden_dtype = hidden_states.dtype
        av_ca_video_timestep, av_ca_audio_timestep = self._get_av_ca_timesteps(
            timestep,
            audio_timestep,
            prompt_timestep,
            audio_prompt_timestep,
        )
        av_ca_video_timestep_for_adaln = self._scale_timestep_for_adaln(
            av_ca_video_timestep
        )
        av_ca_audio_timestep_for_adaln = self._scale_timestep_for_adaln(
            av_ca_audio_timestep
        )
        temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single(
            av_ca_video_timestep_for_adaln.flatten(), hidden_dtype=hidden_dtype
        )
        temb_ca_scale_shift = temb_ca_scale_shift.view(
            batch_size, -1, temb_ca_scale_shift.shape[-1]
        )

        av_ca_gate_factor = self._get_av_ca_gate_timestep_factor()
        temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single(
            av_ca_video_timestep_for_adaln.flatten() * av_ca_gate_factor,
            hidden_dtype=hidden_dtype,
        )
        temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1])

        temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single(
            av_ca_audio_timestep_for_adaln.flatten(),
            hidden_dtype=audio_hidden_states.dtype,
        )
        temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view(
            batch_size, -1, temb_ca_audio_scale_shift.shape[-1]
        )

        temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single(
            av_ca_audio_timestep_for_adaln.flatten() * av_ca_gate_factor,
            hidden_dtype=audio_hidden_states.dtype,
        )
        temb_ca_audio_gate = temb_ca_audio_gate.view(
            batch_size, -1, temb_ca_audio_gate.shape[-1]
        )

        # 4. Prepare prompt embeddings
        if self.caption_projection is not None:
            encoder_hidden_states = self.caption_projection(encoder_hidden_states)
        if self.audio_caption_projection is not None:
            audio_encoder_hidden_states = self.audio_caption_projection(
                audio_encoder_hidden_states
            )

        if _is_npu:
            # If the 'encoder_attention_mask' is provided and it is all ones,
            # it can be set to 'None' to avoid the degradation of performance on the NPU side,
            # where the mask, even though it has no affect,
            # can lead to the introduction of multiple small operators.
            if encoder_attention_mask is not None and torch.all(
                encoder_attention_mask == 1
            ):
                encoder_attention_mask = None

        # 5. Run blocks
        skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ())
        skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ())
        video_self_attn_perturbation_states = None
        audio_self_attn_perturbation_states = None
        a2v_cross_attn_perturbation_states = None
        v2a_cross_attn_perturbation_states = None
        if perturbation_configs is not None:
            block_indices = tuple(
                getattr(block, "idx", -1) for block in self.transformer_blocks
            )
            video_self_attn_perturbation_states = (
                _ltx2_build_batched_perturbation_states(
                    perturbation_configs,
                    "skip_video_self_attn_blocks",
                    block_indices,
                    hidden_states,
                )
            )
            audio_self_attn_perturbation_states = (
                _ltx2_build_batched_perturbation_states(
                    perturbation_configs,
                    "skip_audio_self_attn_blocks",
                    block_indices,
                    audio_hidden_states,
                )
            )
            a2v_cross_attn_perturbation_states = (
                _ltx2_build_batched_perturbation_states(
                    perturbation_configs,
                    "skip_a2v_cross_attn",
                    block_indices,
                    hidden_states,
                )
            )
            v2a_cross_attn_perturbation_states = (
                _ltx2_build_batched_perturbation_states(
                    perturbation_configs,
                    "skip_v2a_cross_attn",
                    block_indices,
                    audio_hidden_states,
                )
            )
        late_layer_start = int(len(self.transformer_blocks) * float(late_layer_ratio))
        for block in self.transformer_blocks:
            block_idx = getattr(block, "idx", -1)
            video_self_attn_perturbation_mask = None
            audio_self_attn_perturbation_mask = None
            a2v_cross_attn_perturbation_mask = None
            v2a_cross_attn_perturbation_mask = None
            skip_video_self_attn = block_idx in skip_video_self_attn_blocks
            skip_audio_self_attn = block_idx in skip_audio_self_attn_blocks
            skip_a2v_cross_attn = disable_a2v_cross_attn
            skip_v2a_cross_attn = disable_v2a_cross_attn
            block_audio_self_attention_mask = audio_self_attention_mask
            if (
                block_idx >= late_layer_start
                and late_audio_self_attention_mask is not None
            ):
                block_audio_self_attention_mask = late_audio_self_attention_mask
            elif block_idx >= late_layer_start and late_layer_ratio < 1.0:
                block_audio_self_attention_mask = None
            if perturbation_configs is not None:
                if not skip_video_self_attn:
                    assert video_self_attn_perturbation_states is not None
                    state = video_self_attn_perturbation_states[block_idx]
                    video_self_attn_perturbation_mask, skip_video_self_attn = state
                if not skip_audio_self_attn:
                    assert audio_self_attn_perturbation_states is not None
                    state = audio_self_attn_perturbation_states[block_idx]
                    audio_self_attn_perturbation_mask, skip_audio_self_attn = state
                if not skip_a2v_cross_attn:
                    assert a2v_cross_attn_perturbation_states is not None
                    state = a2v_cross_attn_perturbation_states[block_idx]
                    a2v_cross_attn_perturbation_mask, skip_a2v_cross_attn = state
                if not skip_v2a_cross_attn:
                    assert v2a_cross_attn_perturbation_states is not None
                    state = v2a_cross_attn_perturbation_states[block_idx]
                    v2a_cross_attn_perturbation_mask, skip_v2a_cross_attn = state
            hidden_states, audio_hidden_states = block(
                hidden_states,
                audio_hidden_states,
                encoder_hidden_states,
                audio_encoder_hidden_states,
                # Keep the first 4 args positional to stay compatible with cache-dit's
                # LTX2 adapter, which treats `audio_hidden_states` as `encoder_hidden_states`
                # under ForwardPattern.Pattern_0.
                temb=temb,
                temb_audio=temb_audio,
                temb_prompt=temb_prompt,
                temb_audio_prompt=temb_audio_prompt,
                temb_ca_scale_shift=temb_ca_scale_shift,
                temb_ca_audio_scale_shift=temb_ca_audio_scale_shift,
                temb_ca_gate=temb_ca_gate,
                temb_ca_audio_gate=temb_ca_audio_gate,
                video_rotary_emb=video_rotary_emb,
                audio_rotary_emb=audio_rotary_emb,
                ca_video_rotary_emb=ca_video_rotary_emb,
                ca_audio_rotary_emb=ca_audio_rotary_emb,
                encoder_attention_mask=encoder_attention_mask,
                audio_encoder_attention_mask=audio_encoder_attention_mask,
                video_self_attention_mask=video_self_attention_mask,
                audio_self_attention_mask=block_audio_self_attention_mask,
                a2v_cross_attention_mask=a2v_cross_attention_mask,
                v2a_cross_attention_mask=v2a_cross_attention_mask,
                skip_video_self_attn=skip_video_self_attn,
                skip_audio_self_attn=skip_audio_self_attn,
                skip_a2v_cross_attn=skip_a2v_cross_attn,
                skip_v2a_cross_attn=skip_v2a_cross_attn,
                video_self_attn_perturbation_mask=video_self_attn_perturbation_mask,
                audio_self_attn_perturbation_mask=audio_self_attn_perturbation_mask,
                a2v_cross_attn_perturbation_mask=a2v_cross_attn_perturbation_mask,
                v2a_cross_attn_perturbation_mask=v2a_cross_attn_perturbation_mask,
                audio_replicated_for_sp=audio_replicated_for_sp,
                video_memory_prefix_len=video_memory_prefix_len,
            )

        # 6. Output layers
        # Video
        scale_shift_values = self.scale_shift_table[None, None].to(
            device=hidden_states.device, dtype=hidden_states.dtype
        ) + embedded_timestep[:, :, None].to(dtype=hidden_states.dtype)
        shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
            hidden_states = self.norm_out(hidden_states)
        hidden_states = hidden_states * (1 + scale) + shift
        hidden_states, _ = self.proj_out(hidden_states)

        # Audio
        audio_scale_shift_values = self.audio_scale_shift_table[None, None].to(
            device=audio_hidden_states.device, dtype=audio_hidden_states.dtype
        ) + audio_embedded_timestep[:, :, None].to(dtype=audio_hidden_states.dtype)
        audio_shift, audio_scale = (
            audio_scale_shift_values[:, :, 0],
            audio_scale_shift_values[:, :, 1],
        )
        with torch.autocast(device_type=audio_hidden_states.device.type, enabled=False):
            audio_hidden_states = self.audio_norm_out(audio_hidden_states)
        audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift
        audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states)
        # Unpatchify if requested (default True for pipeline compatibility)
        return_latents = kwargs.get("return_latents", True)

        if return_latents:
            # Unpatchify Video
            # [B, N, C_out_raw*patch_vol] -> [B, C_out_raw, T, H, W]
            # Requires num_frames, height, width to be known
            if num_frames is not None and height is not None and width is not None:
                p_t, p_h, p_w = self.patch_size
                post_t, post_h, post_w = num_frames // p_t, height // p_h, width // p_w
                b = batch_size
                hidden_states = hidden_states.reshape(
                    b, post_t, post_h, post_w, self.out_channels_raw, p_t, p_h, p_w
                )
                hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(
                    b, self.out_channels_raw, num_frames, height, width
                )

            # Unpatchify Audio
            # [B, N, C_out] -> [B, C_out, T] (or 4D/5D)
            if audio_num_frames is not None:
                b = batch_size
                # simple reshape for 1D patch
                audio_hidden_states = audio_hidden_states.permute(0, 2, 1)  # [B, C, T]

        return hidden_states, audio_hidden_states


# Backward-compatible alias (older internal name).
LTXModel = LTX2VideoTransformer3DModel
EntryClass = LTX2VideoTransformer3DModel