debug

Author: prishajain1

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -61,7 +61,7 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring
-flash_min_seq_length: 4096
+flash_min_seq_length: 0
 dropout: 0.1
 
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -61,7 +61,7 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring
-flash_min_seq_length: 4096
+flash_min_seq_length: 0
 dropout: 0.1
 
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.

src/maxdiffusion/models/attention_flax.py

@@ -225,6 +225,7 @@ def _tpu_flash_attention(
     attention_mask: jax.Array = None,
 ) -> jax.Array:
   """TPU Flash Attention"""
+  jax.debug.print("USing FLASH ATTENTION")
 
   q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
   # This is the case for cross-attn.
@@ -444,6 +445,7 @@ def _apply_attention_dot(
     float32_qk_product: bool,
     use_memory_efficient_attention: bool,
 ):
+  jax.debug.print("Using DOT PRODUCT ATTENTION")
   """Apply Attention."""
   if split_head_dim:
     b = key.shape[0]