Add SplashAttention scheduler config and related improvements

Author: zxhe-sean

src/MaxText/configs/base.yml

@@ -909,6 +909,7 @@ use_max_logit_estimate: -1 # -1 means no estimate, any > 0 value will be used as
 cost_estimate_flops_fwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (forward)
 cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (backward)
 dq_reduction_steps: 0 #the number of reduction steps. For now, only 3 or all the kv steps are supported.
+use_splash_scheduler: False # to use tokamax splash attention scheduler.
 ### Determine if we want to use load balance for context parallelism
 context_parallel_load_balance: True
 context_parallel_strategy: "all_gather" # "all_gather" or "ring"

src/MaxText/configs/types.py

@@ -561,6 +561,7 @@ class SplashAttention(BaseModel):
       0,
       description="the number of reduction steps. For now, only 3 or all the kv steps are supported.",
   )
+  use_splash_scheduler: bool = Field(False, description="Use experimental splash attention scheduler.")
 
 
 class PagedAttention(BaseModel):

src/MaxText/layers/attention_op.py

@@ -1123,6 +1123,7 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
             if config.cost_estimate_flops_bwd >= 0
             else None,
             dq_reduction_steps=config.dq_reduction_steps if config.dq_reduction_steps > 0 else None,
+            use_experimental_scheduler=config.use_splash_scheduler,
         )
       else:
         sa_config = splash_attention_kernel.BlockSizes(