| 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 | def _slice_rope( |
| 23 | pe: tuple[torch.Tensor, torch.Tensor], start: int, end: int |
| 24 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 25 | return pe[0][..., start:end, :], pe[1][..., start:end, :] |
| 26 | |
| 27 | |
| 28 | def update_kv_cache( |
| 29 | cache: dict, |
| 30 | start: int, |
| 31 | k: torch.Tensor, |
| 32 | v: torch.Tensor, |
| 33 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 34 | """Insert one global token range and return the active sink+FIFO window. |
| 35 | |
| 36 | The cache is deliberately inference-only. Repeated denoising forwards for |
| 37 | the same block replace that block in place; the clean refresh forward then |
| 38 | replaces it once more before the next block starts. |
| 39 | """ |
| 40 | if torch.is_grad_enabled(): |
| 41 | raise RuntimeError("causal KV caches are inference-only") |
| 42 | if k.shape != v.shape or k.ndim != 3: |
| 43 | raise ValueError("KV tensors must have matching [batch, tokens, dim] shapes") |
| 44 | |
| 45 | length = int(cache["length"]) |
| 46 | old_positions = cache["positions"][:length] |
| 47 | old_k = cache["k"][:, :length] |
| 48 | old_v = cache["v"][:, :length] |
| 49 | end = start + k.shape[1] |
| 50 | |
| 51 | # A denoising step is a transaction over [start, end): discard a previous |
| 52 | # noisy version of that range while retaining earlier committed history. |
| 53 | keep_old = old_positions < start |
| 54 | positions = torch.cat( |
| 55 | [old_positions[keep_old], torch.arange(start, end, device=k.device)], dim=0 |
| 56 | ) |
| 57 | merged_k = torch.cat([old_k[:, keep_old], k], dim=1) |
| 58 | merged_v = torch.cat([old_v[:, keep_old], v], dim=1) |
| 59 | |
| 60 | local = int(cache.get("local_attn_size", -1)) |
| 61 | sink = int(cache.get("sink_tokens", 0)) |
| 62 | if local >= 0 and positions.numel() > local: |
| 63 | if not 0 <= sink < local: |
| 64 | raise ValueError(f"expected 0 <= sink_tokens < local_attn_size, got {sink}/{local}") |
| 65 | sink_mask = positions < sink |
| 66 | recent_budget = local - int(sink_mask.sum()) |
| 67 | recent_start = max(sink, end - recent_budget) |
| 68 | keep = sink_mask | (positions >= recent_start) |
| 69 | positions = positions[keep] |
| 70 | merged_k = merged_k[:, keep] |
| 71 | merged_v = merged_v[:, keep] |
| 72 | |
| 73 | active = positions.numel() |
| 74 | if active > cache["k"].shape[1]: |
| 75 | raise ValueError(f"KV cache overflow: {active} active tokens exceed capacity {cache['k'].shape[1]}") |
| 76 | cache["k"][:, :active].copy_(merged_k) |
| 77 | cache["v"][:, :active].copy_(merged_v) |
| 78 | cache["positions"][:active].copy_(positions) |
| 79 | cache["length"] = active |
| 80 | return cache["k"][:, :active].clone(), cache["v"][:, :active].clone() |
| 81 | |
| 82 | |
| 83 | class AttentionCallable(Protocol): |
| 84 | def __call__( |
| 85 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 86 | ) -> torch.Tensor: ... |
| 87 | |
| 88 | |
| 89 | class PytorchAttention(AttentionCallable): |
| 90 | def __call__( |
| 91 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 92 | ) -> torch.Tensor: |
| 93 | b, _, dim_head = q.shape |
| 94 | dim_head //= heads |
| 95 | q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v)) |
| 96 | |
| 97 | if mask is not None: |
| 98 | # add a batch dimension if there isn't already one |
| 99 | if mask.ndim == 2: |
| 100 | mask = mask.unsqueeze(0) |
| 101 | # add a heads dimension if there isn't already one |
| 102 | if mask.ndim == 3: |
| 103 | mask = mask.unsqueeze(1) |
| 104 | |
| 105 | out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) |
| 106 | out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) |
| 107 | return out |
| 108 | |
| 109 | |
| 110 | class XFormersAttention(AttentionCallable): |
| 111 | def __call__( |
| 112 | self, |
| 113 | q: torch.Tensor, |
| 114 | k: torch.Tensor, |
| 115 | v: torch.Tensor, |
| 116 | heads: int, |
| 117 | mask: torch.Tensor | None = None, |
| 118 | ) -> torch.Tensor: |
| 119 | if memory_efficient_attention is None: |
| 120 | raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.") |
| 121 | |
| 122 | b, _, dim_head = q.shape |
| 123 | dim_head //= heads |
| 124 | |
| 125 | # xformers expects [B, M, H, K] |
| 126 | q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) |
| 127 | |
| 128 | if mask is not None: |
| 129 | # add a singleton batch dimension |
| 130 | if mask.ndim == 2: |
| 131 | mask = mask.unsqueeze(0) |
| 132 | # add a singleton heads dimension |
| 133 | if mask.ndim == 3: |
| 134 | mask = mask.unsqueeze(1) |
| 135 | # pad to a multiple of 8 |
| 136 | pad = 8 - mask.shape[-1] % 8 |
| 137 | # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk) |
| 138 | # but when using separated heads, the shape has to be (B, H, Nq, Nk) |
| 139 | # in flux, this matrix ends up being over 1GB |
| 140 | # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full) |
| 141 | mask_out = torch.empty( |
| 142 | [mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device |
| 143 | ) |
| 144 | |
| 145 | mask_out[..., : mask.shape[-1]] = mask |
| 146 | # doesn't this remove the padding again?? |
| 147 | mask = mask_out[..., : mask.shape[-1]] |
| 148 | mask = mask.expand(b, heads, -1, -1) |
| 149 | |
| 150 | out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0) |
| 151 | out = out.reshape(b, -1, heads * dim_head) |
| 152 | return out |
| 153 | |
| 154 | |
| 155 | class FlashAttention3(AttentionCallable): |
| 156 | def __call__( |
| 157 | self, |
| 158 | q: torch.Tensor, |
| 159 | k: torch.Tensor, |
| 160 | v: torch.Tensor, |
| 161 | heads: int, |
| 162 | mask: torch.Tensor | None = None, |
| 163 | ) -> torch.Tensor: |
| 164 | if flash_attn_interface is None: |
| 165 | raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.") |
| 166 | |
| 167 | b, _, dim_head = q.shape |
| 168 | dim_head //= heads |
| 169 | |
| 170 | q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) |
| 171 | |
| 172 | if mask is not None: |
| 173 | raise NotImplementedError("Mask is not supported for FlashAttention3") |
| 174 | |
| 175 | out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v) |
| 176 | out = out.reshape(b, -1, heads * dim_head) |
| 177 | return out |
| 178 | |
| 179 | |
| 180 | class AttentionFunction(Enum): |
| 181 | PYTORCH = "pytorch" |
| 182 | XFORMERS = "xformers" |
| 183 | FLASH_ATTENTION_3 = "flash_attention_3" |
| 184 | DEFAULT = "default" |
| 185 | |
| 186 | def __call__( |
| 187 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None |
| 188 | ) -> torch.Tensor: |
| 189 | if self is AttentionFunction.PYTORCH: |
| 190 | return PytorchAttention()(q, k, v, heads, mask) |
| 191 | elif self is AttentionFunction.XFORMERS: |
| 192 | return XFormersAttention()(q, k, v, heads, mask) |
| 193 | elif self is AttentionFunction.FLASH_ATTENTION_3: |
| 194 | return FlashAttention3()(q, k, v, heads, mask) |
| 195 | else: |
| 196 | # Default behavior: XFormers if installed else - PyTorch |
| 197 | return ( |
| 198 | XFormersAttention()(q, k, v, heads, mask) |
| 199 | if memory_efficient_attention is not None |
| 200 | else PytorchAttention()(q, k, v, heads, mask) |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | class Attention(torch.nn.Module): |
| 205 | def __init__( |
| 206 | self, |
| 207 | query_dim: int, |
| 208 | context_dim: int | None = None, |
| 209 | heads: int = 8, |
| 210 | dim_head: int = 64, |
| 211 | norm_eps: float = 1e-6, |
| 212 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 213 | attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT, |
| 214 | apply_gated_attention: bool = False, |
| 215 | ) -> None: |
| 216 | super().__init__() |
| 217 | self.rope_type = rope_type |
| 218 | self.attention_function = attention_function |
| 219 | |
| 220 | inner_dim = dim_head * heads |
| 221 | context_dim = query_dim if context_dim is None else context_dim |
| 222 | |
| 223 | self.heads = heads |
| 224 | self.dim_head = dim_head |
| 225 | |
| 226 | self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) |
| 227 | self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) |
| 228 | |
| 229 | self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) |
| 230 | self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) |
| 231 | self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) |
| 232 | |
| 233 | # Optional per-head gating |
| 234 | if apply_gated_attention: |
| 235 | self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) |
| 236 | else: |
| 237 | self.to_gate_logits = None |
| 238 | |
| 239 | self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) |
| 240 | |
| 241 | def forward( |
| 242 | self, |
| 243 | x: torch.Tensor, |
| 244 | context: torch.Tensor | None = None, |
| 245 | mask: torch.Tensor | None = None, |
| 246 | pe: torch.Tensor | None = None, |
| 247 | k_pe: torch.Tensor | None = None, |
| 248 | perturbation_mask: torch.Tensor | None = None, |
| 249 | all_perturbed: bool = False, |
| 250 | kv_cache: dict | None = None, |
| 251 | kv_cache_start: int = 0, |
| 252 | crossattn_cache: dict | None = None, |
| 253 | ) -> torch.Tensor: |
| 254 | """Multi-head attention with optional RoPE, perturbation masking, and per-head gating. |
| 255 | When ``perturbation_mask`` is all zeros, the expensive query/key path |
| 256 | (linear projections, RMSNorm, RoPE) is skipped entirely and only the |
| 257 | value projection is used as a pass-through. |
| 258 | Args: |
| 259 | x: Query input tensor of shape ``(B, T, query_dim)``. |
| 260 | context: Key/value context tensor of shape ``(B, S, context_dim)``. |
| 261 | Falls back to ``x`` (self-attention) when *None*. |
| 262 | mask: Optional attention mask. Interpretation depends on the attention |
| 263 | backend (additive bias for xformers/PyTorch SDPA). |
| 264 | pe: Rotary positional embeddings applied to both ``q`` and ``k``. |
| 265 | k_pe: Separate rotary positional embeddings for ``k`` only. When |
| 266 | *None*, ``pe`` is reused for keys. |
| 267 | perturbation_mask: Optional mask in ``[0, 1]`` that |
| 268 | blends the attention output with the raw value projection: |
| 269 | ``out = attn_out * mask + v * (1 - mask)``. |
| 270 | **1** keeps the full attention output, **0** bypasses attention |
| 271 | and passes the value projection through unchanged. |
| 272 | *None* or all-ones means standard attention; all-zeros skips |
| 273 | the query/key path entirely for efficiency. |
| 274 | all_perturbed: Whether all perturbations are active for this block. |
| 275 | Returns: |
| 276 | Output tensor of shape ``(B, T, query_dim)``. |
| 277 | """ |
| 278 | context = x if context is None else context |
| 279 | use_attention = not all_perturbed |
| 280 | |
| 281 | v = self.to_v(context) |
| 282 | |
| 283 | if not use_attention: |
| 284 | out = v |
| 285 | else: |
| 286 | q = self.q_norm(self.to_q(x)) |
| 287 | if crossattn_cache is not None: |
| 288 | if not crossattn_cache["is_init"]: |
| 289 | cached_k = self.k_norm(self.to_k(context)) |
| 290 | size = cached_k.shape[1] |
| 291 | crossattn_cache["k"][:, :size].copy_(cached_k) |
| 292 | crossattn_cache["v"][:, :size].copy_(v) |
| 293 | crossattn_cache["length"] = size |
| 294 | crossattn_cache["is_init"] = True |
| 295 | size = int(crossattn_cache["length"]) |
| 296 | k = crossattn_cache["k"][:, :size] |
| 297 | v = crossattn_cache["v"][:, :size] |
| 298 | if pe is not None: |
| 299 | q = apply_rotary_emb(q, pe, self.rope_type) |
| 300 | else: |
| 301 | k = self.k_norm(self.to_k(context)) |
| 302 | local_pe = kv_cache.get("local_rope_pe") if kv_cache is not None else None |
| 303 | local_q_pe = kv_cache.get("local_cross_q_rope_pe") if kv_cache is not None else None |
| 304 | local_k_pe = kv_cache.get("local_cross_k_rope_pe") if kv_cache is not None else None |
| 305 | if local_pe is not None: |
| 306 | k, v = update_kv_cache(kv_cache, kv_cache_start, k, v) |
| 307 | active = k.shape[1] |
| 308 | q_len = q.shape[1] |
| 309 | q = apply_rotary_emb(q, _slice_rope(local_pe, active - q_len, active), self.rope_type) |
| 310 | k = apply_rotary_emb(k, _slice_rope(local_pe, 0, active), self.rope_type) |
| 311 | elif local_q_pe is not None or local_k_pe is not None: |
| 312 | if local_q_pe is None or local_k_pe is None: |
| 313 | raise ValueError("cross-modal RoPE rebase requires both query and key templates") |
| 314 | new_keys = k.shape[1] |
| 315 | k, v = update_kv_cache(kv_cache, kv_cache_start, k, v) |
| 316 | query_slice = kv_cache["local_cross_q_slices"].get((kv_cache_start, kv_cache_start + new_keys)) |
| 317 | if query_slice is None: |
| 318 | raise ValueError("missing local cross-modal query RoPE slice") |
| 319 | q = apply_rotary_emb(q, _slice_rope(local_q_pe, *query_slice), self.rope_type) |
| 320 | k = apply_rotary_emb(k, _slice_rope(local_k_pe, 0, k.shape[1]), self.rope_type) |
| 321 | else: |
| 322 | if pe is not None: |
| 323 | q = apply_rotary_emb(q, pe, self.rope_type) |
| 324 | k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) |
| 325 | if kv_cache is not None: |
| 326 | k, v = update_kv_cache(kv_cache, kv_cache_start, k, v) |
| 327 | |
| 328 | out = self.attention_function(q, k, v, self.heads, mask) # (B, T, H*D) |
| 329 | |
| 330 | if perturbation_mask is not None: |
| 331 | out = out * perturbation_mask + v * (1 - perturbation_mask) |
| 332 | |
| 333 | # Apply per-head gating if enabled |
| 334 | if self.to_gate_logits is not None: |
| 335 | gate_logits = self.to_gate_logits(x) # (B, T, H) |
| 336 | b, t, _ = out.shape |
| 337 | # Reshape to (B, T, H, D) for per-head gating |
| 338 | out = out.view(b, t, self.heads, self.dim_head) |
| 339 | # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0) |
| 340 | gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H) |
| 341 | out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1) |
| 342 | # Reshape back to (B, T, H*D) |
| 343 | out = out.view(b, t, self.heads * self.dim_head) |
| 344 | |
| 345 | return self.to_out(out) |
| 346 |