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

@@ -212,7 +212,8 @@ def test_wan_attention(self):
               None,
               os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
               f"attention={attention_kernel}"
-          ]
+          ],
+          unittest=True
       )
       config = pyconfig.config
       batch_size = 1