trying with flash attn for audio

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -159,7 +159,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -212,7 +212,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -239,7 +239,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -265,7 +265,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )