Trying text_mask 5

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -301,30 +301,37 @@ 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:
+      # DEBUG: Check shapes and values
       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
+      
+      # For CFG, different batch items have different padding patterns (pos vs neg prompts)
+      # We need a per-batch mask, but segment_ids need to be 1D for vmapped_splash
+      # Solution: Use logical OR across batch - a position is valid if ANY batch item needs it
+      # This is safe because zero embeddings at unneeded positions don't affect attention output
       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())
+      kv_mask_any = jnp.max(attention_mask[:, :mask_len], axis=0).astype(jnp.int32)  # (mask_len,)
+      
+      jax.debug.print("[DEBUG _tpu_flash_attention] Using OR across batch, sum: {}", kv_mask_any.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)
-        kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, extra_valid], axis=0)  # (key_seq_len,)
+        kv_mask_any = jnp.concatenate([kv_mask_any, extra_valid], axis=0)  # (key_seq_len,)
       # Pad to kv_padded_len
       if kv_padded_len > key_seq_len:
         padding = jnp.zeros((kv_padded_len - key_seq_len,), dtype=jnp.int32)
-        kv_mask_padded = jnp.concatenate([kv_mask_for_batch, padding], axis=0)  # (kv_padded_len,)
+        kv_mask_padded = jnp.concatenate([kv_mask_any, padding], axis=0)  # (kv_padded_len,)
       else:
-        kv_mask_padded = kv_mask_for_batch
+        kv_mask_padded = kv_mask_any
       # 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)