Trying text_mask 4

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -301,10 +301,17 @@ def wrap_flash_attention(query, key, value):
     # If attention_mask is provided, apply it to kv_segment_ids
     # attention_mask shape: (B, original_kv_seq_len) with 1 for real tokens, 0 for padded
     if attention_mask is not None:
+      jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask shape: {}", attention_mask.shape)
+      jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask[0] sum: {}", attention_mask[0].sum())
+      if attention_mask.shape[0] > 1:
+        jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask[1] sum: {}", attention_mask[1].sum())
+      jax.debug.print("[DEBUG _tpu_flash_attention] key shape: {}, key_seq_len: {}, kv_padded_len: {}", 
+                      key.shape, key_seq_len, kv_padded_len)
       # Take the first item since padding pattern is same across batch (especially with CFG)
       # This keeps kv_segment_ids as (kv_padded_len,) for compatibility with vmapped_splash
       mask_len = min(key_seq_len, attention_mask.shape[1])
       kv_mask_for_batch = attention_mask[0, :mask_len]  # (mask_len,)
+      jax.debug.print("[DEBUG _tpu_flash_attention] Using attention_mask[0], sum: {}", kv_mask_for_batch.sum())
       # If key_seq_len > mask_len, pad the mask with 1s (assume remaining tokens are valid)
       if key_seq_len > mask_len:
         extra_valid = jnp.ones((key_seq_len - mask_len,), dtype=jnp.int32)
@@ -318,6 +325,7 @@ def wrap_flash_attention(query, key, value):
       # Combine with existing kv_segment_ids (which handles block alignment padding)
       # Both are (kv_padded_len,) - element-wise multiplication
       kv_segment_ids = (kv_segment_ids * kv_mask_padded).astype(jnp.int32)
+      jax.debug.print("[DEBUG _tpu_flash_attention] Final kv_segment_ids sum: {}", kv_segment_ids.sum())
 
     segment_ids = splash_attention_kernel.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)