Implement batch-split schedule in deepseek_batchsplit.

PiperOrigin-RevId: 892685342

Author: Kevin Wang

src/maxtext/configs/base.yml

@@ -199,6 +199,7 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
+use_gather_mosaic_kernel: false # whether to use a custom mosaic kernel for token gather ops
 # tunable tiling dimensions used for mlp gmm
 # megablox/jax ragged dot - supports forward pass only (6 configs: `wi_tile_fwd...` and `wo_tile_fwd_...`)
 # tokamax ragged dot - supports all 18 configs

src/maxtext/configs/models/deepseek3-671b-batchsplit.yml

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# model config for DeepSeek V3 - 671B that uses fsdp on two logical axes
+# model config for DeepSeek V3 - 671B that uses batch split schedule
 
 # For DeepSeek default device-limited routing, 
 # please set n_routing_groups=8 and topk_routing_group=4 in your command-line arguments.
@@ -55,17 +55,18 @@ rope_interleave: True
 rope_truncate: True
 rope_attention_scaling: False
 
+use_batch_split_schedule: True
+shard_mode: "explicit"
 override_logical_axis_rules: True
-mesh_axes: ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert', 'context']
+mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'expert', 'context']
 data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert', 'context']]
 logical_axis_rules: [
     ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
     ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
     ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
     ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
-    ['activation_embed_and_logits_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-    ['activation_norm_length', ['context']],
-    ['activation_norm_length_moe', ['context']],
+    ['activation_norm_length', []],
+    ['activation_norm_length_moe', []],
     ['activation_heads', []],
     ['activation_stage', 'stage'],
     ['embed', ['fsdp']],
@@ -81,8 +82,8 @@ logical_axis_rules: [
     ['kv_heads', ['fsdp_transpose']],
     ['heads', ['fsdp_transpose']],
     ['mlp', ['fsdp_transpose']],
-    ['mlp_only_fsdp_transpose', ['fsdp_transpose']],
-    ['expert_only', ['expert']],
-    ['fsdp_transpose_only', ['fsdp_transpose']],
     ['fsdp_transpose_and_expert', ['fsdp_transpose', 'expert']],
+    ['fsdp_transpose_only', ['fsdp_transpose']],
+    ['expert_only', ['expert']],
+    ['diloco', 'diloco'],
 ]

src/maxtext/configs/types.py

@@ -655,6 +655,10 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use Ring of Experts for sparse matmul expert parallelism.",
   )
+  use_gather_mosaic_kernel: bool = Field(
+      False,
+      description="Whether to use a custom mosaic kernel for token gather ops.",
+  )
   use_random_routing: bool = Field(False, description="Whether to use random routing for debugging.")
   interleave_moe_layer_step: int = Field(1, description="Frequency of MoE layers, e.g., 2 means every 2nd layer is MoE.")
   expert_shard_attention_option: Literal["fsdp", "context"] = Field(
@@ -2683,8 +2687,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       self.use_grpo = False
 
     if self.use_batch_split_schedule:
-      if not (self.decoder_block == DecoderBlockType.DEEPSEEK and self.sparse_matmul and self.use_tokamax_gmm):
-        raise ValueError("Batch split only supports deepseek, with `sparse_matmul=True` and `use_tokamax_gmm=True`")
       if self.quantization and not (self.use_qwix_quantization and self.quantization == "fp8_full"):
         raise ValueError(
             "Batch split quantization only supports `use_qwix_quantization=True` and `quantization=fp8_full`"

src/maxtext/kernels/attention/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Attention kernels."""
+
+from maxtext.kernels.attention import splash_attention_kernel

src/maxtext/kernels/attention/splash_attention_kernel.py

@@ -427,9 +427,7 @@ def _wrapped(
       def reshape_activations(activations):
         if activations.ndim == 4:  # pytype: disable=attribute-error
           kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape  # pytype: disable=attribute-error
-          return activations.reshape(
-              kv_heads * q_heads_per_kv_head, q_seq_len, head_dim
-          )  # pytype: disable=attribute-error
+          return activations.reshape(kv_heads * q_heads_per_kv_head, q_seq_len, head_dim)  # pytype: disable=attribute-error
         return activations
 
       def reshape_residuals(residuals):
@@ -1166,10 +1164,7 @@ def _splash_attention_fwd(
     mask_function: MaskFunctionType | None,
     attn_logits_soft_cap: float | None = None,
     interpret: bool = False,
-) -> tuple[
-    tuple[jax.Array],
-    SplashResidualsType,
-]:
+) -> tuple[tuple[jax.Array], SplashResidualsType,]:
   """Forward pass for splash attention."""
   if save_residuals:
     raise NotImplementedError("Higher-order AD not supported")
@@ -1606,7 +1601,6 @@ def init():
   )
 
   def body(i, _):
-
     slice_k = pl.ds(i * bkv_compute, bkv_compute)
     q = q_ref[...]  # We keep q potentially transposed, since it's always RHS
 
@@ -2238,6 +2232,126 @@ def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None):
   )
 
 
+@partial(
+    jax.jit,
+    static_argnames=[
+        "is_mqa",
+        "block_sizes",
+        "save_residuals",
+        "mask_value",
+        "attn_logits_soft_cap",
+        "residual_checkpoint_name",
+        "mask_function",
+        "interpret",
+    ],
+)
+def _splash_attention_manual_fwd(
+    fwd_mask_info: mask_info_lib.MaskInfo,
+    dq_mask_info: mask_info_lib.MaskInfo | None,
+    dkv_mask_info: mask_info_lib.MaskInfo | None,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None = None,
+    sinks: jax.Array | None = None,
+    *,
+    is_mqa: bool,
+    block_sizes: BlockSizes | None,
+    save_residuals: bool,
+    mask_value: float,
+    attn_logits_soft_cap: float | None,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+    interpret: bool,
+) -> SplashCustomReturnType:
+  """Returns both the attention output and logsumexp.
+
+  This is useful when manually controlling remat in the backward pass, as both
+  can be returned as residuals from the forward pass."""
+
+  def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None):
+    if mask_info is None or mask_info.partial_mask_blocks is None:
+      return mask_info
+
+    return mask_info._replace(
+        partial_mask_blocks=mask_info.partial_mask_blocks.reshape(-1, *mask_info.partial_mask_blocks.shape[-2:])
+    )
+
+  if not save_residuals:
+    raise ValueError("Expected save_residuals to be `True`.")
+
+  fwd_mask_info = _collapse_partial_mask_blocks(fwd_mask_info)
+  dq_mask_info = _collapse_partial_mask_blocks(dq_mask_info)
+  dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info)
+  del dq_mask_info, dkv_mask_info
+
+  out, (logsumexp,) = _splash_attention_forward(  # pytype: disable=wrong-arg-types
+      fwd_mask_info,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      is_mqa=is_mqa,
+      block_sizes=block_sizes,
+      residual_checkpoint_name=residual_checkpoint_name,
+      save_residuals=True,
+      mask_function=mask_function,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      interpret=interpret,
+  )
+  return out, logsumexp
+
+
+def _splash_attention_manual_bwd(
+    fwd_mask_info: mask_info_lib.MaskInfo,
+    dq_mask_info: mask_info_lib.MaskInfo | None,
+    dkv_mask_info: mask_info_lib.MaskInfo | None,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    out: jax.Array,
+    logsumexp: jax.Array,
+    do: jax.Array,
+    segment_ids: SegmentIds | None = None,
+    sinks: jax.Array | None = None,
+    *,
+    is_mqa: bool,
+    block_sizes: BlockSizes | None,
+    save_residuals: bool,
+    mask_value: float,
+    attn_logits_soft_cap: float | None,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+    interpret: bool,
+):
+  """Transpose of _splash_attention_manual_fwd that uses attention output and logsumexp."""
+  del fwd_mask_info
+  res = (
+      q,
+      k,
+      v,
+      segment_ids,
+      out,
+      logsumexp,
+      dq_mask_info,
+      dkv_mask_info,
+  )
+  _, _, _, dq, dk, dv, _ = _splash_attention_bwd(
+      save_residuals=save_residuals,
+      mask_value=mask_value,
+      is_mqa=is_mqa,
+      block_sizes=block_sizes,
+      residual_checkpoint_name=residual_checkpoint_name,
+      mask_function=mask_function,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      interpret=interpret,
+      res=res,
+      do=do,
+  )
+  return dq, dk, dv
+
+
 @jax.tree_util.register_pytree_node_class
 class SplashAttentionKernel:
   """Defines a SplashAttention kernel object."""
@@ -2264,6 +2378,26 @@ def __call__(self, *args, **kwargs) -> SplashCustomReturnType:
         **self.kwargs,
     )
 
+  def manual_fwd(self, *args, **kwargs) -> SplashCustomReturnType:
+    return _splash_attention_manual_fwd(
+        self.fwd_mask_info,
+        self.dq_mask_info,
+        self.dkv_mask_info,
+        *args,
+        **kwargs,
+        **self.kwargs,
+    )
+
+  def manual_bwd(self, *args, **kwargs):
+    return _splash_attention_manual_bwd(
+        self.fwd_mask_info,
+        self.dq_mask_info,
+        self.dkv_mask_info,
+        *args,
+        **kwargs,
+        **self.kwargs,
+    )
+
   def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
     """Returns a value that can be used as a shard_map partition spec for the kernel."""
     if self.fwd_mask_info.data_next is not None:

src/maxtext/kernels/sort_activations.py

@@ -24,30 +24,30 @@
 def route(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> jax.Array:
   """Route tokens to selected experts."""
-  return _route_fwd(tokens, selected_experts, use_custom_mosaic_kernel)[0]
+  return _route_fwd(tokens, selected_experts, use_gather_mosaic_kernel)[0]
 
 
 def _route_fwd(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> tuple[jax.Array, jax.Array]:
   return (
-      _route_impl(tokens, selected_experts, use_custom_mosaic_kernel),
+      _route_impl(tokens, selected_experts, use_gather_mosaic_kernel),
       selected_experts,
   )
 
 
 def _route_bwd(
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
     residuals: jax.Array,
     grads: jax.Array,
 ) -> tuple[jax.Array, None]:
   selected_experts = residuals
-  return _unroute_impl(grads, selected_experts, use_custom_mosaic_kernel), None
+  return _unroute_impl(grads, selected_experts, use_gather_mosaic_kernel), None
 
 
 route.defvjp(_route_fwd, _route_bwd)
@@ -57,25 +57,25 @@ def _route_bwd(
 def unroute(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> jax.Array:
-  return _unroute_fwd(tokens, selected_experts, use_custom_mosaic_kernel)[0]
+  return _unroute_fwd(tokens, selected_experts, use_gather_mosaic_kernel)[0]
 
 
 def _unroute_fwd(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> tuple[jax.Array, jax.Array]:
   return (
-      _unroute_impl(tokens, selected_experts, use_custom_mosaic_kernel),
+      _unroute_impl(tokens, selected_experts, use_gather_mosaic_kernel),
       selected_experts,
   )
 
 
-def _unroute_bwd(use_custom_mosaic_kernel: bool, residuals: jax.Array, grads: jax.Array) -> tuple[jax.Array, None]:
+def _unroute_bwd(use_gather_mosaic_kernel: bool, residuals: jax.Array, grads: jax.Array) -> tuple[jax.Array, None]:
   selected_experts = residuals
-  return _route_impl(grads, selected_experts, use_custom_mosaic_kernel), None
+  return _route_impl(grads, selected_experts, use_gather_mosaic_kernel), None
 
 
 unroute.defvjp(_unroute_fwd, _unroute_bwd)
@@ -84,37 +84,32 @@ def _unroute_bwd(use_custom_mosaic_kernel: bool, residuals: jax.Array, grads: ja
 def _route_impl(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> jax.Array:
   """Gather `tokens` according to `selected_experts`."""
   assert (
       tokens.shape[0] == selected_experts.shape[0] and selected_experts.ndim == 2
   ), f"{tokens.shape=}, {selected_experts.shape=}"
-  if use_custom_mosaic_kernel:
-    raise NotImplementedError("Custom Mosaic kernel not implemented.")
   inds = jnp.argsort(jnp.ravel(selected_experts)) // selected_experts.shape[1]
-  return _sort_impl(tokens, inds, use_custom_mosaic_kernel)
+  return _sort_impl(tokens, inds, use_gather_mosaic_kernel)
 
 
 def _unroute_impl(
     tokens: jax.Array,
     selected_experts: jax.Array,
-    use_custom_mosaic_kernel: bool,
+    use_gather_mosaic_kernel: bool,
 ) -> jax.Array:
   """Reverse the routing operation, restoring tokens to their original order."""
   assert tokens.shape[0] == selected_experts.shape[0] * selected_experts.shape[1] and selected_experts.ndim == 2
   inds = jnp.argsort(jnp.argsort(jnp.ravel(selected_experts)))
   return jnp.sum(
       jnp.reshape(
-          _sort_impl(tokens, inds, use_custom_mosaic_kernel),
+          _sort_impl(tokens, inds, use_gather_mosaic_kernel),
           (-1, selected_experts.shape[1]) + tokens.shape[1:],
       ),
       axis=1,
   )
 
 
-def _sort_impl(tokens: jax.Array, inds: jax.Array, use_custom_mosaic_kernel: bool) -> jax.Array:
-  if use_custom_mosaic_kernel:
-    raise NotImplementedError("Custom Mosaic kernel not implemented.")
-  else:
-    return tokens[inds, ...]
+def _sort_impl(tokens: jax.Array, inds: jax.Array, use_gather_mosaic_kernel: bool) -> jax.Array:
+  return tokens[inds, ...]