change the flash_block_sizes

Author: susanbao

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: {