Fixes

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -2,6 +2,8 @@
 hardware: 'tpu'
 skip_jax_distributed_system: False
 attention: 'flash'
+a2v_attention_kernel: 'flash'
+v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 
 precision: 'bf16'
 scan_layers: True

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -106,6 +106,8 @@ def __init__(
       names_which_can_be_saved: list = [],
       names_which_can_be_offloaded: list = [],
       attention_kernel: str = "flash",
+      a2v_attention_kernel: str = "flash",
+      v2a_attention_kernel: str = "dot_product",
       flash_block_sizes: BlockSizes = None,
       flash_min_seq_length: int = 4096,
   ):
@@ -243,7 +245,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="flash",
+        attention_kernel=a2v_attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
         flash_min_seq_length=0,
@@ -270,7 +272,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel=self.attention_kernel,
+        attention_kernel=v2a_attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
         flash_min_seq_length=flash_min_seq_length,
@@ -571,6 +573,8 @@ def __init__(
       names_which_can_be_offloaded: list = [],
       scan_layers: bool = True,
       attention_kernel: str = "flash",
+      a2v_attention_kernel: str = "flash",
+      v2a_attention_kernel: str = "dot_product",
       qk_norm: str = "rms_norm_across_heads",
       flash_block_sizes: BlockSizes = None,
       flash_min_seq_length: int = 4096,
@@ -620,6 +624,8 @@ def __init__(
     self.names_which_can_be_offloaded = names_which_can_be_offloaded
     self.scan_layers = scan_layers
     self.attention_kernel = attention_kernel
+    self.a2v_attention_kernel = a2v_attention_kernel
+    self.v2a_attention_kernel = v2a_attention_kernel
     self.flash_min_seq_length = flash_min_seq_length
 
     _out_channels = self.out_channels or self.in_channels
@@ -813,6 +819,8 @@ def init_block(rngs):
           names_which_can_be_saved=self.names_which_can_be_saved,
           names_which_can_be_offloaded=self.names_which_can_be_offloaded,
           attention_kernel=self.attention_kernel,
+          a2v_attention_kernel=self.a2v_attention_kernel,
+          v2a_attention_kernel=self.v2a_attention_kernel,
           flash_block_sizes=flash_block_sizes,
           flash_min_seq_length=self.flash_min_seq_length,
       )
@@ -846,6 +854,8 @@ def init_block(rngs):
             names_which_can_be_saved=self.names_which_can_be_saved,
             names_which_can_be_offloaded=self.names_which_can_be_offloaded,
             attention_kernel=self.attention_kernel,
+            a2v_attention_kernel=self.a2v_attention_kernel,
+            v2a_attention_kernel=self.v2a_attention_kernel,
             flash_block_sizes=flash_block_sizes,
             flash_min_seq_length=self.flash_min_seq_length,
         )

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1232,22 +1232,27 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
           connectors_graphdef, connectors_state, prompt_embeds_jax, prompt_attention_mask_jax.astype(jnp.bool_)
       )
 
+      video_embeds_sharded = video_embeds
+      audio_embeds_sharded = audio_embeds
+
+      if not self.transformer.scan_layers:
+        activation_axes = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
+        spec = NamedSharding(self.mesh, P(*activation_axes))
+        video_embeds_sharded = jax.device_put(video_embeds, spec)
+        audio_embeds_sharded = jax.device_put(audio_embeds, spec)
+
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
-      for i, t_val in enumerate(timesteps):
+      for i in range(len(timesteps_jax)):
         t = timesteps_jax[i]
 
         # Isolate input sharding to scan_layers=False to avoid affecting the standard path
         latents_jax_sharded = latents_jax
         audio_latents_jax_sharded = audio_latents_jax
-        video_embeds_sharded = video_embeds
-        audio_embeds_sharded = audio_embeds
 
         if not self.transformer.scan_layers:
           activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
           latents_jax_sharded = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names)
           audio_latents_jax_sharded = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names)
-          video_embeds_sharded = jax.lax.with_sharding_constraint(video_embeds, activation_axis_names)
-          audio_embeds_sharded = jax.lax.with_sharding_constraint(audio_embeds, activation_axis_names)
 
         noise_pred, noise_pred_audio = transformer_forward_pass(
             graphdef,

src/maxdiffusion/schedulers/scheduling_flow_match_flax.py

@@ -244,6 +244,7 @@ def _find_timestep_id(self, state: FlowMatchSchedulerState, timestep: jnp.ndarra
       diffs = jnp.abs(state.timesteps[None, :] - timestep[:, None])
       return jnp.argmin(diffs, axis=1)
 
+  # Arguments at indices 0 (self), 5 (to_final), and 6 (return_dict) are kept static for JIT compilation.
   @partial(jax.jit, static_argnums=(0, 5, 6))
   def step(
       self,