some debug added to understand key_seq_len

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -285,6 +285,8 @@ def wrap_flash_attention(query, key, value):
 
     block_kv = max(*block_kv_sizes)
     key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv)
+    print("Key seq len")
+    print(key_seq_len)
     value, _, _ = _pad_data_for_flash(value, heads, block_kv)
 
     mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))