change the flash_block_sizes (#290)

* change the flash_block_sizes

* fix unit test

---------

Co-authored-by: susanbao <sanbao_google_com@t1v-n-216c02cd-w-0.europe-west4-b.c.cloud-tpu-multipod-dev.internal>

Author: susanbao

.github/workflows/UnitTests.yml

@@ -58,6 +58,7 @@ jobs:
         pip show jax jaxlib flax transformers datasets tensorflow tensorflow_datasets
     - name: PyTest
       run: | #--deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py
+        export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
         HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py -x
 #  add_pull_ready:
 #    if: github.ref != 'refs/heads/main'

src/maxdiffusion/configs/base_wan_14b.yml

@@ -64,14 +64,13 @@ flash_min_seq_length: 4096
 dropout: 0.1
 
 flash_block_sizes: {
-  "block_q" : 1024,
-  "block_kv_compute" : 256,
-  "block_kv" : 1024,
-  "block_q_dkv" : 1024,
-  "block_kv_dkv" : 1024,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 1024,
-  "block_kv_dq" : 1024
+  "block_q" : 2048,
+  "block_kv_compute" : 512,
+  "block_kv" : 2048,
+  "block_q_dkv" : 2048,
+  "block_kv_dkv" : 2048,
+  "block_kv_dkv_compute" : 512,
+  "use_fused_bwd_kernel" : True
 }
 # Use on v6e
 # flash_block_sizes: {