better block sizes.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -57,13 +57,13 @@ attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 
 #flash_block_sizes: {}
 flash_block_sizes: {
-  "block_q" : 2048,
-  "block_kv_compute" : 2048,
+  "block_q" : 3024,
+  "block_kv_compute" : 1024,
   "block_kv" : 2048,
-  "block_q_dkv" : 2048,
+  "block_q_dkv" : 3024,
   "block_kv_dkv" : 2048,
   "block_kv_dkv_compute" : 2048,
-  "block_q_dq" : 2048,
+  "block_q_dq" : 3024,
   "block_kv_dq" : 2048
 }
 # GroupNorm groups

src/maxdiffusion/generate_wan.py

@@ -80,7 +80,7 @@ def run(config, pipeline=None, filename_prefix=""):
       slg_start=slg_start,
       slg_end=slg_end,
   )
-  print("compile time: ", (time.perf_counter() - s0))
+  print("generation time: ", (time.perf_counter() - s0))
 
   s0 = time.perf_counter()
   if config.enable_profiler:

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -448,12 +448,20 @@ def __call__(
     return video
 
 
-@jax.jit
-def transformer_forward_pass(graphdef, sharded_state, rest_of_state, latents, timestep, prompt_embeds, is_uncond, slg_mask):
+@partial(jax.jit, static_argnames=("do_classifier_free_guidance", "guidance_scale"))
+def transformer_forward_pass(graphdef, sharded_state, rest_of_state, latents, timestep, prompt_embeds, is_uncond, slg_mask, do_classifier_free_guidance, guidance_scale):
   wan_transformer = nnx.merge(graphdef, sharded_state, rest_of_state)
-  return wan_transformer(
+  noise_pred = wan_transformer(
       hidden_states=latents, timestep=timestep, encoder_hidden_states=prompt_embeds, is_uncond=is_uncond, slg_mask=slg_mask
   )
+  if do_classifier_free_guidance:
+    bsz = latents.shape[0] // 2
+    noise_uncond = noise_pred[bsz:]
+    noise_pred = noise_pred[:bsz]
+    noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
+    latents = latents[:bsz]
+  
+  return noise_pred, latents
 
 
 def run_inference(
@@ -480,13 +488,11 @@ def run_inference(
     if slg_layers and int(slg_start * num_inference_steps) <= step < int(slg_end * num_inference_steps):
       slg_mask = slg_mask.at[jnp.array(slg_layers)].set(True)
     t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
-    # get original batch size before concat in case of cfg.
-    bsz = latents.shape[0]
     if do_classifier_free_guidance:
       latents = jnp.concatenate([latents] * 2)
     timestep = jnp.broadcast_to(t, latents.shape[0])
 
-    noise_pred = transformer_forward_pass(
+    noise_pred, latents = transformer_forward_pass(
         graphdef,
         sharded_state,
         rest_of_state,
@@ -495,12 +501,9 @@ def run_inference(
         prompt_embeds,
         is_uncond=jnp.array(True, dtype=jnp.bool_),
         slg_mask=slg_mask,
+        do_classifier_free_guidance=do_classifier_free_guidance,
+        guidance_scale=guidance_scale
     )
 
-    if do_classifier_free_guidance:
-      noise_uncond = noise_pred[bsz:]
-      noise_pred = noise_pred[:bsz]
-      noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
-      latents = latents[:bsz]
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
   return latents