added enable_profiler flag in yml file, fix in attention_flax.py for kv_max_block_size

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -65,6 +65,7 @@ ici_data_parallelism: 1
 ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
 ici_context_parallelism: 1
 ici_tensor_parallelism: 1
+enable_profiler: False
 
 replicate_vae: False
 

src/maxdiffusion/models/attention_flax.py

@@ -235,8 +235,7 @@ def _tpu_flash_attention(
   q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
   # This is the case for cross-attn.
   if key.shape[1] != query.shape[1]:
-    assert key.shape[1] % 128 == 0
-    kv_max_block_size = key.shape[1]
+    kv_max_block_size = ((key.shape[1] + 127) // 128) * 128
   else:
     kv_max_block_size = q_max_block_size
   # ensure that for cross attention we override the block sizes.