MoE fp8 backward: in-kernel drhs operand quantization + configurable gradient dtype

[ultrons]

Stacked on #4895 (token quantized all-gather) and #4969 (ring collectives + x_sorted save); the base branch is #4895's head, so the diff shows this change plus #4969's commit — review the top commit.

What

Two default-off knobs for performance improvement of fp8 MoE training at dropless (worst-case) ragged buffer sizes:

• moe_bwd_inkernel_quant — quantize the drhs tgmm's operands inside the kernel (per-gm-tile-per-channel e4m3, implemented in tgmm_v2's quantize_operands path) instead of with dense XLA-level quantize ops. The dense quantize/amax ops process every row of the ragged buffer at its static size — at ragged_buffer_factor<=0 (dropless worst case) that is tokens * top_k rows regardless of how many are valid — while the in-kernel path touches only the tiles covered by group_sizes. It subsumes three buffer-sized ops: the per-row lhs re-quantize, the drhs_dout *= lhs.scale multiply, and the per-N cotangent quantize. A second effect: because no reduction ever reads past the valid rows, the amax can never ingest uninitialized buffer tail rows (with the dense path, a stale-row amax=Inf propagates Inf*0=NaN into the weight gradient; we root-caused a training NaN to exactly this).
• bwd_quantization_dtype — the gradient fp8 dtype for the fp8_full qwix recipe: e5m2 (default, existing behavior) or e4m3.

Requires use_tokamax_gmm=true use_gmm_v2=true and an fp8 qwix bwd_qtype; falls back to the XLA-level quantize otherwise. Flags off reproduce the existing behavior.

Repro

DeepSeek-V3 671B on v7x, the #4895 recipe at the dropless configuration this lever targets:

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=-1 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_bwd_inkernel_quant=true bwd_quantization_dtype=e4m3 \
  dataset_type=synthetic use_random_routing=true steps=20

• The win scales with the gap between the buffer's static size and the balanced size; at small ragged_buffer_factor (e.g. 2.0) the dense ops are cheap and the flag is roughly neutral. Use it for dropless (ragged_buffer_factor=-1) configurations.
• The in-kernel scales are finer-grained than the XLA path's, so losses match the flag-off run closely but not bit-exactly.

Validation

• AOT compile (virtual tpu7x topology) green: flag-on at the dropless per_device_batch_size=1.0 / ragged_buffer_factor=-1 config above, and flag-off on the unmodified Add token quantized all gather flow #4895 recipe.
• 512-chip DeepSeek-V3 runs with this path train cleanly at ragged_buffer_factor=-1 with no buffer sanitization and no non-finite losses; a 500-step real-data curve on this backward matches the dense-quantize reference at every eval point.

[gemini-code-assist]

Code Review

This pull request introduces several Mixture of Experts (MoE) performance optimizations, including in-kernel backward quantization for ragged kernels, a configurable backward quantization dtype, a rematerialization policy (moe_x_sorted) to save routed expert inputs, and a Pallas-based bidirectional ring all-gather kernel for backward cotangent reduce-scatters. The code review feedback highlights a critical synchronization bug in the ring all-gather kernel where a semaphore wait is needed to prevent forwarding stale data, a performance optimization to avoid redundant dense quantization of drhs_dout when in-kernel quantization is active, and a configuration improvement to use a Literal type for the backward quantization dtype.

[gemini-code-assist]

In a store-and-forward ring collective, each device forwards the block it received in the previous step. At step i (for i > 0), the block to be sent was written to the local o_ref by the neighbor in step i - 1. Without waiting for the local recv_sem to be signaled, a device can read and forward uninitialized or stale data, leading to silent data corruption. Adding a semaphore wait ensures correctness.

Suggested change

	for i in range(s - 1):
	for li, (si, nb, sign, ss, sl) in enumerate(lanes):
	start = block_start(my - i if sign < 0 else my + i)
	if prev[li] is not None:
	prev[li].wait()
	prev[li] = pltpu.async_remote_copy(
	o_ref.at[idx(start, ss, sl)], o_ref.at[idx(start, ss, sl)],
	send_sem.at[si], recv_sem.at[si], device_id=nb, device_id_type=MESH)
	for i in range(s - 1):
	for li, (si, nb, sign, ss, sl) in enumerate(lanes):
	start = block_start(my - i if sign < 0 else my + i)
	if i > 0:
	pl.semaphore_wait(recv_sem.at[si], i)
	if prev[li] is not None:
	prev[li].wait()
	prev[li] = pltpu.async_remote_copy(
	o_ref.at[idx(start, ss, sl)], o_ref.at[idx(start, ss, sl)],
	send_sem.at[si], recv_sem.at[si], device_id=nb, device_id_type=MESH)

[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]

Calling _bwd_quantize_gradient unconditionally quantizes both dlhs_dout and drhs_dout. Since the in-kernel path is designed to avoid dense XLA-level quantization of drhs_dout, calling this helper and discarding the second return value still executes the expensive dense quantization of drhs_dout. Directly quantizing dlhs_dout avoids this redundant computation.

Suggested change

	if inkernel_drhs:
	# the in-kernel tgmm quantizes drhs_dout ITSELF; keep the raw cotangent for the drhs side.
	dlhs_dout, _ = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule)
	if inkernel_drhs:
	# the in-kernel tgmm quantizes drhs_dout ITSELF; keep the raw cotangent for the drhs side.
	dlhs_dout = qpl.quantize(
	dlhs_dout,
	quantization_rule.bwd_qtype,
	channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
	calibration_method=quantization_rule.bwd_calibration_method,
	)

[ultrons]

Good catch — the discarded second output still traced the dense drhs quantize (XLA DCE removes it, but relying on DCE for the exact op this flag eliminates is fragile). Applied: the in-kernel branch now quantizes only the dlhs cotangent inline.

[gemini-code-assist]

Using a strict Literal type instead of a generic str ensures that invalid quantization data types are caught early during configuration validation.

Suggested change

	bwd_quantization_dtype: str = Field(
	"e5m2",
	description=(
	"fp8 dtype for the BACKWARD (gradient) quantization in the fp8_full qwix recipe: 'e5m2' "
	"(default) or 'e4m3'."
	),
	)
	bwd_quantization_dtype: Literal["e5m2", "e4m3"] = Field(
	"e5m2",
	description=(
	"fp8 dtype for the BACKWARD (gradient) quantization in the fp8_full qwix recipe: 'e5m2' "
	"(default) or 'e4m3'."
	),
	)

[ultrons]

Applied — Literal["e5m2", "e4m3"].