Offload MoE collectives and ragged-sort index math to the TPU SparseCore

src/maxtext/configs/base.yml

@@ -294,6 +294,22 @@ norm_topk_prob: false # boolean to enable the top-k probability normalization. q
 
 # when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
 moe_fsdp_use_two_stage_all_gather: false
+# Comma-separated list of MoE ops to run on the TPU SparseCore instead of the TensorCore, freeing
+# TensorCore cycles for the expert GEMMs. Empty (the default) keeps everything on the TensorCore.
+# Supported targets, or "all" for every one of them:
+#   fsdp_all_gather - the MoE weight all-gather over the fsdp / fsdp_transpose axes
+#   ep_collectives  - the expert-parallel activation all-gathers and ragged all-to-alls
+#   ragged_sort     - the routing index math (argsorts, group sizes, offsets) in the ragged sort kernels
+# Requires a TPU with a SparseCore (v5p, v6e, tpu7x or newer); MaxText raises otherwise.
+#
+# Offloading never changes results. For collectives it is not merely a hint though: XLA's
+# SparseCore collective-offload pass reads the annotation as "force this one" and CHECK-fails,
+# aborting the compile, on an annotated collective the chip cannot lower. MaxText therefore
+# only annotates a collective its SparseCore is known to run (utils/sparsecore.py). Offloading
+# an all-gather needs Ironwood, and a reduce-scatter transposes to one in the backward pass, so on
+# v5p/v6e the fsdp_all_gather target logs a warning and is ignored, and ep_collectives is left with
+# just its ragged all-to-alls. ragged_sort is not a collective and runs everywhere.
+moe_sparse_core_offload_targets: ""
 # Shard the expert dimension of the MLP weights on the FSDP axis.
 # This configuration is recommended only when num_experts is a multiple of fsdp_parallelism
 shard_exp_on_fsdp: false

src/maxtext/configs/types.py

@@ -37,6 +37,7 @@
 from maxtext.utils import elastic_utils
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, HF_IDS
 from maxtext.utils import accelerator_to_spec_map
+from maxtext.utils import sparsecore
 from pydantic.config import ConfigDict
 from pydantic.fields import Field
 from pydantic.functional_validators import field_validator, model_validator
@@ -979,6 +980,13 @@ class MoEGeneral(BaseModel):
       False,
       description="Use two separate All-Gather calls for MoE weights sharded on both FSDP and FSDP-transpose.",
   )
+  moe_sparse_core_offload_targets: str = Field(
+      "",
+      description="Comma-separated list of MoE ops to run on the TPU SparseCore instead of the TensorCore. "
+      f"Supported targets: {', '.join(sparsecore.OFFLOAD_TARGETS)}; 'all' enables every one of them. "
+      "Empty (the default) keeps everything on the TensorCore. Requires a TPU with a SparseCore; targets "
+      "whose collectives that SparseCore cannot run are warned about and ignored.",
+  )
   shard_exp_on_fsdp: bool = Field(
       False,
       description="Shard the expert dimension of the MLP weights on the FSDP axis, "
@@ -3283,6 +3291,23 @@ def _validate_check_vma_is_supported(self):
           f"Found other ICI axes enabled: {active}."
       )
 
+  def _validate_sparse_core_offload(self):
+    """Validates moe_sparse_core_offload_targets against the target hardware."""
+    # Raises on unrecognized target names.
+    targets = sparsecore.parse_offload_targets(self.moe_sparse_core_offload_targets)
+    if not targets:
+      return
+    if not sparsecore.has_sparse_core(self.compile_topology, self.hardware):
+      raise ValueError(
+          f"moe_sparse_core_offload_targets={self.moe_sparse_core_offload_targets!r} requires a TPU with a "
+          "SparseCore (v5p, v6e, tpu7x or newer), but the target hardware has none. Set it to '' to keep "
+          "these ops on the TensorCore."
+      )
+    # Warns now, at startup, for any target this chip cannot serve, rather than
+    # leaving the first one to surface mid-trace. Unserviceable targets are
+    # dropped instead of rejected so the same config runs on every chip.
+    sparsecore.supported_offload_targets(self.moe_sparse_core_offload_targets, self.compile_topology, self.hardware)
+
   def validate_ragged_buffer_factor(self):
     if self.ragged_buffer_factor <= 0:
       return  # Not using a ragged buffer factor
@@ -4807,6 +4832,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       )
 
     self._validate_check_vma_is_supported()
+    self._validate_sparse_core_offload()
 
     # Final string-to-enum conversions if they haven't been coerced by pydantic yet.
     if isinstance(self.decoder_block, str):