Add experimental settings in tokamax

Author: Perseus14

src/maxdiffusion/models/attention_flax.py

@@ -271,7 +271,8 @@ def convert_to_tokamax_splash_config(
     residual_checkpoint_name: str | None = None,
     attn_logits_soft_cap: float | None = None,
     fuse_reciprocal: bool = True,
-    use_base2_exp: bool = False,
+    use_base2_exp: bool = True,
+    use_experimental_scheduler: bool = True,
     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,