Merge branch 'main' into revert-for-preview

Author: coolkp

docs/attention_blocks_flowchart.md

@@ -501,6 +501,7 @@ def get_flash_block_sizes(config):
   """Create custom flash attention BlockSizes."""
   flash_block_sizes = None
   if len(config.flash_block_sizes.keys()) > 0:
+    use_fused_bwd_kernel = config.flash_block_sizes.get("use_fused_bwd_kernel", False)
     flash_block_sizes = splash_attention_kernel.BlockSizes(
         block_q=config.flash_block_sizes["block_q"],
         block_kv_compute=config.flash_block_sizes["block_kv_compute"],

docs/attention_blocks_flowchart.png

@@ -189,16 +189,16 @@ def _tpu_flash_attention(
   if flash_block_sizes:
     block_sizes = flash_block_sizes
   else:
-    block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
     block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=block_size_q,
+        block_q=min(q_max_block_size, query.shape[2]),
         block_kv_compute=min(kv_max_block_size, key.shape[2]),
         block_kv=min(kv_max_block_size, key.shape[2]),
-        block_q_dkv=block_size_q,
+        block_q_dkv=min(q_max_block_size, query.shape[2]),
         block_kv_dkv=min(kv_max_block_size, key.shape[2]),
         block_kv_dkv_compute=min(kv_max_block_size, query.shape[2]),
-        block_q_dq=min(q_max_block_size, query.shape[2]),
-        block_kv_dq=min(kv_max_block_size, query.shape[2]),
+        block_q_dq=None if attention_kernel == "tokamax_flash" else block_sizes.block_q_dq,
+        block_kv_dq=None if attention_kernel == "tokamax_flash" else min(kv_max_block_size, query.shape[2]),
+        use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
     )
   num_fsdp_shards = mesh.shape["fsdp"]
   query = _reshape_data_for_flash(query, heads)

src/maxdiffusion/max_utils.py

@@ -181,25 +181,24 @@ def test_wan_block(self):
     assert dummy_output.shape == dummy_hidden_states.shape
 
   def test_wan_attention(self):
-    for attention_kernel in ["flash", "tokamax_flash"]:
-      pyconfig.initialize(
-          [
-              None,
-              os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
-              f"attention={attention_kernel}"
-          ],
-          unittest=True
-      )
-      config = pyconfig.config
-      batch_size = 1
-      channels = 16
-      frames = 21
-      height = 90
-      width = 160
-      hidden_states_shape = (batch_size, frames, height, width, channels)
-      dummy_hidden_states = jnp.ones(hidden_states_shape)
-      wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=128, patch_size=[1, 2, 2], max_seq_len=1024)
-      dummy_rotary_emb = wan_rot_embed(dummy_hidden_states)
+    pyconfig.initialize(
+        [
+            None,
+            os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
+        ],
+        unittest=True,
+    )
+    config = pyconfig.config
+
+    batch_size = 1
+    channels = 16
+    frames = 21
+    height = 90
+    width = 160
+    hidden_states_shape = (batch_size, frames, height, width, channels)
+    dummy_hidden_states = jnp.ones(hidden_states_shape)
+    wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=128, patch_size=[1, 2, 2], max_seq_len=1024)
+    dummy_rotary_emb = wan_rot_embed(dummy_hidden_states)
 
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)
@@ -425,4 +424,4 @@ def test_quantize_transformer_disabled(self, mock_quantize_model):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  absltest.main()