MoE ring-of-experts: TC Pallas ring all-gather for the combine backward + x_sorted remat save

[ultrons]

Stacked on #4895 (token quantized all-gather); the base branch is that PR's head, so this diff shows only the incremental change.

What

Two default-off flags for performance improvement on expert-parallel MoE (ring-of-experts path), plus one new kernel file:

• moe_ring_cotangent_ag — the backward cotangent all-gather of the ring-of-experts combine reduce-scatter runs on a TensorCore Pallas bidirectional ring kernel instead of the XLA collective. As an XLA collective this all-gather can be placed on the SparseCore collective-offload queue and serialize behind other SC work; issuing its ICI DMAs from the TensorCore lets the scheduler overlap it with backward compute. The forward is unchanged (plain psum_scatter), so a remat recompute re-traces the plain collective and no Pallas DMA runs inside a rematted region. Numerically equal to lax.all_gather (pure tiled data move).
• moe_x_sorted (RematLocation, default remat) — device saves the routed (expert-sorted) MoE input and its small routing/metadata bundle across the remat boundary, so the backward loads them instead of re-running the dispatch token all-gather and sort. Batch-size sensitive: intended for small per-device batch; at larger per-device batch the save exceeds HBM (see repro notes).

New file src/maxtext/kernels/ring_ag.py: bidirectional store-and-forward ring all-gather (TC Pallas, ICI DMA), used inside the MoE shard_map with a caller-supplied collective_id (avoids barrier-semaphore collisions with other in-flight Pallas collectives) and an explicit CostEstimate so the latency-hiding scheduler accounts for the DMA cost.

Both compose with #4895's moe_quantize_token_all_gather (the moe_x_sorted tag is applied leaf-wise so a QArray input is saved correctly).

Repro

DeepSeek-V3 671B on v7x, the recipe from #4895 with the new flags added:

python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml \
  model_name=deepseek3-671b per_device_batch_size=1.0 max_target_length=4096 \
  ici_fsdp_parallelism=64 ici_expert_parallelism=8 \
  attention=flash sparse_matmul=true megablox=true use_tokamax_gmm=true use_gmm_v2=true \
  use_tokamax_splash=true use_ring_of_experts=true use_ragged_sort=true \
  use_custom_sort_vjp=true ragged_buffer_factor=2.0 num_moe_token_chunks=2 \
  remat_policy=custom decoder_layer_input=offload context=device out_proj=device \
  quantization=fp8_full use_qwix_quantization=true \
  weight_quantization_calibration_method=fixed,-224,224 \
  act_quantization_calibration_method=fixed,-224,224 \
  bwd_quantization_calibration_method=absmax \
  moe_quantize_token_all_gather=true \
  moe_ring_cotangent_ag=true moe_x_sorted=device \
  dataset_type=synthetic use_random_routing=true steps=20

• moe_ring_cotangent_ag: any EP degree with ring-of-experts; loss matches the flag-off run (the backward all-gather is bit-exact).
• moe_x_sorted=device: use with small per-device batch. AOT compile-checked: fits at per_device_batch_size=1.0 on a 4x8x8 mesh; at per_device_batch_size=4.0 the saved activations exceed HBM (drop this flag or keep remat there).
• Flags off reproduce the existing behavior byte-for-byte (moe_x_sorted tags are inert under remat; the ring branches are not taken).

Validation

• AOT compile (virtual tpu7x topology, no hardware) green for: the Add token quantized all gather flow #4895 recipe + moe_ring_cotangent_ag at per_device_batch_size=4.0, and both flags at per_device_batch_size=1.0.
• Ring kernel semantics validated on v7x against lax.all_gather(..., tiled=True) (bit-exact).
• Loss parity vs flag-off verified on a 512-chip DeepSeek-V3 run (identical lm_loss trajectory).

[gemini-code-assist]

Code Review

This pull request introduces a bidirectional store-and-forward ring all-gather kernel implemented in TensorCore Pallas (ring_ag.py) and integrates it into the MoE layer to optimize the backward cotangent all-gather. It also adds configuration options for this kernel and a new remat policy (moe_x_sorted) to save routed expert inputs. The review feedback highlights two critical issues in the new Pallas kernel: a missing semaphore wait on recv_sem that could cause a race condition during data forwarding, and a redundant duplicate call to pltpu.make_async_copy when copying the local shard.

[gemini-code-assist]

In the bidirectional store-and-forward ring, device my must wait for the incoming chunk from its neighbor to be fully written to its local buffer before forwarding it in the next step. Without waiting on recv_sem.at[si], a race condition occurs where the device may read and forward stale or unwritten data from o_ref.

To fix this, add a semaphore wait on recv_sem.at[si] for steps i > 0.

      if prev[li] is not None:
        prev[li].wait()
      if i > 0:
        pl.semaphore_wait(recv_sem.at[si], i)

[ultrons]

No race: pltpu.make_async_remote_copy(...).wait() waits on both semaphores — the send semaphore and the local recv_sem (the split wait_send() / wait_recv() forms exist precisely because plain wait() covers both). So at step i, prev[li].wait() blocks on recv_sem.at[si], which is signaled exactly by the step i-1 incoming block from the ring neighbor — the block step i forwards. The signal/wait accounting balances: one arrival per step per lane, one recv-wait consumed per step, and the trailing for d in prev: d.wait() consumes the final arrival. Adding an extra pl.semaphore_wait(recv_sem.at[si], i) would double-count the arrivals and deadlock. The kernel is validated bit-exact against jax.lax.all_gather(..., tiled=True) on v7x hardware from 8 to 512 chips.

[gemini-code-assist]

The local shard copy is defined and created twice using pltpu.make_async_copy. The first call starts the copy but is never waited on, while the second call creates a redundant copy operation and waits on it. This results in duplicate DMA transfers of the same local shard, wasting memory bandwidth and potentially causing race conditions.

Instead, create the AsyncCopy descriptor once, start it, and then wait on it.

  op = pltpu.make_async_copy(
      w_ref.at[_full_index(w_ref, gather_dim, 0, chunk)],
      o_ref.at[_full_index(o_ref, gather_dim, my_start, chunk)],
      local_sem)
  op.start()
  op.wait()

[ultrons]

There is no duplicate transfer — make_async_copy(...) builds a descriptor and .wait() only waits the semaphore (it does not issue a DMA), so the second construction was a wait-by-reconstruction, not a second copy. That said, the single-object form is clearer; applied in the latest push.