Fix fp8 MoE on the sparse_matmul path

src/maxtext/layers/moe.py

@@ -1499,12 +1499,14 @@ def get_tokamax_group_sizes(group_sizes, inputs, _kernel):
         )
 
     def get_quantization_dtypes():
-      lhs_quantize_dtype, rhs_quantize_dtype = None, None
-      if self.quant is not None:
-        quant_dg = self.quant.quant_dg
-        lhs_quantize_dtype = quant_dg.fwd.dg_quantizer.lhs.numerics.get_dtype()
-        rhs_quantize_dtype = quant_dg.fwd.dg_quantizer.rhs.numerics.get_dtype()
-      return lhs_quantize_dtype, rhs_quantize_dtype
+      # AQT describes its numerics through a `quant_dg`, while the fp8 schemes just name the
+      # dtype they quantize to. Under qwix there is no quantization object here at all, and
+      # the gmm takes its rule from the qwix rule instead.
+      quant_dg = getattr(self.quant, "quant_dg", None)
+      if quant_dg is not None:
+        return quant_dg.fwd.dg_quantizer.lhs.numerics.get_dtype(), quant_dg.fwd.dg_quantizer.rhs.numerics.get_dtype()
+      quantize_dtype = getattr(self.quant, "quantize_dtype", None)
+      return quantize_dtype, quantize_dtype
 
     def gmm(
         inputs,

src/maxtext/layers/quantizations.py

@@ -307,6 +307,9 @@ class Fp8Quantization(Quantization):
   """Configures Fp8 quantization for NVIDIA GPUs"""
 
   quant_mode = "train"
+  # The forward dtype, for callers that quantize an operand themselves rather than through
+  # `dot_general_cls` or `einsum`, such as the grouped matmul in MoE.
+  quantize_dtype = jnp.float8_e4m3fn
 
   def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Returns dot_general configured with aqt params."""
@@ -396,6 +399,7 @@ class NANOOFp8Quantization(Quantization):
   """Configures NANOO Fp8 quantization for AMD MI300/MI325 GPUs"""
 
   quant_mode = "train"
+  quantize_dtype = jnp.float8_e4m3fnuz
 
   def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Returns dot_general configured with aqt params."""

tests/integration/train_tests.py

@@ -76,17 +76,37 @@ class TrainTests(unittest.TestCase):
       "sharding_tolerance=0.1",
   ]
 
-  # Routes the MoE layer through dense_matmul, which is what runs wherever the megablox and
-  # ragged kernels are unavailable.
-  _moe_model_overrides = [
+  _moe_expert_overrides = [
       "decoder_block=mixtral",
       "num_experts=4",
       "num_experts_per_tok=2",
       "base_moe_mlp_dim=32",
+  ]
+
+  # Routes the MoE layer through dense_matmul, which is what runs wherever the megablox and
+  # ragged kernels are unavailable.
+  _moe_model_overrides = _moe_expert_overrides + [
       "sparse_matmul=False",
       "megablox=False",
   ]
 
+  # The sparse_matmul path, where megablox builds and uses the gmm quantization rule for real
+  # rather than falling back to ragged_dot.
+  _moe_sparse_model_overrides = _moe_expert_overrides + [
+      "sparse_matmul=True",
+      "megablox=True",
+  ]
+
+  # Every operand has to divide into its tile, and the default tiles are far larger than a
+  # downscaled model. The embedding dim is 28 or 32 depending on the device count, so 4 is
+  # the largest tile that fits it either way.
+  _megablox_tile_overrides = [
+      f"{matrix}_tile_{direction}_{dim}={size}"
+      for matrix in ("wi", "wo")
+      for direction in ("fwd", "dlhs", "drhs")
+      for dim, size in (("batch_seq", 16), ("embed_dim", 4), ("mlp_dim", 16))
+  ]
+
   _qwen3_overrides = [
       "override_model_config=True",
       "base_num_decoder_layers=2",
@@ -195,6 +215,20 @@ class TrainTests(unittest.TestCase):
       ]
       + _small_model_overrides
       + _moe_model_overrides,
+      "moe_sparse": [  # tests a MoE model on the sparse_matmul path, to be combined with a quantization
+          None,
+          get_test_config_path(),
+          f"base_output_directory={_base_output_directory}",
+          "run_name=runner_test",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
+          "steps=2",
+          "enable_checkpointing=False",
+          "enable_goodput_recording=False",
+          rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}",
+      ]
+      + _small_model_overrides
+      + _moe_sparse_model_overrides
+      + _megablox_tile_overrides,
       "te_fp8_delayedscaling": [  # tests base config with te_fp8_delayedscaling
           None,
           get_test_config_path(),
@@ -368,6 +402,22 @@ def test_moe_fp8_token_dropping(self):
     # capacity_factor > 0 adds the dispatch and combine einsums to the ones above.
     train_main(TrainTests.CONFIGS["moe"] + ["quantization=fp8", "capacity_factor=1.25"])
 
+  # The sparse_matmul tests below carry no hardware marker for the same reasons. What they cover
+  # is an attribute read during tracing rather than anything a kernel does, and megablox runs the
+  # quantized grouped matmul on CPU through its interpret mode.
+  @pytest.mark.integration_test
+  def test_moe_fp8_sparse_matmul(self):
+    train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=fp8"])
+
+  @pytest.mark.integration_test
+  def test_moe_nanoo_fp8_sparse_matmul(self):
+    train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=nanoo_fp8"])
+
+  # int8 takes the `quant_dg` branch of the same read, which the fp8 tests never reach.
+  @pytest.mark.integration_test
+  def test_moe_int8_sparse_matmul(self):
+    train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=int8"])
+
   @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
   @pytest.mark.integration_test
   @pytest.mark.gpu_only

tests/unit/moe_test.py

@@ -1936,6 +1936,68 @@ def test_quantized_dense_matmul_tracks_unquantized(self, quantization):
     self.assertFalse(np.allclose(actual, expected))
 
 
+class SparseMatmulQuantizationTest(parameterized.TestCase):
+  """The dtypes RoutedMoE hands the grouped matmul on the sparse_matmul path."""
+
+  @parameterized.named_parameters(
+      ("fp8", "fp8", False, jnp.float8_e4m3fn),
+      ("nanoo_fp8", "nanoo_fp8", False, jnp.float8_e4m3fnuz),
+      ("int8", "int8", False, jnp.int8),
+      ("unquantized", "", False, None),
+      # Under qwix a non-fp8_full scheme declares no gmm rule, so the gmm is left alone.
+      ("fp8_under_qwix", "fp8", True, None),
+      ("int8_under_qwix", "int8", True, None),
+  )
+  def test_gmm_quantize_dtypes(self, quantization, use_qwix_quantization, expected_dtype):
+    cfg = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="sparse_matmul_quantization_test",
+        enable_checkpointing=False,
+        decoder_block="mixtral",
+        num_experts=4,
+        num_experts_per_tok=2,
+        base_emb_dim=64,
+        base_mlp_dim=32,
+        base_moe_mlp_dim=32,
+        dtype="float32",
+        weight_dtype="float32",
+        sparse_matmul=True,
+        megablox=True,
+        max_target_length=8,
+        per_device_batch_size=1,
+        quantization=quantization,
+        use_qwix_quantization=use_qwix_quantization,
+    )
+    devices_array = maxtext_utils.create_device_mesh(cfg)
+    model = moe.RoutedMoE(
+        config=cfg,
+        num_experts=cfg.num_experts,
+        num_experts_per_tok=cfg.num_experts_per_tok,
+        mesh=Mesh(devices_array, cfg.mesh_axes),
+        kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+        kernel_axes=("embed", "mlp"),
+        dtype=jnp.float32,
+        quant=configure_quantization(cfg),
+        rngs=nnx.Rngs(0),
+    )
+
+    calls = []
+
+    def record_gmm(**kwargs):
+      calls.append(kwargs)
+      return jnp.zeros((kwargs["lhs"].shape[0], kwargs["rhs"].shape[-1]), dtype=kwargs["preferred_element_type"])
+
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (1, 8, cfg.base_emb_dim), dtype=jnp.float32)
+    with mock.patch.object(moe.mblx, "gmm", record_gmm):
+      with nn_partitioning.axis_rules(cfg.logical_axis_rules):
+        model(inputs)
+
+    self.assertNotEmpty(calls)
+    for kwargs in calls:
+      self.assertEqual(kwargs["lhs_quantize_dtype"], expected_dtype)
+      self.assertEqual(kwargs["rhs_quantize_dtype"], expected_dtype)
+
+
 def make_moe(cfg, mesh):
   return moe.RoutedMoE(
       config=cfg,