MoE: native FP8 QAT for routed-expert GMMs (FP8 forward + FP8-weight backward dLHS)

src/maxtext/configs/base.yml

@@ -347,6 +347,17 @@ fp8_moe_fake_quant: false
 # once for the scanned layer stack, before the scan. gather8: also all-gather the fp8 codes and dequantize after the
 # gather. Env PARAMBOLE_FQ_HOIST (0, absmax, 1 = gather8) selects a mode while this is off.
 fp8_fake_quant_hoist: "off"
+# With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8, run the routed-expert wi/wo forward GMMs in FP8 (gathered
+# codes + per-channel scales into gmm_v2; the activation is quantized in the kernel with one float32 absmax scale per
+# row over the full K dim, and accumulated in float32). Backward: bf16 DLHS / TGMM at the dequantized weight, unless
+# fp8_moe_native_bwd_dlhs dequantizes it inside the DLHS kernel.
+# Needs sparse_matmul, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, use_tokamax_gmm, use_gmm_v2 and fixed tiles
+# (use_gmm_v2_heuristic_tiling=false), and runs only inside the decoder's layer scan (where gather8's absmax is
+# attached): a routed-MoE call outside it raises.
+fp8_moe_native_fwd: false
+# fp8_moe_native_fwd: the backward DLHS GMMs read the fp8 weight codes and dequantize them in VMEM inside gmm_v2
+# (bit-identical; no XLA dequantization of the expert weights). Needs moe_gmm_v2_dlhs_transpose_rhs and fixed tiles.
+fp8_moe_native_bwd_dlhs: false
 # Trainer side: the rollout serves fp8_moe experts, so weight sync quantizes them (and emits their scales).
 rollout_fp8_moe: false
 load_balance_loss_weight: 0.0 # weight for the load balance loss

src/maxtext/configs/types.py

@@ -1937,6 +1937,66 @@ def _off_for_false_fp8_fake_quant_hoist(v: Any) -> Any:
       _off_for_false_fp8_fake_quant_hoist
   )
 
+  fp8_moe_native_fwd: bool = Field(
+      False,
+      description=(
+          "Native FP8 forward for the routed experts. With fp8_moe_fake_quant and fp8_fake_quant_hoist=gather8,"
+          " the wi and wo forward GMMs run on the MXU in float8_e4m3fn: the gathered weight codes and their"
+          " per-channel scales feed gmm_v2 directly (no dequantization), and gmm_v2 quantizes the bf16 activation"
+          " in the kernel with one dynamic float32 absmax scale per row over the full K dim (the rollout FP8 MoE kernel's"
+          " decode scheme) and accumulates in float32. The backward runs"
+          " bf16 DLHS / TGMM at the dequantized weight (straight-through weight gradient), unless"
+          " fp8_moe_native_bwd_dlhs moves the DLHS dequantization into the kernel. Runs only inside the decoder's"
+          " layer scan, where gather8's absmax is attached; a RoutedMoE call outside it (unscanned layers, KV-cache"
+          " decode) raises instead of silently running fake QAT."
+      ),
+  )
+  fp8_moe_native_bwd_dlhs: bool = Field(
+      False,
+      description=(
+          "fp8_moe_native_fwd: the backward DLHS GMMs (wi, and wo's fused-dSwiGLU DLHS, including the ring-of-experts"
+          " one) read the float8 weight codes and dequantize each weight block in VMEM inside gmm_v2, with the same"
+          " expression (codes.astype(float32) * scale).astype(bf16), instead of an XLA dequantization of every"
+          " expert weight before the kernel. Bit-identical results; the weight-gradient TGMMs are unchanged. Needs"
+          " moe_gmm_v2_dlhs_transpose_rhs and fixed tiles (use_gmm_v2_heuristic_tiling=False)."
+      ),
+  )
+
+  @model_validator(mode="after")
+  def validate_fp8_moe_native_fwd(self) -> "MoEGeneral":
+    """Rejects fp8_moe_native_fwd / fp8_moe_native_bwd_dlhs without the config they require.
+
+    The FP8 GMMs replace the gather8 dequantization on the fused gmm_v2 SwiGLU / dSwiGLU path only and need fixed
+    tiles (float32 accumulator); the in-kernel DLHS dequantization also needs the transposed DLHS rhs.
+    """
+    if self.fp8_moe_native_fwd:
+      _reject_unless(
+          "fp8_moe_native_fwd=True",
+          [
+              (self.fp8_moe_fake_quant, "fp8_moe_fake_quant=True"),
+              (self.fp8_fake_quant_hoist == "gather8", "fp8_fake_quant_hoist=gather8"),
+              (self.sparse_matmul, "sparse_matmul=True"),
+              (self.use_tokamax_gmm and self.use_gmm_v2, "use_tokamax_gmm=True and use_gmm_v2=True"),
+              (self.moe_wi_gmm_fused_swiglu, "moe_wi_gmm_fused_swiglu=True"),
+              (self.moe_wo_gmm_fused_dswiglu, "moe_wo_gmm_fused_dswiglu=True"),
+              # gmm_v2's heuristic tiling budgets a bfloat16 accumulator when the lhs is quantized in the kernel,
+              # but the native forward accumulates in float32, so heuristic tiles could exceed VMEM.
+              (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"),
+          ]
+          + self._fp8_moe_native_fwd_unsplit_k_requirements()
+          + self._fp8_moe_native_fwd_tile_requirements(),
+      )
+    if self.fp8_moe_native_bwd_dlhs:
+      _reject_unless(
+          "fp8_moe_native_bwd_dlhs=True",
+          [
+              (self.fp8_moe_native_fwd, "fp8_moe_native_fwd=True"),
+              (self.moe_gmm_v2_dlhs_transpose_rhs, "moe_gmm_v2_dlhs_transpose_rhs=True"),
+              (not self.use_gmm_v2_heuristic_tiling, "use_gmm_v2_heuristic_tiling=False"),
+          ],
+      )
+    return self
+
   rollout_fp8_moe: bool = Field(
       False,
       description=(
@@ -2098,6 +2158,115 @@ def _moe_mlp_dim_if_unsharded(self) -> int | None:
       return None
     return 2 ** get_individual_scales(scale)[2] * base
 
+  def _moe_emb_dim_if_unsharded(self) -> int | None:
+    """The routed experts' input dim (wi's K dim) if no mesh axis shards it in the sparse_matmul GMMs, else None."""
+    if self.use_2d_fsdp_sharding or getattr(self, "num_moe_emb_chunks", 0) > 0:
+      return None
+    axes = {"tensor"}
+    axes.update(axes_for_logical(self.logical_axis_rules or [], "activation_embed_moe"))
+    for axis in axes:
+      if getattr(self, f"ici_{axis}_parallelism", None) != 1 or getattr(self, f"dcn_{axis}_parallelism", None) != 1:
+        return None
+    if self.moe_expert_input_dim > 0:
+      return self.moe_expert_input_dim
+    base, scale = self.base_emb_dim, self.global_parameter_scale
+    if base <= 0 or scale <= 0 or scale & (scale - 1):
+      return None
+    return 2 ** get_individual_scales(scale)[0] * base
+
+  def _fp8_moe_native_fwd_unsplit_k_requirements(self) -> list[tuple[bool, str]]:
+    """fp8_moe_native_fwd scales each activation row over the whole K dim, so neither forward GMM may see a K slice.
+
+    wi contracts the embed dim and wo the MoE MLP dim. A mesh axis that splits one of them in the expert shard_map, or
+    num_moe_emb_chunks, would make each shard or chunk quantize its slice with its own scale, and make_gmm_configs
+    cannot tell: it only sees the local size_k. The axes are the ones _moe_emb_dim_if_unsharded and
+    _moe_mlp_dim_if_unsharded read, under the training rules and separate eval rules. Parallelism -1 is filled in from
+    the device count later and custom_mesh_and_rule replaces logical_axis_rules after this validator, so RoutedMoE
+    checks the traced weight partition specs too.
+    """
+    rules = [self.logical_axis_rules or []]
+    eval_axis_rules = getattr(self, "logical_axis_rules_for_eval", None)
+    if eval_axis_rules:
+      rules.append(eval_axis_rules)
+
+    def split_by(axes: set[str], logical: str) -> list[str]:
+      for rule in rules:
+        axes = axes | set(axes_for_logical(rule, logical))
+      return [
+          f"{kind}_{axis}_parallelism={degree}"
+          for axis in sorted(axes)
+          for kind in ("ici", "dcn")
+          if (degree := getattr(self, f"{kind}_{axis}_parallelism", None)) is not None and degree > 1
+      ]
+
+    emb_split = split_by({"tensor"}, "activation_embed_moe")
+    mlp_split = split_by({"tensor", "tensor_sequence", "autoregressive"}, "mlp_no_fsdp")
+    emb_chunks = getattr(self, "num_moe_emb_chunks", 0)
+    return [
+        (not emb_split, f"an unsplit embed dim, wi's K (got {', '.join(emb_split)})"),
+        (not mlp_split, f"an unsplit MoE MLP dim, wo's K (got {', '.join(mlp_split)})"),
+        (emb_chunks <= 0, f"num_moe_emb_chunks=0, an unchunked wi K (got {emb_chunks})"),
+    ]
+
+  def _fp8_moe_native_fwd_tile_requirements(self) -> list[tuple[bool, str]]:
+    """Tile rules for fp8_moe_native_fwd's full-K (lhs_quant_block_size=0) forward GMMs, checked at config time.
+
+    gmm_v2 quantizes one [tile_m, tile_k] lhs slice per K step, so full-K dynamic scaling needs tile_k == size_k in
+    both forward GMMs (wi contracts emb_dim, wo contracts the MoE MLP dim); make_gmm_configs enforces the same rule
+    when traced. When separate eval logical axis rules are set, RoutedMoE's fwd_tile reads eval_wi_tile_fwd_embed_dim
+    and eval_wo_tile_fwd_mlp_dim in eval_step, so those are checked too.
+    """
+    if self.use_gmm_v2_heuristic_tiling:
+      return []
+    eval_rule = getattr(self, "custom_mesh_and_rule_for_eval", CustomRule.DEFAULT)
+    eval_axis_rules = getattr(self, "logical_axis_rules_for_eval", None)
+    separate_eval_rules = eval_rule not in (
+        CustomRule.DEFAULT,
+        getattr(self, "custom_mesh_and_rule", CustomRule.DEFAULT),
+    ) or bool(eval_axis_rules and eval_axis_rules != self.logical_axis_rules)
+    reqs: list[tuple[bool, str]] = []
+    emb = self._moe_emb_dim_if_unsharded()
+    if emb is not None:
+      reqs.append(
+          (
+              self.wi_tile_fwd_embed_dim == emb,
+              f"wi_tile_fwd_embed_dim={emb}, the embed dim (one K tile for full-K activation quantization;"
+              f" got {self.wi_tile_fwd_embed_dim})",
+          )
+      )
+      if separate_eval_rules and all(
+          getattr(self, f"ici_{ax}_parallelism", None) == 1 and getattr(self, f"dcn_{ax}_parallelism", None) == 1
+          for ax in axes_for_logical(eval_axis_rules or [], "activation_embed_moe")
+      ):
+        reqs.append(
+            (
+                self.eval_wi_tile_fwd_embed_dim == emb,
+                f"eval_wi_tile_fwd_embed_dim={emb} when evaluation uses separate logical axis rules"
+                f" (got {self.eval_wi_tile_fwd_embed_dim})",
+            )
+        )
+    n = self._moe_mlp_dim_if_unsharded()
+    if n is not None:
+      reqs.append(
+          (
+              self.wo_tile_fwd_mlp_dim == n,
+              f"wo_tile_fwd_mlp_dim={n}, the MoE MLP dim (one K tile for full-K activation quantization;"
+              f" got {self.wo_tile_fwd_mlp_dim})",
+          )
+      )
+      if separate_eval_rules and all(
+          getattr(self, f"ici_{ax}_parallelism", None) == 1 and getattr(self, f"dcn_{ax}_parallelism", None) == 1
+          for ax in axes_for_logical(eval_axis_rules or [], "mlp_no_fsdp")
+      ):
+        reqs.append(
+            (
+                self.eval_wo_tile_fwd_mlp_dim == n,
+                f"eval_wo_tile_fwd_mlp_dim={n} when evaluation uses separate logical axis rules"
+                f" (got {self.eval_wo_tile_fwd_mlp_dim})",
+            )
+        )
+    return reqs
+
   @model_validator(mode="after")
   def validate_moe_roe_packed_routing_gather(self) -> "MoEGeneral":
     """The packed gather replaces the per-chunk routing gathers of the pipelined local-routing chunk loop."""

src/maxtext/kernels/megablox/__init__.py

@@ -15,9 +15,13 @@
 
 from maxtext.kernels.megablox.ops import (
     TcUnsortCfg,
+    dequantize_fp8_codes,
     gmm,
     gmm_dswiglu,
     gmm_dswiglu_bwd,
+    gmm_dswiglu_fp8,
+    gmm_fp8_forward,
     gmm_silu,
+    gmm_silu_fp8,
     gmm_silu_pre_act_tile_n,
 )