Test on attention type and automatically modify flash block sizes object when 'tokamax_flash' requested

Signed-off-by: Kunjan Patel <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/tests/wan_transformer_test.py

@@ -240,7 +240,7 @@ def test_wan_attention(self):
  query_dim=query_dim,
  heads=40,
  dim_head=128,
- attention_kernel=config.attention,
+ attention_kernel=attention_kernel,
  mesh=mesh,
  flash_block_sizes=flash_block_sizes,
  )