Dev rocm mx type

src/maxtext/configs/base.yml

@@ -134,6 +134,7 @@ dtype: "bfloat16"
 # 'nanoo_fp8' for 8-bit floating-point gemms on amd mi300/mi325 gpus.
 # 'fp8_full' for fp8 quantization with static scaling.
 # 'serve_fp8_weight' to load and use FP8 checkpoint directly at inference time.
+# 'jax_mxfp8' / 'jax_mxfp4' for OCP MX training through jax.nn.scaled_dot_general.
 quantization: ""
 # used to configure constant_bound_config in aqt lib for static scaling, e.g. constant_bound_config='0.5, 0.5, 0.5, 0.5, 0.5, 0.5'
 constant_bound_config: ""

src/maxtext/configs/types.py

@@ -103,6 +103,8 @@ class QuantizationType(str, Enum):
   TE_FP8_DS = "te_fp8_delayedscaling"
   TE_FP8_CS = "te_fp8_currentscaling"
   TE_MXFP8 = "te_mxfp8"
+  JAX_MXFP8 = "jax_mxfp8"
+  JAX_MXFP4 = "jax_mxfp4"
   TE_NVFP4 = "te_nvfp4"
   TE_NVFP4_NO_RHT = "te_nvfp4_no_rht"
   SERVE_FP8_WEIGHT = "serve_fp8_weight"
@@ -4448,6 +4450,16 @@ def validate_shard_embed_moe_on_fsdp(self) -> "MaxTextConfig":
       )
     return self
 
+  def _validate_jax_mx_quantization(self) -> None:
+    """Rejects jax_mxfp8 and jax_mxfp4 settings that would silently leave GEMMs unquantized or fail obscurely in JAX."""
+    name = self.quantization.value
+    if self.use_qwix_quantization:
+      raise ValueError(f"quantization='{name}' does not support use_qwix_quantization=True.")
+    if self.shard_mode == ShardMode.EXPLICIT:
+      raise ValueError(f"quantization='{name}' does not support shard_mode='explicit'.")
+    if self.num_experts > 1:
+      raise ValueError(f"quantization='{name}' does not support MoE models (num_experts > 1).")
+
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
     """
@@ -4744,6 +4756,8 @@ def get_num_target_devices():
       if self.quantization not in (
           "fp8",
           "nanoo_fp8",
+          "jax_mxfp8",
+          "jax_mxfp4",
       ) and not self.quantization.startswith("te_"):
         logger.warning(
             "WARNING: AQT quantization is deprecated and will be removed in a future release. "
@@ -6068,6 +6082,9 @@ def _stream_chunked(tokenize_data: bool) -> bool:
             "counterparts) to 1, or turn off `shard_optimizer_over_data`."
         )
 
+    if self.quantization in ("jax_mxfp8", "jax_mxfp4"):
+      self._validate_jax_mx_quantization()
+
     # Diloco params
     # Resolve dcn_diloco_parallelism=-1 if left unspecified, using the same convention as dcn_data_parallelism.
     # num_diloco_replicas must be computed after this resolution, so we resolve it here rather than

src/maxtext/layers/quantizations.py

@@ -162,6 +162,45 @@ def einsum(self, dtype: DType = jnp.float32):
     """Placeholder for einsum implementation in subclasses."""
 
 
+class JaxBlockScaledDotGeneral(nnx.Module):
+  """Stateless OCP MX dot_general: `jax.nn.scaled_dot_general` with one config for both operands and the gradient."""
+
+  def __init__(self, mode: str):
+    self.mode = mode
+
+  def __call__(
+      self,
+      lhs: jax.Array,
+      rhs: jax.Array,
+      dimension_numbers: jax.lax.DotDimensionNumbers,
+      precision: jax.lax.PrecisionLike = None,
+      preferred_element_type: jax.typing.DTypeLike | None = None,
+  ) -> jax.Array:
+    # jax.nn.scaled_dot_general takes no precision argument.
+    del precision
+    output_dtype = jnp.result_type(lhs, rhs) if preferred_element_type is None else preferred_element_type
+    config = jax.nn.get_scaled_dot_general_config(self.mode)
+    return jax.nn.scaled_dot_general(
+        lhs, rhs, dimension_numbers, preferred_element_type=output_dtype, configs=[config, config, config]
+    )
+
+
+@dataclass
+class JaxBlockScaledQuantization(Quantization):
+  """Configures OCP MX quantization through `jax.nn.scaled_dot_general`."""
+
+  quant_mode = "train"
+  mode: str
+
+  def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
+    del mesh_axes
+    return functools.partial(JaxBlockScaledDotGeneral, mode=self.mode)
+
+  def einsum(self, *args, **kwargs):
+    del args, kwargs
+    raise ValueError(f"quantization='jax_{self.mode}' does not support einsum.")
+
+
 def infer_scale_granularity(scale_shape: Tuple[int, ...]) -> Tuple[str, int | None]:
   """Infers quantization scheme ('per_tensor', 'per_channel', or 'block_wise') from scale shape."""
   nontrivial_axes = [i for i, s in enumerate(scale_shape) if s > 1]
@@ -899,6 +938,8 @@ def _get_quant_config(config):
   """Set quantization params based on user configuration."""
   if not config.quantization or config.quantization == "":
     return None
+  if config.quantization in ("jax_mxfp8", "jax_mxfp4"):
+    return config.quantization
   if config.quantization == "int8":
     return _get_int8_quant_config(config)
   if config.quantization == "intmp":
@@ -959,7 +1000,9 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"):
     return None
   quant_cfg = _get_quant_config(config)
   if quant_cfg:
-    if quant_cfg == "fp8":
+    if quant_cfg in ("jax_mxfp8", "jax_mxfp4"):
+      return JaxBlockScaledQuantization(mode=quant_cfg.removeprefix("jax_"))
+    elif quant_cfg == "fp8":
       return Fp8Quantization()
     elif quant_cfg == "nanoo_fp8":
       return NANOOFp8Quantization()

tests/unit/jax_mx_quantization_test.py

@@ -0,0 +1,130 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the native JAX MX recipes, quantization=jax_mxfp8 and jax_mxfp4."""
+
+import sys
+import unittest
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.configs import pyconfig
+from maxtext.layers import linears
+from maxtext.layers import quantizations
+from tests.utils.test_helpers import get_test_config_path
+
+
+def _supports(mode: str) -> bool:
+  try:
+    jax.nn.get_scaled_dot_general_config(mode)
+  except ValueError:
+    return False
+  return True
+
+
+def _modes() -> list[str]:
+  return [mode for mode in ("mxfp8", "mxfp4") if _supports(mode)]
+
+
+def _config(**kwargs):
+  """A tiny dense model whose contracting extents are all multiples of 32, with jax_mxfp8 by default."""
+  init_kwargs = {
+      "run_name": "test",
+      "dataset_type": "synthetic",
+      "enable_checkpointing": False,
+      "enable_goodput_recording": False,
+      "steps": 1,
+      "per_device_batch_size": 1,
+      "use_qwix_quantization": False,
+      "skip_jax_distributed_system": True,
+      "base_emb_dim": 64,
+      "base_num_query_heads": 2,
+      "base_num_kv_heads": 2,
+      "head_dim": 32,
+      "base_mlp_dim": 128,
+      "base_num_decoder_layers": 1,
+      "max_target_length": 32,
+      "quantization": "jax_mxfp8",
+  } | kwargs
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs)
+
+
+class JaxMxQuantizationTest(unittest.TestCase):
+  """The quantization object and the layer it builds."""
+
+  def test_configure_quantization(self):
+    for mode in _modes():
+      with self.subTest(mode=mode):
+        quant = quantizations.configure_quantization(_config(quantization=f"jax_{mode}"))
+        self.assertIsInstance(quant, quantizations.JaxBlockScaledQuantization)
+        self.assertEqual(quant.mode, mode)
+
+  def test_layer_matches_scaled_dot_general(self):
+    x = jax.random.normal(jax.random.key(1), (2, 32, 64), jnp.float32).astype(jnp.bfloat16)
+    for mode in _modes():
+      with self.subTest(mode=mode):
+        self._check_layer_matches_scaled_dot_general(mode, x)
+
+  def _check_layer_matches_scaled_dot_general(self, mode, x):
+    """DenseGeneral with the MX backend computes exactly what jax.nn.scaled_dot_general computes."""
+    layer = linears.DenseGeneral(
+        in_features_shape=64,
+        out_features_shape=128,
+        weight_dtype=jnp.float32,
+        dtype=jnp.bfloat16,
+        quant=quantizations.JaxBlockScaledQuantization(mode=mode),
+        rngs=nnx.Rngs(0),
+    )
+    kernel = layer.kernel[...].astype(jnp.bfloat16)
+    config = jax.nn.get_scaled_dot_general_config(mode)
+
+    def reference(lhs):
+      return jax.nn.scaled_dot_general(
+          lhs, kernel, (((2,), (0,)), ((), ())), preferred_element_type=jnp.bfloat16, configs=[config, config, config]
+      )
+
+    def loss(fn, lhs):
+      return fn(lhs).astype(jnp.float32).sum()
+
+    y = layer(x)
+    self.assertEqual(y.dtype, jnp.bfloat16)
+    np.testing.assert_array_equal(np.asarray(y, np.float32), np.asarray(reference(x), np.float32))
+    dx = jax.grad(loss, argnums=1)(layer, x)
+    dx_ref = jax.grad(loss, argnums=1)(reference, x)
+    np.testing.assert_array_equal(np.asarray(dx, np.float32), np.asarray(dx_ref, np.float32))
+
+
+class JaxMxConfigValidationTest(unittest.TestCase):
+  """Settings that would silently leave GEMMs unquantized or fail obscurely in JAX are rejected at config time."""
+
+  def assert_rejected(self, regex, **kwargs):
+    # pydantic's ValidationError, which pyconfig raises, is a ValueError.
+    with self.assertRaisesRegex(ValueError, regex):
+      _config(**kwargs)
+
+  def test_qwix_flag(self):
+    self.assert_rejected("does not support use_qwix_quantization=True", use_qwix_quantization=True)
+
+  def test_explicit_sharding(self):
+    self.assert_rejected("does not support shard_mode='explicit'", shard_mode="explicit")
+
+  def test_moe(self):
+    self.assert_rejected("does not support MoE models", num_experts=8, num_experts_per_tok=2, base_moe_mlp_dim=128)
+
+
+if __name__ == "__main__":
+  unittest.main()