ltx2 ¶

def get_patch_grid_bounds(
    self,
    output_shape: AudioLatentShape,
    device: Optional[torch.device] = None,
) -> torch.Tensor:
    """Get patch grid bounds for audio RoPE computation.

    Args:
        output_shape: Shape of the audio latent tensor
        device: Device to create tensors on
    """
    start_timings = self._get_audio_latent_time_in_sec(
        self.shift,
        output_shape.frames + self.shift,
        torch.float32,
        device,
    )
    start_timings = start_timings.unsqueeze(0).expand(output_shape.batch,
                                                      -1).unsqueeze(1)

    end_timings = self._get_audio_latent_time_in_sec(
        self.shift + 1,
        output_shape.frames + self.shift + 1,
        torch.float32,
        device,
    )
    end_timings = end_timings.unsqueeze(0).expand(output_shape.batch,
                                                  -1).unsqueeze(1)

    return torch.stack([start_timings, end_timings], dim=-1)

def forward(
    self,
    video: TransformerArgs | None,
    audio: TransformerArgs | None,
    video_original_seq_len: int | None = None,
    audio_original_seq_len: int | None = None,
    skip_cross_modal_attn: bool = False,
    skip_video_self_attn: bool | torch.Tensor = False,
    skip_audio_self_attn: bool | torch.Tensor = False,
) -> tuple[TransformerArgs | None, TransformerArgs | None]:
    """Forward pass for transformer block.

    Args:
        video: Video transformer args
        audio: Audio transformer args
        video_original_seq_len: Original unpadded video sequence length
        audio_original_seq_len: Original unpadded audio sequence length
        skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
            attention (used for the modality-isolated CFG pass).
        skip_video_self_attn: If True, skip video self-attention
            in this block (used for STG perturbed pass).
        skip_audio_self_attn: If True, skip audio self-attention
            in this block (used for STG perturbed pass).
    """
    vx = video.x if video is not None else None
    ax = audio.x if audio is not None else None

    run_vx = video is not None and video.enabled and vx.numel() > 0
    run_ax = audio is not None and audio.enabled and ax.numel() > 0
    run_a2v = run_vx and run_ax
    run_v2a = run_ax and run_vx

    def _build_attn_keep_mask(
        values: torch.Tensor,
        skip_flag: bool | torch.Tensor,
    ) -> torch.Tensor:
        """STG keep-mask for self-attention.

        Returns a per-sample multiplier (1 keep, 0 perturb) broadcastable
        over ``values``. Only the configured ``stg_block_idx`` is perturbed.
        Supports both the LTX-2.0 bool flag and an LTX-2.3 per-sample
        tensor ``[B]`` (1/True == perturb). When skip_flag is False/0 the
        multiplier is all ones, reproducing the un-skipped LTX-2.0 path.
        """
        bsz = values.shape[0]
        keep = torch.ones((bsz, ), device=values.device, dtype=values.dtype)
        if self.idx == self.stg_block_idx:
            if torch.is_tensor(skip_flag):
                if skip_flag.ndim == 0:
                    perturb = skip_flag.reshape(1).expand(bsz)
                else:
                    if skip_flag.shape[0] != bsz:
                        raise ValueError(
                            "Per-sample STG mask batch size mismatch: "
                            f"got {skip_flag.shape[0]}, expected {bsz}")
                    perturb = skip_flag.reshape(bsz)
                keep = 1.0 - perturb.to(device=values.device,
                                        dtype=values.dtype)
            elif bool(skip_flag):
                keep.zero_()
        return keep.view(bsz, *([1] * (values.ndim - 1)))

    if run_vx:
        vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
            self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
        )
        norm_vx = _rms_norm_dispatch(vx, eps=self.norm_eps) * (
            1 + vscale_msa) + vshift_msa
        video_self_attn_mask = _build_attn_keep_mask(vx, skip_video_self_attn)
        if self.use_distributed_attention:
            vx_attn1_out = self.attn1(
                norm_vx,
                pe=video.positional_embeddings,
                original_seq_len=video_original_seq_len,
            )
        else:
            vx_attn1_out = self.attn1(norm_vx, pe=video.positional_embeddings)
        vx = vx + vx_attn1_out * vgate_msa * video_self_attn_mask
        # Text cross-attention: no SP mask needed (text is replicated).
        vx = vx + self._apply_text_cross_attention(
            x=vx,
            context=video.context,
            attn=self.attn2,
            scale_shift_table=self.scale_shift_table,
            prompt_scale_shift_table=getattr(
                self, "prompt_scale_shift_table", None),
            timestep=video.timesteps,
            prompt_timestep=video.prompt_timestep,
            context_mask=video.context_mask,
        )

    if run_ax:
        ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
            self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
        )
        norm_ax = _rms_norm_dispatch(ax, eps=self.norm_eps) * (
            1 + ascale_msa) + ashift_msa
        audio_self_attn_mask = _build_attn_keep_mask(ax, skip_audio_self_attn)
        if self.use_distributed_attention:
            ax_attn1_out = self.audio_attn1(
                norm_ax,
                pe=audio.positional_embeddings,
                original_seq_len=audio_original_seq_len,
            )
        else:
            ax_attn1_out = self.audio_attn1(norm_ax, pe=audio.positional_embeddings)
        ax = ax + ax_attn1_out * agate_msa * audio_self_attn_mask
        # Text cross-attention: no SP mask needed (text is replicated).
        ax = ax + self._apply_text_cross_attention(
            x=ax,
            context=audio.context,
            attn=self.audio_attn2,
            scale_shift_table=self.audio_scale_shift_table,
            prompt_scale_shift_table=getattr(
                self, "audio_prompt_scale_shift_table", None),
            timestep=audio.timesteps,
            prompt_timestep=audio.prompt_timestep,
            context_mask=audio.context_mask,
        )

    if (run_a2v or run_v2a) and not skip_cross_modal_attn:
        vx_norm3 = _rms_norm_dispatch(vx, eps=self.norm_eps)
        ax_norm3 = _rms_norm_dispatch(ax, eps=self.norm_eps)

        (
            scale_ca_audio_hidden_states_a2v,
            shift_ca_audio_hidden_states_a2v,
            scale_ca_audio_hidden_states_v2a,
            shift_ca_audio_hidden_states_v2a,
            gate_out_v2a,
        ) = self.get_av_ca_ada_values(
            self.scale_shift_table_a2v_ca_audio,
            ax.shape[0],
            audio.cross_scale_shift_timestep,
            audio.cross_gate_timestep,
        )

        (
            scale_ca_video_hidden_states_a2v,
            shift_ca_video_hidden_states_a2v,
            scale_ca_video_hidden_states_v2a,
            shift_ca_video_hidden_states_v2a,
            gate_out_a2v,
        ) = self.get_av_ca_ada_values(
            self.scale_shift_table_a2v_ca_video,
            vx.shape[0],
            video.cross_scale_shift_timestep,
            video.cross_gate_timestep,
        )

        if run_a2v:
            vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
            ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
            # Audio-to-video cross-attention: Q from video (sharded), K/V from audio
            # For SP, gather audio context from all ranks before cross-attention
            if self.use_distributed_attention:
                # Gather audio context from all SP ranks
                ax_context = sequence_model_parallel_all_gather(ax_scaled, dim=1)
                # Video Q: video tokens divide evenly (adjusted upstream), so simple slicing works
                video_q_pe = None
                if video.cross_positional_embeddings is not None:
                    sp_rank = get_sp_parallel_rank()
                    local_seq_len = vx_scaled.shape[1]
                    pe_tuple = video.cross_positional_embeddings
                    start_idx = sp_rank * local_seq_len
                    end_idx = start_idx + local_seq_len
                    video_q_pe = tuple(pe[:, :, start_idx:end_idx, :] for pe in pe_tuple)
                # Audio K: gathered context may have padding, trim to original length
                audio_k_pe = audio.cross_positional_embeddings
                if audio_k_pe is not None:
                    original_audio_len = audio_k_pe[0].shape[2]
                    ax_context = ax_context[:, :original_audio_len, :]
                vx = vx + (
                    self.audio_to_video_attn(
                        vx_scaled,
                        context=ax_context,
                        pe=video_q_pe,
                        k_pe=audio_k_pe,
                    )
                    * gate_out_a2v
                )
            else:
                ax_context = ax_scaled
                video_q_pe = video.cross_positional_embeddings
                audio_k_pe = audio.cross_positional_embeddings
                vx = vx + (
                    self.audio_to_video_attn(
                        vx_scaled,
                        context=ax_context,
                        pe=video_q_pe,
                        k_pe=audio_k_pe,
                    )
                    * gate_out_a2v
                )

        if run_v2a:
            ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
            vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
            # Video-to-audio cross-attention: Q from audio, K/V from video
            # For SP, gather both modalities and run cross-attention on full sequences
            # (audio is short, so no need to keep it sharded for cross-attention)
            if self.use_distributed_attention:
                # Gather both modalities to full sequences
                ax_full = sequence_model_parallel_all_gather(ax_scaled, dim=1)
                vx_context = sequence_model_parallel_all_gather(vx_scaled, dim=1)
                # Use full positional embeddings directly
                audio_q_pe = audio.cross_positional_embeddings
                video_k_pe = video.cross_positional_embeddings
                # Trim gathered tensors to original lengths (may have padding from sharding)
                if audio_q_pe is not None:
                    original_audio_len = audio_q_pe[0].shape[2]
                    ax_full = ax_full[:, :original_audio_len, :]
                if video_k_pe is not None:
                    original_video_len = video_k_pe[0].shape[2]
                    vx_context = vx_context[:, :original_video_len, :]
                # Run cross-attention on full audio
                v2a_out_full = self.video_to_audio_attn(
                    ax_full,
                    context=vx_context,
                    pe=audio_q_pe,
                    k_pe=video_k_pe,
                )
                # Shard the output back to match ax's sharding
                v2a_out, _ = sequence_model_parallel_shard(v2a_out_full, dim=1)
                ax = ax + v2a_out * gate_out_v2a
            else:
                vx_context = vx_scaled
                audio_q_pe = audio.cross_positional_embeddings
                video_k_pe = video.cross_positional_embeddings
                ax = ax + (
                    self.video_to_audio_attn(
                        ax_scaled,
                        context=vx_context,
                        pe=audio_q_pe,
                        k_pe=video_k_pe,
                    )
                    * gate_out_v2a
                )

    if run_vx:
        # slice(3, 6) selects the FFN shift/scale/gate. Identical to
        # slice(3, None) for the 6-row LTX-2.0 table, but correct for the
        # 9-row cross_attention_adaln table (rows 6..9 are the cross-attn).
        vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
            self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
        )
        vx_scaled = _rms_norm_dispatch(vx, eps=self.norm_eps) * (
            1 + vscale_mlp) + vshift_mlp
        vx = vx + self.ff(vx_scaled) * vgate_mlp

    if run_ax:
        ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
            self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
        )
        ax_scaled = _rms_norm_dispatch(ax, eps=self.norm_eps) * (
            1 + ascale_mlp) + ashift_mlp
        ax = ax + self.audio_ff(ax_scaled) * agate_mlp

    if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
        video_sum = vx.float().sum().item() if vx is not None else 0.0
        audio_sum = ax.float().sum().item() if ax is not None else 0.0
        _debug_block_log_line(
            f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} "
            f"audio_sum={audio_sum:.6f}"
        )

    # Register FSDP2 backward hooks on output tensors (module-level hooks don't
    # fire for dataclass outputs, so we must hook the tensors directly)
    self._register_fsdp_backward_hooks_on_output(vx, ax)

    return (
        replace(video, x=vx) if video is not None else None,
        replace(audio, x=ax) if audio is not None else None,
    )

def forward(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    original_seq_len: int,
    replicated_q: torch.Tensor | None = None,
    replicated_k: torch.Tensor | None = None,
    replicated_v: torch.Tensor | None = None,
    gate_compress: torch.Tensor | None = None,
    ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Forward pass with LTX-2 style RoPE application.

    Args:
        q: Query tensor [batch_size, seq_len, num_heads, head_dim]
        k: Key tensor [batch_size, seq_len, num_heads, head_dim]
        v: Value tensor [batch_size, seq_len, num_heads, head_dim]
        replicated_q: Replicated query tensor for text tokens
        replicated_k: Replicated key tensor
        replicated_v: Replicated value tensor
        original_seq_len: Original (unpadded) full sequence length
        ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]

    Returns:
        Tuple of output tensor and optional replicated output
    """
    assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
    batch_size, _, num_heads, _ = q.shape
    local_rank = get_sp_parallel_rank()
    world_size = get_sp_world_size()

    forward_context: ForwardContext = get_forward_context()
    ctx_attn_metadata = forward_context.attn_metadata

    use_vsa = self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN

    if use_vsa:
        assert (
            replicated_q is None and replicated_k is None
            and replicated_v is None
        ), "Replicated QKV is not supported for VSA now"
        if gate_compress is None:
            raise ValueError(
                "gate_compress must be provided when using VIDEO_SPARSE_ATTN"
            )

        qkvg = torch.cat([q, k, v, gate_compress],
                         dim=0)  # [4*batch, seq_len, heads, head_dim]
        qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
                                                     scatter_dim=2,
                                                     gather_dim=1)
        pad_seq_len = qkvg.shape[1] - original_seq_len
        qkvg = qkvg[:, :original_seq_len, :, :]

        if ltx_freqs_cis is not None:
            cos, sin = ltx_freqs_cis
            heads_per_rank = num_heads // world_size
            head_start = local_rank * heads_per_rank
            head_end = head_start + heads_per_rank
            cos_local = cos[:, head_start:head_end, :original_seq_len, :]
            sin_local = sin[:, head_start:head_end, :original_seq_len, :]
            qk_part = qkvg[:batch_size * 2].transpose(1, 2)
            qk_part = apply_ltx_rotary_emb_4d(
                qk_part, (cos_local, sin_local), self.rope_type)
            qkvg[:batch_size * 2] = qk_part.transpose(1, 2)

        qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata)
        q, k, v, gate_compress = qkvg.chunk(4, dim=0)
        output = self.attn_impl.forward(  # type: ignore[call-arg]
            q, k, v, gate_compress, ctx_attn_metadata)
        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

        output = sequence_model_parallel_all_to_all_4D(output,
                                                       scatter_dim=1,
                                                       gather_dim=2)
        return output, None

    # Stack QKV
    qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]

    # Redistribute heads across sequence dimension
    qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)

    # After all-to-all, each rank has the full sequence but only a subset of heads
    pad_seq_len = qkv.shape[1] - original_seq_len
    qkv = qkv[:, :original_seq_len, :, :]

    # Apply LTX-2 style RoPE after all-to-all (when we have full sequence)
    if ltx_freqs_cis is not None:
        cos, sin = ltx_freqs_cis
        heads_per_rank = num_heads // world_size
        head_start = local_rank * heads_per_rank
        head_end = head_start + heads_per_rank
        # Slice to this rank's heads: [B, H/SP, T, D]
        cos_local = cos[:, head_start:head_end, :original_seq_len, :]
        sin_local = sin[:, head_start:head_end, :original_seq_len, :]

        # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
        qk_part = qkv[:batch_size * 2]
        # Transpose to [2*B, H/SP, T, D] for LTX RoPE application
        qk_part = qk_part.transpose(1, 2)
        qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type)
        # Transpose back to [2*B, T, H/SP, D]
        qkv[:batch_size * 2] = qk_part.transpose(1, 2)

    # Apply backend-specific preprocess_qkv
    qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)

    # Concatenate with replicated QKV if provided
    if replicated_q is not None:
        assert replicated_k is not None and replicated_v is not None
        replicated_qkv = torch.cat(
            [replicated_q, replicated_k, replicated_v],
            dim=0)  # [3, seq_len, num_heads, head_dim]
        heads_per_rank = num_heads // world_size
        replicated_qkv = replicated_qkv[:, :, local_rank *
                                        heads_per_rank:(local_rank + 1) *
                                        heads_per_rank]
        qkv = torch.cat([qkv, replicated_qkv], dim=1)

    q, k, v = qkv.chunk(3, dim=0)

    output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)

    # Redistribute back if using sequence parallelism
    replicated_output = None
    if replicated_q is not None:
        split_idx = original_seq_len
        replicated_output = output[:, split_idx:]
        output = output[:, :split_idx]
        replicated_output = sequence_model_parallel_all_gather(
            replicated_output.contiguous(), dim=2)

    # Apply backend-specific postprocess_output
    output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)

    output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

    output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)

    return output, replicated_output