added attention_kernel param

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -92,10 +92,12 @@ def __init__(
         precision: jax.lax.Precision = None,
         names_which_can_be_saved: list = [],
         names_which_can_be_offloaded: list = [],
+        attention_kernel: str = "flash",
     ):
         self.dim = dim
         self.norm_eps = norm_eps
         self.norm_elementwise_affine = norm_elementwise_affine
+        self.attention_kernel = attention_kernel
 
         # 1. Self-Attention (video and audio)
         self.norm1 = nnx.RMSNorm(self.dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -109,7 +111,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         self.audio_norm1 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -123,7 +126,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         # 2. Prompt Cross-Attention
@@ -139,7 +143,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         self.audio_norm2 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -154,7 +159,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -170,7 +176,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         self.video_to_audio_norm = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -185,7 +192,8 @@ def __init__(
             out_bias=attention_out_bias,
             eps=norm_eps,
             dtype=dtype,
-            mesh=mesh
+            mesh=mesh,
+            attention_kernel=self.attention_kernel
         )
 
         # 4. Feed Forward
@@ -523,6 +531,7 @@ def __init__(
         names_which_can_be_saved: list = [],
         names_which_can_be_offloaded: list = [],
         scan_layers: bool = True,
+        attention_kernel: str = "flash",
     ):
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -568,6 +577,7 @@ def __init__(
         self.names_which_can_be_saved = names_which_can_be_saved
         self.names_which_can_be_offloaded = names_which_can_be_offloaded
         self.scan_layers = scan_layers
+        self.attention_kernel = attention_kernel
 
         _out_channels = self.out_channels or self.in_channels
         _audio_out_channels = self.audio_out_channels or self.audio_in_channels
@@ -723,6 +733,7 @@ def init_block(rngs):
                 precision=self.precision,
                 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,
             )
 
         if self.scan_layers:
@@ -754,6 +765,7 @@ def init_block(rngs):
                     precision=self.precision,
                     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,
                 )
                 blocks.append(block)
             self.transformer_blocks = nnx.List(blocks)

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -258,6 +258,64 @@ def test_transformer_3d_model_instantiation_and_forward(self):
         self.assertEqual(sample.shape, (self.batch_size, self.seq_len, self.out_channels))
         self.assertEqual(audio_sample.shape, (self.batch_size, 128, self.audio_in_channels))
 
+    def test_transformer_3d_model_dot_product_attention(self):
+        """Verifies LTX2VideoTransformer3DModel full instantiation and forward pass with dot_product attention."""
+        
+        # 1. Instantiate Model with dot_product kernel
+        model = LTX2VideoTransformer3DModel(
+            rngs=nnx.Rngs(0),
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            patch_size=self.patch_size,
+            patch_size_t=self.patch_size_t,
+            num_attention_heads=self.num_attention_heads,
+            attention_head_dim=self.attention_head_dim,
+            cross_attention_dim=self.cross_attention_dim,
+            audio_in_channels=self.audio_in_channels,
+            audio_out_channels=self.audio_out_channels,
+            audio_patch_size=self.audio_patch_size,
+            audio_patch_size_t=self.audio_patch_size_t,
+            audio_num_attention_heads=self.audio_num_attention_heads,
+            audio_attention_head_dim=self.audio_attention_head_dim,
+            audio_cross_attention_dim=self.audio_cross_attention_dim,
+            num_layers=1, # Reduced layers for speed
+            config=self.config,
+            scan_layers=False,
+            mesh=self.mesh,
+            attention_kernel="dot_product"
+        )
+        
+        # 2. Inputs
+        hidden_states = jnp.ones((self.batch_size, self.seq_len, self.in_channels)) * 0.5
+        audio_hidden_states = jnp.ones((self.batch_size, 128, self.audio_in_channels)) * 0.5
+        timestep = jnp.array([1.0]) # (B,)
+        
+        encoder_hidden_states = jnp.zeros((self.batch_size, 128, 32)) # (B, Lc, Dc)
+        audio_encoder_hidden_states = jnp.zeros((self.batch_size, 128, 32))
+        encoder_attention_mask = jnp.ones((self.batch_size, 128), dtype=jnp.float32)
+        audio_encoder_attention_mask = jnp.ones((self.batch_size, 128), dtype=jnp.float32)
+        
+        # Forward
+        with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+            output = model(
+                hidden_states=hidden_states,
+                audio_hidden_states=audio_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                audio_encoder_hidden_states=audio_encoder_hidden_states,
+                timestep=timestep,
+                num_frames=self.num_frames,
+                height=self.height,
+                width=self.width,
+                audio_num_frames=128,
+                fps=24.0,
+                return_dict=True,
+                encoder_attention_mask=encoder_attention_mask,
+                audio_encoder_attention_mask=audio_encoder_attention_mask
+            )
+            
+        self.assertEqual(output.sample.shape, hidden_states.shape)
+        self.assertEqual(output.audio_sample.shape, audio_hidden_states.shape)
+
     def test_scan_remat_parity(self):
         """
         Verifies that scan_layers=True produces identical output to scan_layers=False.