Perf(LTX2): Comprehensive XLA, Memory, and Transformer Code Quality Optimizations

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -1,7 +1,12 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
-attention: 'flash'
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring 
+
+# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
+ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+ulysses_attention_chunks: 1
 a2v_attention_kernel: 'flash'
 v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 

src/maxdiffusion/configs/ltx2_video.yml

@@ -1,7 +1,12 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
-attention: 'flash'
+attention: 'flash'  # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring 
+
+# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
+ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+ulysses_attention_chunks: 1
 a2v_attention_kernel: 'dot_product'
 v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 

src/maxdiffusion/generate_ltx2.py

@@ -237,7 +237,9 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
 
   # Export videos
   for i in range(len(videos)):
-    video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4"
+    model_name = getattr(config, "model_name", "ltx2") or "ltx2"
+    model_name_prefix = model_name.replace(".", "_")
+    video_path = f"{filename_prefix}{model_name_prefix}_output_{getattr(config, 'seed', 0)}_{i}.mp4"
     audio_i = audios[i] if audios is not None else None
 
     audio_format = getattr(config, "audio_format", "s16")

src/maxdiffusion/models/attention_flax.py

@@ -1914,7 +1914,7 @@ def __init__(
       self,
       mesh: Mesh,
       attention_kernel: str,
-      scale: int,
+      scale: float,
       heads: int,
       dim_head: int,
       use_memory_efficient_attention: bool = False,
@@ -2009,7 +2009,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
 class AttentionOp(nn.Module):
   mesh: Mesh
   attention_kernel: str
-  scale: int
+  scale: float
   heads: int
   dim_head: int
   use_memory_efficient_attention: bool = False

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -88,13 +88,8 @@ def apply_split_rotary_emb(x: Array, freqs: Tuple[Array, Array]) -> Array:
   first_x = split_x[..., 0, :]
   second_x = split_x[..., 1, :]
 
-  cos_u = jnp.expand_dims(cos, axis=-2)
-  sin_u = jnp.expand_dims(sin, axis=-2)
-
-  out = split_x * cos_u
-
-  out_first = out[..., 0, :] - second_x * sin_u.squeeze(-2)
-  out_second = out[..., 1, :] + first_x * sin_u.squeeze(-2)
+  out_first = first_x * cos - second_x * sin
+  out_second = second_x * cos + first_x * sin
 
   out = jnp.stack([out_first, out_second], axis=-2)
   out = out.reshape(*out.shape[:-2], last_dim)
@@ -176,12 +171,6 @@ def prepare_video_coords(
     patch_ends = grid + patch_size_delta
 
     # Combine start and end coordinates
-    latent_coords = jnp.stack([grid, patch_ends], axis=-1)  # [3, N_F, N_H, N_W, 2]
-    latent_coords = latent_coords.transpose(1, 2, 3, 0, 4)  # [N_F, N_H, N_W, 3, 2]
-    latent_coords = latent_coords.reshape(-1, 3, 2)  # [num_patches, 3, 2]
-    latent_coords = jnp.expand_dims(latent_coords, 0)  # [1, num_patches, 3, 2]
-    latent_coords = jnp.tile(latent_coords, (batch_size, 1, 1, 1))  # [B, num_patches, 3, 2]
-
     latent_coords = jnp.stack([grid, patch_ends], axis=-1)  # [3, N_F, N_H, N_W, 2]
     latent_coords = latent_coords.reshape(3, -1, 2)  # [3, num_patches, 2]
     latent_coords = jnp.expand_dims(latent_coords, 0)  # [1, 3, num_patches, 2]
@@ -352,6 +341,8 @@ def __init__(
       flash_min_seq_length: int = 4096,
       sharding_specs: Optional[LTX2DiTShardingSpecs] = None,
       gated_attn: bool = False,
+      ulysses_shards: int = -1,
+      ulysses_attention_chunks: int = 1,
   ):
     self.heads = heads
     self.rope_type = rope_type
@@ -445,17 +436,32 @@ def __init__(
           dtype=dtype,
       )
 
+    is_self_attention = context_dim is None
+    if not is_self_attention:
+      if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring"):
+        attention_kernel = "tokamax_flash"  # do not use ring attention for cross attention
+      if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"):
+        attention_kernel = "ulysses_custom"  # plain ulysses (no ring) for cross attention
+
+      axis_names_q = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_KV_LENGTH, common_types.D_KV)
+    else:
+      axis_names_q = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV)
+
     self.attention_op = NNXAttentionOp(
         mesh=mesh,
         attention_kernel=attention_kernel,
         scale=dim_head**-0.5,
         heads=heads,
         dim_head=dim_head,
         dtype=dtype,
-        axis_names_q=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV),
-        axis_names_kv=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV),
+        axis_names_q=axis_names_q,
+        axis_names_kv=axis_names_kv,
         flash_block_sizes=flash_block_sizes,
         flash_min_seq_length=flash_min_seq_length,
+        ulysses_shards=ulysses_shards,
+        ulysses_attention_chunks=ulysses_attention_chunks,
     )
 
   def __call__(
@@ -485,7 +491,7 @@ def __call__(
     # 3. Apply RoPE
     with jax.named_scope("Apply RoPE"):
       if rotary_emb is not None:
-        if hasattr(self, "rope_type") and self.rope_type == "split":
+        if self.rope_type == "split":
           # Split RoPE: passing full freqs [B, H, S, D//2]
           # apply_split_rotary_emb handles reshaping query/key
 

src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py

@@ -20,11 +20,9 @@
 import jax
 from torchax import interop, default_env
 
-# --- Monkeypatch transformers masking_utils to avoid torchax integer tracing bug ---
+import contextlib
 import transformers.masking_utils
 
-_orig_sliding_window_overlay = transformers.masking_utils.sliding_window_overlay
-
 
 def _patched_sliding_window_overlay(sliding_window: int):
   # pylint: disable=unused-argument
@@ -44,6 +42,16 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
   return inner_mask
 
 
+@contextlib.contextmanager
+def patch_sliding_window_overlay():
+  orig = transformers.masking_utils.sliding_window_overlay
+  transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay
+  try:
+    yield
+  finally:
+    transformers.masking_utils.sliding_window_overlay = orig
+
+
 class TorchaxGemma3TextEncoder(interop.JittableModule):
   """
   A jittable Torchax module for wrapping the HuggingFace PyTorch
@@ -57,8 +65,7 @@ def __call__(
       self, input_ids: jax.Array, attention_mask: jax.Array, output_hidden_states: bool = True
   ) -> Tuple[jax.Array, ...]:
     # Dynamically patch transformers.masking_utils only during the duration of this call
-    transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay
-    try:
+    with patch_sliding_window_overlay():
       with default_env():
         input_ids = interop.torch_view(input_ids)
         attention_mask = interop.torch_view(attention_mask)
@@ -72,9 +79,6 @@ def __call__(
             output_hidden_states=output_hidden_states,
         )
       return interop.jax_view(output)
-    finally:
-      # Restore original behavior to prevent side effects on other potential models in same env
-      transformers.masking_utils.sliding_window_overlay = _orig_sliding_window_overlay
 
   @staticmethod
   def _forward_inner(model, input_ids, attention_mask, output_hidden_states=True):