| 1 | from enum import Enum |
| 2 | from typing import Protocol |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb |
| 7 | |
| 8 | memory_efficient_attention = None |
| 9 | flash_attn_interface = None |
| 10 | try: |
| 11 | from xformers.ops import memory_efficient_attention |
| 12 | except ImportError: |
| 13 | memory_efficient_attention = None |
| 14 | try: |
| 15 | # FlashAttention3 and XFormersAttention cannot be used together |
| 16 | if memory_efficient_attention is None: |
| 17 | import flash_attn_interface |
| 18 | except ImportError: |
| 19 | flash_attn_interface = None |
| 20 | |
| 21 | |
| 22 | class AttentionCallable(Protocol): |
| 23 | def __call__( |
| 24 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 25 | ) -> torch.Tensor: ... |
| 26 | |
| 27 | |
| 28 | class PytorchAttention(AttentionCallable): |
| 29 | def __call__( |
| 30 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 31 | ) -> torch.Tensor: |
| 32 | b, _, dim_head = q.shape |
| 33 | dim_head //= heads |
| 34 | q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v)) |
| 35 | |
| 36 | if mask is not None: |
| 37 | # add a batch dimension if there isn't already one |
| 38 | if mask.ndim == 2: |
| 39 | mask = mask.unsqueeze(0) |
| 40 | # add a heads dimension if there isn't already one |
| 41 | if mask.ndim == 3: |
| 42 | mask = mask.unsqueeze(1) |
| 43 | |
| 44 | out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) |
| 45 | out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) |
| 46 | return out |
| 47 | |
| 48 | |
| 49 | class XFormersAttention(AttentionCallable): |
| 50 | def __call__( |
| 51 | self, |
| 52 | q: torch.Tensor, |
| 53 | k: torch.Tensor, |
| 54 | v: torch.Tensor, |
| 55 | heads: int, |
| 56 | mask: torch.Tensor | None = None, |
| 57 | ) -> torch.Tensor: |
| 58 | if memory_efficient_attention is None: |
| 59 | raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.") |
| 60 | |
| 61 | b, _, dim_head = q.shape |
| 62 | dim_head //= heads |
| 63 | |
| 64 | # xformers expects [B, M, H, K] |
| 65 | q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) |
| 66 | |
| 67 | if mask is not None: |
| 68 | # add a singleton batch dimension |
| 69 | if mask.ndim == 2: |
| 70 | mask = mask.unsqueeze(0) |
| 71 | # add a singleton heads dimension |
| 72 | if mask.ndim == 3: |
| 73 | mask = mask.unsqueeze(1) |
| 74 | # pad to a multiple of 8 |
| 75 | pad = 8 - mask.shape[-1] % 8 |
| 76 | # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk) |
| 77 | # but when using separated heads, the shape has to be (B, H, Nq, Nk) |
| 78 | # in flux, this matrix ends up being over 1GB |
| 79 | # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full) |
| 80 | mask_out = torch.empty( |
| 81 | [mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device |
| 82 | ) |
| 83 | |
| 84 | mask_out[..., : mask.shape[-1]] = mask |
| 85 | # doesn't this remove the padding again?? |
| 86 | mask = mask_out[..., : mask.shape[-1]] |
| 87 | mask = mask.expand(b, heads, -1, -1) |
| 88 | |
| 89 | out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0) |
| 90 | out = out.reshape(b, -1, heads * dim_head) |
| 91 | return out |
| 92 | |
| 93 | |
| 94 | class FlashAttention3(AttentionCallable): |
| 95 | def __call__( |
| 96 | self, |
| 97 | q: torch.Tensor, |
| 98 | k: torch.Tensor, |
| 99 | v: torch.Tensor, |
| 100 | heads: int, |
| 101 | mask: torch.Tensor | None = None, |
| 102 | ) -> torch.Tensor: |
| 103 | if flash_attn_interface is None: |
| 104 | raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.") |
| 105 | |
| 106 | b, _, dim_head = q.shape |
| 107 | dim_head //= heads |
| 108 | |
| 109 | q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) |
| 110 | |
| 111 | if mask is not None: |
| 112 | raise NotImplementedError("Mask is not supported for FlashAttention3") |
| 113 | |
| 114 | out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v) |
| 115 | out = out.reshape(b, -1, heads * dim_head) |
| 116 | return out |
| 117 | |
| 118 | |
| 119 | class AttentionFunction(Enum): |
| 120 | PYTORCH = "pytorch" |
| 121 | XFORMERS = "xformers" |
| 122 | FLASH_ATTENTION_3 = "flash_attention_3" |
| 123 | DEFAULT = "default" |
| 124 | |
| 125 | def __call__( |
| 126 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 127 | ) -> torch.Tensor: |
| 128 | if self is AttentionFunction.PYTORCH: |
| 129 | return PytorchAttention()(q, k, v, heads, mask) |
| 130 | elif self is AttentionFunction.XFORMERS: |
| 131 | return XFormersAttention()(q, k, v, heads, mask) |
| 132 | elif self is AttentionFunction.FLASH_ATTENTION_3: |
| 133 | return FlashAttention3()(q, k, v, heads, mask) |
| 134 | else: |
| 135 | # Default behavior: XFormers if installed else - PyTorch |
| 136 | return ( |
| 137 | XFormersAttention()(q, k, v, heads, mask) |
| 138 | if memory_efficient_attention is not None |
| 139 | else PytorchAttention()(q, k, v, heads, mask) |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | class Attention(torch.nn.Module): |
| 144 | def __init__( |
| 145 | self, |
| 146 | query_dim: int, |
| 147 | context_dim: int | None = None, |
| 148 | heads: int = 8, |
| 149 | dim_head: int = 64, |
| 150 | norm_eps: float = 1e-6, |
| 151 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 152 | attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT, |
| 153 | apply_gated_attention: bool = False, |
| 154 | ) -> None: |
| 155 | super().__init__() |
| 156 | self.rope_type = rope_type |
| 157 | self.attention_function = attention_function |
| 158 | |
| 159 | inner_dim = dim_head * heads |
| 160 | context_dim = query_dim if context_dim is None else context_dim |
| 161 | |
| 162 | self.heads = heads |
| 163 | self.dim_head = dim_head |
| 164 | |
| 165 | self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) |
| 166 | self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) |
| 167 | |
| 168 | self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) |
| 169 | self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) |
| 170 | self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) |
| 171 | |
| 172 | # Optional per-head gating |
| 173 | if apply_gated_attention: |
| 174 | self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) |
| 175 | else: |
| 176 | self.to_gate_logits = None |
| 177 | |
| 178 | self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) |
| 179 | |
| 180 | def forward( |
| 181 | self, |
| 182 | x: torch.Tensor, |
| 183 | context: torch.Tensor | None = None, |
| 184 | mask: torch.Tensor | None = None, |
| 185 | pe: torch.Tensor | None = None, |
| 186 | k_pe: torch.Tensor | None = None, |
| 187 | perturbation_mask: torch.Tensor | None = None, |
| 188 | all_perturbed: bool = False, |
| 189 | ) -> torch.Tensor: |
| 190 | """Multi-head attention with optional RoPE, perturbation masking, and per-head gating. |
| 191 | When ``perturbation_mask`` is all zeros, the expensive query/key path |
| 192 | (linear projections, RMSNorm, RoPE) is skipped entirely and only the |
| 193 | value projection is used as a pass-through. |
| 194 | Args: |
| 195 | x: Query input tensor of shape ``(B, T, query_dim)``. |
| 196 | context: Key/value context tensor of shape ``(B, S, context_dim)``. |
| 197 | Falls back to ``x`` (self-attention) when *None*. |
| 198 | mask: Optional attention mask. Interpretation depends on the attention |
| 199 | backend (additive bias for xformers/PyTorch SDPA). |
| 200 | pe: Rotary positional embeddings applied to both ``q`` and ``k``. |
| 201 | k_pe: Separate rotary positional embeddings for ``k`` only. When |
| 202 | *None*, ``pe`` is reused for keys. |
| 203 | perturbation_mask: Optional mask in ``[0, 1]`` that |
| 204 | blends the attention output with the raw value projection: |
| 205 | ``out = attn_out * mask + v * (1 - mask)``. |
| 206 | **1** keeps the full attention output, **0** bypasses attention |
| 207 | and passes the value projection through unchanged. |
| 208 | *None* or all-ones means standard attention; all-zeros skips |
| 209 | the query/key path entirely for efficiency. |
| 210 | all_perturbed: Whether all perturbations are active for this block. |
| 211 | Returns: |
| 212 | Output tensor of shape ``(B, T, query_dim)``. |
| 213 | """ |
| 214 | context = x if context is None else context |
| 215 | use_attention = not all_perturbed |
| 216 | |
| 217 | v = self.to_v(context) |
| 218 | |
| 219 | if not use_attention: |
| 220 | out = v |
| 221 | else: |
| 222 | q = self.to_q(x) |
| 223 | k = self.to_k(context) |
| 224 | |
| 225 | q = self.q_norm(q) |
| 226 | k = self.k_norm(k) |
| 227 | |
| 228 | if pe is not None: |
| 229 | q = apply_rotary_emb(q, pe, self.rope_type) |
| 230 | k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) |
| 231 | |
| 232 | out = self.attention_function(q, k, v, self.heads, mask) # (B, T, H*D) |
| 233 | |
| 234 | if perturbation_mask is not None: |
| 235 | out = out * perturbation_mask + v * (1 - perturbation_mask) |
| 236 | |
| 237 | # Apply per-head gating if enabled |
| 238 | if self.to_gate_logits is not None: |
| 239 | gate_logits = self.to_gate_logits(x) # (B, T, H) |
| 240 | b, t, _ = out.shape |
| 241 | # Reshape to (B, T, H, D) for per-head gating |
| 242 | out = out.view(b, t, self.heads, self.dim_head) |
| 243 | # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0) |
| 244 | gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H) |
| 245 | out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1) |
| 246 | # Reshape back to (B, T, H*D) |
| 247 | out = out.view(b, t, self.heads * self.dim_head) |
| 248 | |
| 249 | return self.to_out(out) |
| 250 |