passing two options as configs

Author: eltsai

src/maxdiffusion/configs/base_wan_14b.yml

@@ -62,6 +62,8 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses
+use_base2_exp: False
+use_experimental_scheduler: False
 flash_min_seq_length: 0
 
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -61,6 +61,8 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring, ulysses
+use_base2_exp: False
+use_experimental_scheduler: False
 flash_min_seq_length: 0
 
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.

src/maxdiffusion/configs/base_wan_27b.yml

@@ -62,6 +62,8 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring, ulysses
+use_base2_exp: False
+use_experimental_scheduler: False
 flash_min_seq_length: 4096
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -61,6 +61,8 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring, ulysses
+use_base2_exp: False
+use_experimental_scheduler: False
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -61,6 +61,8 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring, ulysses
+use_base2_exp: False
+use_experimental_scheduler: False
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/models/attention_flax.py

@@ -272,6 +272,7 @@ def convert_to_tokamax_splash_config(
     attn_logits_soft_cap: float | None = None,
     fuse_reciprocal: bool = True,
     use_base2_exp: bool = False,
+    use_experimental_scheduler: bool = False,
     max_logit_const: float | None = None,
     interpret: bool = False,
     dq_reduction_steps: int | None = None,
@@ -294,6 +295,7 @@ def convert_to_tokamax_splash_config(
       attn_logits_soft_cap=attn_logits_soft_cap,
       fuse_reciprocal=fuse_reciprocal,
       use_base2_exp=use_base2_exp,
+      use_experimental_scheduler=use_experimental_scheduler,
       max_logit_const=max_logit_const,
       interpret=interpret,
       dq_reduction_steps=dq_reduction_steps,
@@ -314,6 +316,8 @@ def _tpu_flash_attention(
     mask_padding_tokens: bool = True,
     residual_checkpoint_name: str | None = None,
     attention_mask: jax.Array = None,
+    use_base2_exp: bool = False,
+    use_experimental_scheduler: bool = False,
 ) -> jax.Array:
   """TPU Flash Attention"""
 
@@ -399,7 +403,12 @@ def wrap_flash_attention(query, key, value):
       splash_kernel = tokamax_splash_attention_kernel.make_splash_mha(
           mask=mask,
           q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
-          config=convert_to_tokamax_splash_config(block_sizes, residual_checkpoint_name=residual_checkpoint_name),
+          config=convert_to_tokamax_splash_config(
+              block_sizes,
+              residual_checkpoint_name=residual_checkpoint_name,
+              use_base2_exp=use_base2_exp,
+              use_experimental_scheduler=use_experimental_scheduler,
+          ),
           save_residuals=False,
       )
     elif attention_kernel == "tokamax_ring":
@@ -409,7 +418,12 @@ def wrap_flash_attention(query, key, value):
       splash_kernel = tokamax_ring_attention_kernel.make_ring_attention(
           mask=mask,
           is_mqa=False,
-          config=convert_to_tokamax_splash_config(block_sizes, residual_checkpoint_name=residual_checkpoint_name),
+          config=convert_to_tokamax_splash_config(
+              block_sizes,
+              residual_checkpoint_name=residual_checkpoint_name,
+              use_base2_exp=use_base2_exp,
+              use_experimental_scheduler=use_experimental_scheduler,
+          ),
           save_residuals=False,
           ring_axis="context",
           rotate_segment_ids=False,  # We don't rotate segment ids in tokamax ring attention because our segment ids is for padding each kv shard has same segment ids
@@ -741,6 +755,8 @@ def _apply_attention(
     mask_padding_tokens: bool = True,
     residual_checkpoint_name: str | None = None,
     attention_mask: Array = None,
+    use_base2_exp: bool = False,
+    use_experimental_scheduler: bool = False,
 ):
   """Routes to different attention kernels."""
   _check_attention_inputs(query, key, value)
@@ -789,6 +805,8 @@ def _apply_attention(
         mask_padding_tokens=mask_padding_tokens,
         residual_checkpoint_name=residual_checkpoint_name,
         attention_mask=attention_mask,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
   elif "ring" in attention_kernel:
     return _tpu_flash_attention(
@@ -983,8 +1001,12 @@ def __init__(
       quant: Quant = None,
       mask_padding_tokens: bool = True,
       residual_checkpoint_name: str | None = None,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
   ):
     self.dpa_layer = None
+    self.use_base2_exp = use_base2_exp
+    self.use_experimental_scheduler = use_experimental_scheduler
     if attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
 
@@ -1045,6 +1067,8 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         mask_padding_tokens=self.mask_padding_tokens,
         residual_checkpoint_name=self.residual_checkpoint_name,
         attention_mask=attention_mask,
+        use_base2_exp=self.use_base2_exp if hasattr(self, "use_base2_exp") else False,
+        use_experimental_scheduler=self.use_experimental_scheduler if hasattr(self, "use_experimental_scheduler") else False,
     )
 
 
@@ -1063,6 +1087,8 @@ class AttentionOp(nn.Module):
   flash_block_sizes: BlockSizes = None
   dtype: DType = jnp.float32
   quant: Quant = None
+  use_base2_exp: bool = False
+  use_experimental_scheduler: bool = False
 
   def setup(self):
     self.dpa_layer = None
@@ -1108,6 +1134,8 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         flash_block_sizes=self.flash_block_sizes,
         dpa_layer=self.dpa_layer,
         attention_mask=attention_mask,
+        use_base2_exp=self.use_base2_exp,
+        use_experimental_scheduler=self.use_experimental_scheduler,
     )
 
 
@@ -1144,6 +1172,8 @@ def __init__(
       enable_jax_named_scopes: bool = False,
       added_kv_proj_dim: Optional[int] = None,  # New for I2V
       image_seq_len: Optional[int] = None,  # New for I2V
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
   ):
     if attention_kernel in {"flash", "cudnn_flash_te"} and mesh is None:
       raise ValueError(f"The flash attention kernel requires a value for mesh, but mesh is {self.mesh}")
@@ -1186,6 +1216,8 @@ def __init__(
         quant=quant,
         mask_padding_tokens=mask_padding_tokens,
         residual_checkpoint_name=residual_checkpoint_name,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
     # None axes corresponds to the stacked weights across all blocks
     # because of the use of nnx.vmap and nnx.scan.

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -291,6 +291,8 @@ def __init__(
       dropout: float = 0.0,
       mask_padding_tokens: bool = True,
       enable_jax_named_scopes: bool = False,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
   ):
     self.enable_jax_named_scopes = enable_jax_named_scopes
 
@@ -315,6 +317,8 @@ def __init__(
         mask_padding_tokens=mask_padding_tokens,
         residual_checkpoint_name="self_attn",
         enable_jax_named_scopes=enable_jax_named_scopes,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     # 1. Cross-attention
@@ -339,6 +343,8 @@ def __init__(
         mask_padding_tokens=mask_padding_tokens,
         residual_checkpoint_name="cross_attn",
         enable_jax_named_scopes=enable_jax_named_scopes,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
     assert cross_attn_norm is True
     self.norm2 = FP32LayerNorm(rngs=rngs, dim=dim, eps=eps, elementwise_affine=True)
@@ -486,6 +492,8 @@ def __init__(
       mask_padding_tokens: bool = True,
       scan_layers: bool = True,
       enable_jax_named_scopes: bool = False,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
   ):
     inner_dim = num_attention_heads * attention_head_dim
     out_channels = out_channels or in_channels
@@ -547,6 +555,8 @@ def init_block(rngs):
           enable_jax_named_scopes=enable_jax_named_scopes,
           added_kv_proj_dim=added_kv_proj_dim,
           image_seq_len=image_seq_len,
+          use_base2_exp=use_base2_exp,
+          use_experimental_scheduler=use_experimental_scheduler,
       )
 
     self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy)

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -139,6 +139,8 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   wan_config["mask_padding_tokens"] = config.mask_padding_tokens
   wan_config["scan_layers"] = config.scan_layers
   wan_config["enable_jax_named_scopes"] = config.enable_jax_named_scopes
+  wan_config["use_base2_exp"] = config.use_base2_exp
+  wan_config["use_experimental_scheduler"] = config.use_experimental_scheduler
 
   # 2. eval_shape - will not use flops or create weights on device
   # thus not using HBM memory.