Add Tokamax ring MHA attention for TPU

docs/guides/optimization/sharding.md

@@ -256,7 +256,9 @@ Context parallelism is similar to FSDP except we shard the sequence dimension of
 
 Care needs to be taken to shard the sequence dimension for attention - only the queries are sharded by sequence, the keys and values need to be all-gathered to perform the full computation. Additionally if we naively shard the sequence dimension then the attention computation is not evenly distributed due to the lower triangular causal mask - shards corresponding to later queries have more non-zero mask and thus become the bottleneck. Instead we “stripe” the inputs, so that the first shard has the first and last chunk of the sequence, the second shard has the second and second to last, etc. This striping is done on the initial data inputs (instead of every layer), so it is a small cost.
 
-Note in general there are many flavors of CP such as ring attention, which in theory can hide all of the comms (as opposed to this implementation where the KV all gathers are probably exposed). This all gather is relatively cheap so we have implemented this flavor for now, a good trade-off of complexity and performance. Currently TPUs only support this all gather strategy `context_parallel_strategy=all_gather`, but GPUs support both an `all_gather` strategy or a `ring` strategy which will perform the computation and communication in chunks and ideally overlap in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips.
+Note in general there are many flavors of CP such as ring attention, which in theory can hide all of the comms (as opposed to all gather CP where the KV all gathers are probably exposed). This all gather is relatively cheap so it remains a good trade-off of complexity and performance.
+
+MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips.
 
 ### CP Arithmetic Intensity
 

src/maxtext/configs/types.py

@@ -3261,15 +3261,60 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
     context_parallel_size = getattr(self, f"ici_{self.context_sharding}_parallelism", 1) * getattr(
         self, f"dcn_{self.context_sharding}_parallelism", 1
     )
-    if context_parallel_size > 1 and self.context_parallel_strategy.lower() == "ring":
-      if "gpu" not in self.hardware:
+    context_parallel_strategy = self.context_parallel_strategy.lower()
+    if (
+        context_parallel_strategy == "ring"
+        and "gpu" not in self.hardware
+        and "tpu" not in self.hardware
+        and context_parallel_size > 1
+    ):
+      raise ValueError(
+          "Ring context parallelism strategy (context_parallel_strategy='ring') is only supported on GPUs "
+          "or TPU with attention=flash and use_tokamax_splash=True."
+      )
+    if context_parallel_strategy == "ring" and "gpu" not in self.hardware and "tpu" in self.hardware:
+      if context_parallel_size <= 1:
+        raise ValueError("TPU Tokamax ring attention requires context_parallel_size > 1.")
+      if self.context_sharding != "context":
+        raise ValueError("TPU Tokamax ring attention requires context_sharding='context'.")
+      if self.dq_reduction_steps not in (0, 3):
+        raise ValueError("TPU Tokamax ring attention requires dq_reduction_steps to be 0 or 3.")
+      if self.max_target_length % (context_parallel_size * context_parallel_size) != 0:
         raise ValueError(
-            "Ring context parallelism strategy (context_parallel_strategy='ring') is only supported on GPUs."
+            "TPU Tokamax ring attention requires max_target_length to be divisible by context_parallel_size squared."
         )
+      if self.attention != "flash":
+        raise ValueError("TPU ring context parallelism requires attention=flash.")
+      if not self.use_tokamax_splash:
+        raise ValueError("TPU ring context parallelism requires use_tokamax_splash=True.")
+      if self.use_jax_splash:
+        raise ValueError("TPU ring context parallelism requires use_jax_splash=False.")
+      if self.attention_type != "global":
+        raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.")
+      if self.packing:
+        raise ValueError("TPU Tokamax ring attention does not support packing yet.")
+      if self.context_parallel_load_balance:
+        raise ValueError("TPU Tokamax ring attention does not support context_parallel_load_balance yet.")
+      if self.use_ragged_attention:
+        raise ValueError("TPU Tokamax ring attention does not support ragged attention.")
+      if self.attention_sink:
+        raise ValueError("TPU Tokamax ring attention does not support attention sinks.")
+      if self.use_indexer:
+        raise ValueError("TPU Tokamax ring attention does not support sparse indexer masks.")
+      if self.use_chunked_prefill:
+        raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.")
+      if self.moba:
+        raise ValueError("TPU Tokamax ring attention does not support MoBA.")
+      if self.use_multimodal:
+        raise ValueError("TPU Tokamax ring attention does not support multimodal attention.")
+      if self.use_qk_clip:
+        raise ValueError("TPU Tokamax ring attention does not support QK-Clip statistics yet.")
+      if self.enable_dropout and self.dropout_rate > 0.0:
+        raise ValueError("TPU Tokamax ring attention does not support dropout yet.")
     # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only.
-    # The AUTO + packing case (which training resolves to STRIPED) is not validated here
-    # because test code paths may load the same config but use a different reorder path.
-    # Training's runtime path in max_utils.reorder_causal_load_balanced enforces this.
+    # The AUTO + packing case, which training resolves to STRIPED, is not
+    # validated here because test code paths may load the same config but use a
+    # different reorder path. Training's runtime path enforces this.
     if (
         context_parallel_size > 1
         and "gpu" not in self.hardware