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

src/maxtext/configs/types.py

@@ -884,6 +884,35 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use Ring of Experts for sparse matmul expert parallelism.",
   )
+  moe_bwd_inkernel_quant: bool = Field(
+      False,
+      description=(
+          "Quantize the MoE backward-gmm operands INSIDE the ragged kernels instead of with dense "
+          "XLA-level quantize ops. The dense quantize/amax ops process every row of the ragged "
+          "buffer at its STATIC size (worst-case at ragged_buffer_factor<=0); the in-kernel path "
+          "touches only the valid group_sizes rows/tiles, and no reduction ever reads "
+          "uninitialized buffer tail rows. drhs: tgmm quantizes BOTH operands in-kernel "
+          "(per-gm-tile-per-channel e4m3). Requires use_tokamax_gmm + use_gmm_v2 + an fp8 qwix "
+          "bwd_qtype; falls back to the XLA quantize otherwise. For performance improvement at "
+          "dropless (worst-case) ragged buffer sizes."
+      ),
+  )
+  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'."
+      ),
+  )
+  moe_ring_cotangent_ag: bool = Field(
+      False,
+      description=(
+          "Run the BACKWARD cotangent all-gather of the ring-of-experts combine reduce-scatter on "
+          "a TensorCore Pallas ring kernel instead of the XLA collective (which can serialize on "
+          "the SparseCore collective-offload queue). Forward is unchanged. For performance "
+          "improvement; numerically equal to lax.all_gather."
+      ),
+  )
   moe_quantize_token_all_gather: bool = Field(
       False,
       description="Whether to quantize token activations to FP8 before All-Gather across EP shards in Ring of Experts.",
@@ -1321,6 +1350,15 @@ class RematAndOffload(BaseModel):
       RematLocation.REMAT,
       description="Remat policy for the first part of a gated MoE's output.",
   )
+  moe_x_sorted: RematLocation = Field(
+      RematLocation.REMAT,
+      description=(
+          "Remat policy for the routed (post-dispatch, expert-sorted) MoE input plus its small "
+          "routing/metadata bundle. 'device' saves them across the remat boundary so the backward "
+          "does not re-run the dispatch token all-gather and sort; the expert GMMs re-run from the "
+          "saved tensor. Default 'remat' recomputes (existing behavior)."
+      ),
+  )
   moe_mlpwi_1: RematLocation = Field(
       RematLocation.REMAT,
       description="Remat policy for the second part of a gated MoE's output.",
@@ -3423,6 +3461,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "context",
           "mlpwi",
           "moe_mlpwi_0",
+          "moe_x_sorted",
           "moe_mlpwi_1",
           "moe_mlpwo",
           "mlpwi_0",
@@ -4564,6 +4603,7 @@ def set_derived_values_and_validate(self) -> "RLConfig":
           "context",
           "mlpwi",
           "moe_mlpwi_0",
+          "moe_x_sorted",
           "moe_mlpwi_1",
           "moe_mlpwo",
           "mlpwi_0",

src/maxtext/kernels/megablox/ops.py

@@ -76,6 +76,7 @@ def gmm(
     use_manual_quantization: bool = False,  # used in batchsplit
     use_gmm_v2: bool = False,
     partial_sum: jnp.ndarray | None = None,
+    bwd_inkernel_quant: bool = False,
 ):
   """Grouped matrix multiplication operation."""
   if interpret is None:
@@ -106,7 +107,7 @@ def gmm(
   gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0]  # pylint: disable=C3001
   gmm_fwd_bwd = jax.custom_vjp(
       gmm_fwd_bwd,
-      nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15),
+      nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 17),
   )
   gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype))
   return gmm_fwd_bwd(
@@ -127,6 +128,7 @@ def gmm(
       rhs_vma_axes,
       use_gmm_v2,
       partial_sum,
+      bwd_inkernel_quant,
   )
 
 
@@ -163,6 +165,7 @@ def _gmm_fwd(
     rhs_vma_axes: tuple = tuple(),
     use_gmm_v2: bool = False,
     partial_sum: jnp.ndarray | None = None,
+    bwd_inkernel_quant: bool = False,
 ) -> tuple[
     jnp.ndarray,
     tuple[
@@ -462,6 +465,7 @@ def _gmm_bwd(
     lhs_vma_axes: tuple,
     rhs_vma_axes: tuple,
     use_gmm_v2: bool,
+    bwd_inkernel_quant: bool,
     residual: tuple[
         jnp.ndarray | qpl.QArray,
         jnp.ndarray | qpl.QArray,
@@ -493,14 +497,44 @@ def _gmm_bwd(
   #  - dlhs_dout: the incoming gradient used to calculate dlhs.
   #  - drhs_dout: the incoming gradient used to calculate drhs.
 
+  # moe_bwd_inkernel_quant: run the drhs tgmm with BOTH operands quantized in-kernel
+  # (per-gm-tile-per-channel), touching only rows covered by group_sizes. This subsumes three
+  # dense buffer-sized XLA ops -- the per-row lhs re-quantize, the drhs_dout *= lhs.scale
+  # multiply, and the per-N cotangent quantize -- and, because no reduction ever reads past the
+  # valid rows, removes the NaN hazard of amax over uninitialized ragged-buffer tail rows.
+  inkernel_drhs = (
+      bwd_inkernel_quant
+      and use_tokamax_backend
+      and use_gmm_v2
+      and quantization_rule is not None
+      and bool(quantization_rule.bwd_qtype)
+  )
+
   # 1. Scale Application & QArray Unwrapping
   dlhs_dout, drhs_dout, lhs, rhs = _bwd_prepare_inputs(
-      grad, residual_lhs, residual_rhs, group_sizes, use_gmm_v2, transpose_rhs, quantization_rule
+      grad, residual_lhs, residual_rhs, group_sizes, use_gmm_v2, transpose_rhs, quantization_rule,
+      skip_lhs_quant=inkernel_drhs,
   )
 
   # 2. Backward Pass Quantization
   if quantization_rule:
-    dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule)
+    if inkernel_drhs:
+      # the in-kernel tgmm quantizes drhs_dout ITSELF; quantize ONLY the dlhs cotangent here
+      # (calling the two-sided helper would emit the dense drhs quantize just to discard it).
+      if quantization_rule.bwd_qtype:
+        dlhs_dout = qpl.quantize(
+            # pyrefly: ignore[bad-argument-type]
+            dlhs_dout,
+            quantization_rule.bwd_qtype,
+            channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
+            calibration_method=quantization_rule.bwd_calibration_method,
+        )
+      if not isinstance(drhs_dout, qpl.QArray) and drhs_dout.dtype != lhs.dtype:
+        # tgmm requires equal operand widths; the in-kernel path reads the RAW cotangent, so
+        # carry it at the activation width (halves the kernel's cotangent read bytes vs f32).
+        drhs_dout = drhs_dout.astype(lhs.dtype)
+    else:
+      dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule)
 
   # 3. DLHS Gradient Execution
   dlhs = _compute_dlhs(
@@ -534,6 +568,7 @@ def _gmm_bwd(
       interpret,
       rhs_vma_axes,
       quantization_rule,
+      inkernel_quant=inkernel_drhs,
   )
 
   # 5. Output Formatting
@@ -572,8 +607,14 @@ def _bwd_prepare_inputs(
     use_gmm_v2: bool,
     transpose_rhs: bool,
     quantization_rule: qwix.QtRule | None,
+    skip_lhs_quant: bool = False,
 ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, jnp.ndarray, jnp.ndarray]:
-  """Prepares backward operands."""
+  """Prepares backward operands.
+
+  `skip_lhs_quant=True` (bwd_inkernel_quant) keeps the lhs as the raw wide array: the drhs
+  tgmm quantizes BOTH operands in-kernel over valid gm tiles only, so the dense per-row XLA
+  quantize here (and the drhs_dout *= lhs.scale multiply below) would be buffer-sized overhead.
+  """
 
   # dlhs_dout and drhs_dout can be different when quantization is enabled.
   dlhs_dout = grad
@@ -597,7 +638,7 @@ def _bwd_prepare_inputs(
 
   # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype
   # in the residual tuple. In BWD, we explicitly quantize lhs.
-  if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray):
+  if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not skip_lhs_quant:
     lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
         lhs,
         quantization_rule.act_qtype,
@@ -816,12 +857,16 @@ def _compute_drhs(
     interpret: bool,
     rhs_vma_axes: tuple,
     quantization_rule: qwix.QtRule | None,
+    inkernel_quant: bool = False,
 ) -> jnp.ndarray:
   """Routes execution of DRHS based on backend choices."""
   if use_tokamax_backend and not use_gmm_v2:
     drhs = _drhs_run_tokamax_v1(drhs_dout, lhs, group_sizes, rhs_dtype, use_manual_quantization)
   elif use_tokamax_backend and use_gmm_v2:
-    drhs = _drhs_run_tokamax_v2(drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling)
+    drhs = _drhs_run_tokamax_v2(
+        drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling,
+        quantize_operands=inkernel_quant,
+    )
   else:
     drhs = _drhs_run_megablox(
         drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, interpret, rhs_vma_axes
@@ -888,6 +933,7 @@ def _drhs_run_tokamax_v2(
     num_actual_groups: int,
     rhs_dtype: jax.typing.DTypeLike,
     tiling: tuple,
+    quantize_operands: bool = False,
 ) -> jnp.ndarray:
   """Executes Tokamax TGMM V2 backend for DRHS = LHS^T @ DRHS_dout."""
   drhs_rhs = drhs_dout.qvalue if isinstance(drhs_dout, qpl.QArray) else drhs_dout
@@ -909,6 +955,7 @@ def _drhs_run_tokamax_v2(
       preferred_element_type=rhs_dtype,  # pyrefly: ignore[bad-argument-type]
       group_offset=group_offset,
       tile_info=custom_drhs_tiling,
+      quantize_operands=quantize_operands and rhs_scale is None,
   )
 
 

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

@@ -191,6 +191,7 @@ def make_tgmm_configs(
     out_dtype: jnp.dtype,
     acc_dtype: jnp.dtype | None,
     target_zero_ref_bytes: int,
+    quantize_operands: bool = False,
 ):
   """Fills the GMM config for the TGMM kernel."""
   assert out_dtype, "out_dtype cannot be None"
@@ -233,15 +234,24 @@ def make_tgmm_configs(
       size_lhs_sublane=size_lhs_sublane,
   )
 
+  # moe_bwd_inkernel_quant: quantize_operands puts a per-(gm-tile x channel) dynamic e4m3
+  # quantize of BOTH operands inside the inner kernel (signalled to tgmm_inner_kernel via
+  # lhs_cfgs.quant_dtype). Mutually exclusive with a pre-computed per-N rhs_scale.
+  if quantize_operands:
+    assert rhs_scale is None, "quantize_operands is mutually exclusive with rhs_scale"
+    inkernel_q_dtype = jnp.float8_e4m3fn.dtype
+  else:
+    inkernel_q_dtype = None
+
   rhs_quant_block_size_m = size_m
   rhs_cfgs = gmm_v2.InputConfigs(
-      quant_dtype=None,
+      quant_dtype=inkernel_q_dtype,
       quant_block_size=rhs_quant_block_size_m,
       dtype=rhs.dtype,
       has_scale=(rhs_scale is not None),
   )
   lhs_cfgs = gmm_v2.InputConfigs(
-      quant_dtype=None,
+      quant_dtype=inkernel_q_dtype,
       quant_block_size=-1,
       dtype=lhs.dtype,
   )
@@ -334,12 +344,43 @@ def _matmul(is_new_group: bool, is_group_changing: bool):
     rhs_mask = jnp.logical_and(m_start_local <= rhs_iota, rhs_iota < m_end_local)
     rhs_masked = jnp.where(rhs_mask, tiled_rhs_ref[...], 0)
 
-    acc = jax.lax.dot_general(
-        lhs_masked,
-        rhs_masked,
-        (((0,), (0,)), ((), ())),
-        preferred_element_type=jnp.float32,
-    )
+    if cfgs.lhs_cfgs.quant_dtype is not None:
+      # moe_bwd_inkernel_quant: quantize BOTH operands in-kernel with per-(gm-tile x channel)
+      # dynamic scales and rescale the partial product by the scale outer-product before
+      # accumulation (the tgmm_block per-segment pattern, fused into the tile loop -- no dense
+      # XLA-level quantize/amax over the ragged buffer, only valid gm tiles pay). The masked
+      # rows are zero, so they neither perturb the amax nor the product.
+      q_dtype = cfgs.lhs_cfgs.quant_dtype
+      dtype_max = float(jnp.finfo(q_dtype).max)
+      lhs_f = lhs_masked.astype(jnp.float32)
+      rhs_f = rhs_masked.astype(jnp.float32)
+      lhs_scale = jnp.max(jnp.abs(lhs_f), axis=0) / dtype_max  # [tile_k] f32
+      rhs_scale = jnp.max(jnp.abs(rhs_f), axis=0) / dtype_max  # [tile_n] f32
+      # A near-zero scale would give 0 * inf = NaN. An `== 0` guard is NOT enough: for any column
+      # whose amax is nonzero but below ~1.3e-36, `1/scale` OVERFLOWS f32 to inf, the guard does not
+      # fire, and every exactly-zero element in that column becomes 0*inf = NaN. Masked rows are set
+      # to exactly 0 above, so MORE masked rows = more NaN sites -- imbalanced (real) routing makes
+      # small groups and mostly-masked tiles, so it is strictly more exposed than a balanced
+      # synthetic router. Guard on the smallest scale with a finite reciprocal instead.
+      _recip_min = jnp.float32(1.0) / jnp.finfo(jnp.float32).max
+      lhs_inv = jnp.where(lhs_scale > _recip_min, 1.0 / lhs_scale, 0.0)
+      rhs_inv = jnp.where(rhs_scale > _recip_min, 1.0 / rhs_scale, 0.0)
+      lhs_q = (lhs_f * lhs_inv.reshape(1, -1)).astype(q_dtype)
+      rhs_q = (rhs_f * rhs_inv.reshape(1, -1)).astype(q_dtype)
+      acc = jax.lax.dot_general(
+          lhs_q,
+          rhs_q,
+          (((0,), (0,)), ((), ())),
+          preferred_element_type=jnp.float32,
+      )
+      acc = acc * lhs_scale.reshape(-1, 1) * rhs_scale.reshape(1, -1)
+    else:
+      acc = jax.lax.dot_general(
+          lhs_masked,
+          rhs_masked,
+          (((0,), (0,)), ((), ())),
+          preferred_element_type=jnp.float32,
+      )
 
     if not is_new_group:
       acc += acc_ref[...]
@@ -642,6 +683,7 @@ def validate_tgmm_inputs(
         "precision",
         "preferred_element_type",
         "acc_dtype",
+        "quantize_operands",
     ],
 )
 def tgmm_v2(
@@ -658,6 +700,7 @@ def tgmm_v2(
     precision: jax.lax.Precision = jax.lax.Precision.DEFAULT,
     preferred_element_type: jnp.dtype | None = None,
     acc_dtype: jnp.dtype | None = None,
+    quantize_operands: bool = False,
 ):
   """Computes a transposed grouped matrix multiplication.
 
@@ -710,6 +753,7 @@ def tgmm_v2(
       out_dtype=preferred_element_type,  # pyrefly: ignore[bad-argument-type]
       acc_dtype=acc_dtype,
       target_zero_ref_bytes=target_zero_ref_bytes,
+      quantize_operands=quantize_operands,
   )
   dims = cfgs.dims
   tiles = cfgs.tiles