registering a2v and v2a attention

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -2,7 +2,7 @@
 hardware: 'tpu'
 skip_jax_distributed_system: False
 attention: 'ulysses'
-a2v_attention_kernel: 'ulysses'
+a2v_attention_kernel: 'dot_product'
 v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 
 precision: 'bf16'

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -127,6 +127,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
   ltx2_config["dtype"] = config.activations_dtype
   ltx2_config["weights_dtype"] = config.weights_dtype
   ltx2_config["attention_kernel"] = config.attention
+  ltx2_config["a2v_attention_kernel"] = getattr(config, "a2v_attention_kernel", "flash")
+  ltx2_config["v2a_attention_kernel"] = getattr(config, "v2a_attention_kernel", "dot_product")
   ltx2_config["precision"] = get_precision(config)
   ltx2_config["flash_block_sizes"] = get_flash_block_sizes(config)
   ltx2_config["flash_min_seq_length"] = getattr(config, "flash_min_seq_length", 4096)