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

[NuojCheng]

Description

Adds moe_sparse_core_offload_targets, an opt-in, comma-separated list of routed-MoE ops to run on the TPU SparseCore instead of the TensorCore, freeing TensorCore cycles for the expert GEMMs. Default is "", so nothing changes for existing configs.

target	what moves to the SparseCore
fsdp_all_gather	the MoE weight all-gather over the fsdp / fsdp_transpose axes
ep_collectives	expert-parallel activation all-gathers and ragged all-to-alls, including the ring-of-experts dispatch/combine
ragged_sort	the routing index math in the ragged sort kernels

all enables every target. Unknown targets are rejected at config time, as is any target on hardware without a SparseCore (checked against compile_topology / hardware, so an AOT compile is validated the same way a real run is).

This is a redesign of #5125, which does not currently work: jax.experimental.compute_on.compute_on is a context manager whose only parameter is compute_type, so compute_on(fn, compute_type="tpu_sparsecore") raises TypeError: compute_on() got multiple values for argument 'compute_type' the first time the feature is exercised. Its is_non_gen7_tpu gate is also inverted for this purpose in both directions — it disables the feature on v5p and v6e, which do have SparseCores, and fails open on GPU (a3) — and its _fsdp_all_gather_with_rs backward is a maybe_shard_with_pspec resharding constraint rather than a reduce-scatter.

Design notes

Collectives are not a hint

The single most important finding here, and the thing that shaped the whole design.

For ordinary compute the annotation is advisory: XLA offloads what it can lower and quietly leaves the rest on the TensorCore. For collectives it is not. XLA's SparseCore collective-offload pass reads _xla_compute_type="sparseoffload" as force this one, and when it selects an annotated collective the chip cannot lower it CHECK-fails — aborting the compiler process, not falling back:

F sparse_core_collective_offload.cc:586] Candidate rejected: instruction has compute type
annotation sparseoffload but the operation is currently not supported on SC.
%all-gather-start = ...

Probed one collective per process on v5p with libtpu 0.0.46 (the toolchain CI uses), through a minimal shard_map, over 1D bf16[4096], 2D bf16[32,1024] and 3D bf16[8,1024,256]:

annotated collective	result
all_gather	aborts the compile at every rank
psum (all-reduce)	compiles, annotation silently stripped (sparseoffload=0)
psum_scatter (reduce-scatter)	compiles, genuinely offloaded (sparseoffload=3, async_execution_thread="sparsecore")
ragged_all_to_all	accepted

Note this happens under the default flags on v5p — --xla_tpu_enable_sparse_core_collective_offload_all_gather is already true there with libtpu 0.0.46, which contradicts the comments in benchmarks/xla_flags_library.py saying these default on only from Ironwood. An earlier draft of this PR assumed the pass was dormant on v5p; it is not, and the fsdp_all_gather TPU test took the whole pytest process down with the trace above.

The annotation survives AD, so transposition pairs move together

The compute type is snapshotted onto the jaxpr equation (jax/_src/core.py, JaxprEqnContext), and AD carries it onto the transposed equation. Verified both directions:

• annotating a forward all_gather produces an annotated — and genuinely offloaded — backward reduce-scatter;
• annotating a psum_scatter produces an annotated all-gather in transpose(jvp(...)), which is just as fatal as annotating one directly.

So a collective may only be annotated when both ends of its transposition pair are offloadable. utils/sparsecore.py encodes this as _TRANSPOSE_OF, and on v5p it is why reduce-scatter is disabled despite the probe above showing it works standalone.

Capability model

supports_collective_offload(collective, ...) gates each collective on the lowest chip generation both it and its transpose are known to work on (_MIN_GENERATION; all-gather is Ironwood-and-later per xla_flags_library.py and the original SparseCore MoE work). supported_offload_targets(...) then drops any target whose one required collective the chip cannot serve — with a warning, not an error, so the same config file runs unchanged across chip generations.

chip	fsdp_all_gather	ep_collectives	ragged_sort
tpu7x and later	offloaded	fully offloaded	offloaded
v5p, v6e	warned and ignored	ragged all-to-alls and their index math offloaded; all-gathers / reduce-scatters left on the TensorCore	offloaded
v4, v5e, GPU, CPU	rejected at config time (no SparseCore)	ditto	ditto

Chip detection reads JAX's own table (pltpu.get_tpu_info_for_chip) for both SparseCore presence and generation, rather than pattern-matching device strings, so it stays correct as chips are added.

Where the annotation goes

compute_on tags ops as they are traced, so there has to be an op to tag. The MoE weight all-gather had none: the weights entered the sparse_matmul shard_map with an in_spec that dropped the FSDP mesh axes and GSPMD synthesized the collective at that boundary. When fsdp_all_gather is live the weights are instead handed to the shard_map still FSDP-sharded and the identical gather is performed inside the manual region, where it is a real jax.lax.all_gather that can carry a compute type. _fsdp_weight_all_gather_plan derives the (dim, axes) gathers from the source and target PartitionSpecs and falls back to the previous implicit boundary gather whenever the transition is not a pure all-gather, the weights are quantized, or explicitly_weight_ag() already hand-wrote its own gather.

The restructuring is deliberately tied to the target being live, not merely requested: on v5p/v6e the target is dropped, so the graph is byte-identical to moe_sparse_core_offload_targets: "". Making the gather explicit is arguably a win in its own right — with it the weight-gradient collectives go from 5 all-reduces to 3 all-reduces plus 2 reduce-scatters, gradients bit-identical — but that is an unrelated change and shouldn't ride in as a consolation prize under a flag named for offloading. Happy to be argued out of this and let the rewrite apply wherever the target is requested.

all_gather's default to="varying" is load-bearing: it transposes to the psum_scatter that reduces each shard's partial weight gradient. to="invarying" transposes to a bare dynamic_slice_in_dim with no psum, silently dropping the FSDP gradient reduction — the forward value stays correct, so this only shows up in a gradient comparison, which is what the equivalence tests below exist to catch.

One more trap the plan has to dodge: the manual gather makes the shard_map body vary over the FSDP axis, but maybe_replicate_incompatible_batch strips that axis off the output pspec when the batch is not divisible by the mesh, and check_vma=True then rejects the region outright. The plan is only taken when every gathered axis still appears in the out specs; otherwise it logs and leaves the gather to the partitioner.

ragged_sort annotates index math only

Argsorts, one-hot histograms, group sizes/offsets, and permutations of 1D index/weight vectors move; the Pallas ragged_gather / ragged_gather_reduce kernels and everything touching the hidden dimension stay on the TensorCore, since that is dense vector work the SparseCore is not the right place for. ragged_sort involves no collectives, so it is unaffected by everything above and runs on every SparseCore chip.

AOT

offload() deliberately does not check local devices. An ahead-of-time compile for a SparseCore topology usually runs on a host with no TPU at all, and its HLO has to match what the real run compiles; the config validator has already established that the target chip has a SparseCore, and the frontend attribute is inert on backends that don't consume it.

Tests

tests/unit/sparsecore_test.py (50 CPU tests): target-string parsing, chip-table detection for v5p/v6e/tpu7x/v4/v5e/GPU/unknown topologies, the per-collective capability gate including the transpose rule, warn-and-drop of unserviceable targets, the PartitionSpec → all-gather-plan derivation including the cases that must decline (gaining an axis, dropping a major axis, swapped axes, a pspec naming more dims than the array has), and config validation.

Two of those run the real lowering on CPU and assert _xla_compute_type = "sparseoffload" is present when offloading and absent when not. That coverage exists because the compute-type API is the part of this feature most likely to move under us and every test that could catch it moving used to be TPU-gated — a JAX upgrade could have broken the annotation with CI green.

SparseCoreOffloadTest in tests/unit/moe_test.py (@pytest.mark.tpu_only, 6 cases): for each target, compares loss and every parameter gradient against the same config with offloading off, and asserts the baseline has zero annotations while the offloaded HLO has some. Gradients are constrained back to the parameter sharding so the weight-gradient collectives actually appear — this is the test that catches the to="invarying" trap. The parameterized fsdp_all_gather case skips where the chip cannot offload an all-gather, since the target is dropped there and there would be nothing to assert.

That skip would have left the riskiest code here untested on the only hardware available, so test_explicit_fsdp_weight_gather_matches_the_implicit_one patches the capability gate out and runs the comparison anyway. The rewrite's numerics have nothing to do with the SparseCore; the per-collective gate inside sparsecore.offload still suppresses the annotation so XLA does not abort. It asserts the annotation count matches what the chip supports, and that the baseline has no reduce-scatter while the rewritten graph does — without which the whole comparison would be the baseline against itself.

$ pytest tests/unit/sparsecore_test.py -q                       # CPU, jax 0.11.1 and nightly
50 passed
$ pytest tests/unit/moe_test.py -q -k SparseCoreOffloadTest     # 4x v5p, 0.11.1 and nightly
5 passed, 1 skipped
$ pytest tests/unit/moe_test.py -q                              # 4x v5p, 0.11.1
5 failed, 32 passed, 43 skipped

The 5 failures (test_gmm_grad_equivalence_tokamax_v2_fp8_{dynamic_ep4,static_ep1,static_ep1_qag,static_ep4} and test_shard_embed_moe_on_fsdp) reproduce identically at the merge base in a clean worktree with the same venv — pre-existing and unrelated.

Gradient equivalence

Separate harness over 7 parallelism/feature combinations, mixtral-8x7b on 4x v5p, comparing loss and all 5 parameter gradients against offloading off. Run on both jax 0.11.1 and nightly. (The fsdp4 row is the no-op-on-v5p case; the rewrite itself is covered by the unit test above.)

variant	loss	grads bit-exact	worst rel. err
fsdp4 + fsdp_all_gather	bit-exact	5/5	0.000e+00
ep4 + ep_collectives	bit-exact	5/5	0.000e+00
ep4 + ragged_sort	bit-exact	5/5	0.000e+00
ep4 + ragged + all	bit-exact	5/5	0.000e+00
fsdp2 x ep2 + all	bit-exact	5/5	0.000e+00
ep4 + ring-of-experts + chunking + all	bit-exact	5/5	0.000e+00
ep4 + ring-of-experts + ragged + all	bit-exact	5/5	0.000e+00

Benchmarks

mixtral-8x7b MoE layer, 4x v5p, bf16, megablox, base_emb_dim=4096, base_mlp_dim=1024, per-device batch 4, seqlen 512; median of 30 iterations after warmup. HLO counts are AR = all-reduce, AG = all-gather, A2A = all-to-all, RS = reduce-scatter, SO = sparseoffload annotations.

Default XLA flags:

variant	median	vs. its baseline	HLO
baseline_fsdp4	6.25 ms	—	AR 5, AG 4, SO 0
fsdp4 + fsdp_all_gather	6.25 ms	+0.05%	AR 5, AG 4, SO 0
baseline_ep4	8.63 ms	—	AR 3, AG 2, A2A 4, SO 0
ep4 + ep_collectives	8.55 ms	-0.97%	AR 2, AG 2, A2A 4, RS 1, SO 14
baseline_ep4_ragged	8.88 ms	—	AR 3, AG 2, A2A 4, SO 0
ep4 + ragged + ragged_sort	8.82 ms	-0.68%	AR 2, AG 2, A2A 4, RS 1, SO 10
ep4 + ragged + all	8.85 ms	-0.34%	AR 2, AG 2, A2A 4, RS 1, SO 24
baseline_fsdp2_ep2	7.44 ms	—	AR 6, AG 5, A2A 4, SO 0
fsdp2 x ep2 + all	7.34 ms	-1.34%	AR 2, AG 5, A2A 4, RS 4, SO 24

The fsdp4 row is the backward-compatibility datapoint, not a result: on v5p the target is dropped, so the HLO is identical to the baseline and the 0.05% is noise.

With the SparseCore op-offload flags on (--xla_tpu_enable_offloading_{sort,gather,reduce,scatter,reshape,copy}_to_sparsecore=true), which also speed up the baselines:

variant	median	vs. its baseline
baseline_ep4_ragged	8.74 ms	—
ep4 + ragged + ragged_sort	8.64 ms	-1.17%
ep4 + ragged + all	8.62 ms	-1.37%
baseline_ep4	8.21 ms	—
ep4 + ep_collectives	8.23 ms	+0.24%

Numerics at benchmark scale: losses across all 19 runs take one of two float64-printed values, 1.2e-10 apart in relative terms, and which one a run lands on does not track the offload setting — the same unmodified baseline gives one value under default flags and the other with the op-offload flags on. That is collective-decomposition reassociation in XLA, not the annotation.

Backward compatibility: the default moe_sparse_core_offload_targets: "" produces HLO with zero sparseoffload annotations and takes the original implicit-boundary-gather path, asserted by test_disabled_offload_leaves_the_hlo_untouched.

Shortcomings

The only SparseCore hardware available for this work was v5p, where the SparseCore can take the ragged all-to-alls but not the all-gather. So fsdp_all_gather is exercised here only for correctness and for its no-op behaviour; its throughput case rests on Ironwood. _MIN_GENERATION is deliberately conservative and should be revisited on that hardware — in particular an all-reduce entry may want adding once there is a chip where offloading one does something, and reduce-scatter is currently disabled on v5p purely because its transpose is an all-gather, not because the reduce-scatter itself fails.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for offloading specific Mixture of Experts (MoE) operations—such as FSDP weight all-gathers, expert-parallel collectives, and ragged sort index math—to the TPU SparseCore instead of the TensorCore. This is controlled via a new configuration parameter moe_sparse_core_offload_targets and supported by a new sparsecore.py utility module. The review feedback suggests several robustness improvements: wrapping the compile-time topology check in a try-except block to handle potential JAX/Pallas version mismatches, checking for None values of pspec in _pspec_axes_per_dim to prevent a TypeError, and unpacking single-element tuples to plain strings in jax.lax.all_gather for broader compatibility.

[gemini-code-assist]

The compile-time topology check relies on pltpu.get_tpu_info_for_chip which can raise ValueError or AttributeError depending on the JAX/Pallas version or environment. Wrapping the entire block in a try-except block catching these exceptions prevents unexpected compile-time crashes.

Suggested change

	if compile_topology:
	try:
	spec = accelerator_to_spec_map.get_system_characteristics(compile_topology)
	except ValueError:
	return None
	if spec.platform != "tpu":
	return None
	chip_version = _chip_version(compile_topology)
	if chip_version is None:
	return None
	# SparseCore presence does not depend on the Megacore split, so 1 core per
	# logical device is a valid probe for every chip version.
	return pltpu.get_tpu_info_for_chip(chip_version, 1).sparse_core
	if compile_topology:
	try:
	spec = accelerator_to_spec_map.get_system_characteristics(compile_topology)
	if spec.platform != 'tpu':
	return None
	chip_version = _chip_version(compile_topology)
	if chip_version is None:
	return None
	# SparseCore presence does not depend on the Megacore split, so 1 core per
	# logical device is a valid probe for every chip version.
	return pltpu.get_tpu_info_for_chip(chip_version, 1).sparse_core
	except (ValueError, AttributeError, RuntimeError):
	return None

[gemini-code-assist]

If pspec is None (representing fully replicated sharding in JAX), len(pspec) will raise a TypeError. Adding a check for pspec is not None ensures robust handling of replicated shardings.

Suggested change

	def _pspec_axes_per_dim(pspec, ndim):
	"""Normalizes a PartitionSpec into one axis-name tuple per array dim."""
	per_dim = []
	for i in range(ndim):
	axis = pspec[i] if i < len(pspec) else None
	if axis is None:
	per_dim.append(())
	elif isinstance(axis, str):
	per_dim.append((axis,))
	else:
	per_dim.append(tuple(axis))
	return per_dim
	def _pspec_axes_per_dim(pspec, ndim):
	'''Normalizes a PartitionSpec into one axis-name tuple per array dim.'''
	per_dim = []
	for i in range(ndim):
	axis = pspec[i] if pspec is not None and i < len(pspec) else None
	if axis is None:
	per_dim.append(())
	elif isinstance(axis, str):
	per_dim.append((axis,))
	else:
	per_dim.append(tuple(axis))
	return per_dim

[gemini-code-assist]

Passing a single-element tuple as axis_name to jax.lax.all_gather can sometimes cause issues or different code paths in older JAX versions. Unpacking single-element tuples to a plain string is safer and more compatible.

Suggested change

	for w, plan in zip((w0, w1, wo), weight_ag_plans):
	for dim, axes in plan:
	w = jax.lax.all_gather(w, axes, axis=dim, tiled=True)
	gathered.append(w)
	for w, plan in zip((w0, w1, wo), weight_ag_plans):
	for dim, axes in plan:
	axis_name = axes[0] if len(axes) == 1 else axes
	w = jax.lax.all_gather(w, axis_name, axis=dim, tiled=True)
	gathered.append(w)

[codecov]

Codecov Report

❌ Patch coverage is 89.11565% with 32 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/sparsecore.py	84.12%	15 Missing and 5 partials ⚠️
src/maxtext/layers/moe.py	90.00%	3 Missing and 3 partials ⚠️
src/maxtext/utils/sharding.py	82.35%	3 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!