diff --git a/docs/guides/optimization/benchmark_and_performance.md b/docs/guides/optimization/benchmark_and_performance.md
index 7f50614672..cd72f24a54 100644
--- a/docs/guides/optimization/benchmark_and_performance.md
+++ b/docs/guides/optimization/benchmark_and_performance.md
@@ -191,7 +191,7 @@ Different quantization recipes are available, including` "int8", "fp8", "fp8_ful
 
 For v6e and earlier generation TPUs, use the "int8" recipe. For v7x and later generation TPUs, use "fp8_full". GPUs should use “fp8_gpu” for NVIDIA and "nanoo_fp8" for AMD.
 
-See [](quantization-doc).
+See [Quantization guide](../../reference/core_concepts/quantization.md).
 
 ### Choose sharding strategy
 
diff --git a/docs/reference/core_concepts/moe_configuration.md b/docs/reference/core_concepts/moe_configuration.md
index d83000fd93..171fd1680d 100644
--- a/docs/reference/core_concepts/moe_configuration.md
+++ b/docs/reference/core_concepts/moe_configuration.md
@@ -83,6 +83,8 @@ Dropping:
 
 `use_tokamax_gmm`: If enabled, use Tokamax library's Ragged Dot for matmul. Recommended for dropless configurations.
 
+`use_gmm_v2`: If enabled, use the Tokamax GMM v2 kernel for grouped matrix multiplication. Requires `use_tokamax_gmm` to be True.
+
 `megablox`: If enabled, use Megablox for sparse matrix operations. Effective only when `use_tokamax_gmm` is False.
 
 `capacity_factor`: A scalar multiplier for expert capacity. Effective only when `sparse_matmul` is False.
diff --git a/docs/reference/models/supported_models_and_architectures.md b/docs/reference/models/supported_models_and_architectures.md
index 4ff8927db5..93cf729198 100644
--- a/docs/reference/models/supported_models_and_architectures.md
+++ b/docs/reference/models/supported_models_and_architectures.md
@@ -12,7 +12,7 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX.
 
 - **Supported Precisions**: FP32, BF16, INT8, and FP8.
 - **Ahead-of-Time Compilation (AOT)**: For faster model development/prototyping and earlier OOM detection.
-- **Quantization**: Via **Qwix** (recommended) and AQT (deprecated). See the [Quantization guide](quantization-doc).
+- **Quantization**: Via **Qwix** (recommended) and AQT (deprecated). See the [Quantization guide](../core_concepts/quantization.md).
 - **Diagnostics**: Simple logging via `max_logging`, profiling in **XProf**, and visualization in **TensorBoard**.
 - **Multi-Token Prediction (MTP)**: Enables token efficient training with multi-token prediction.
 - **Elastic Training**: Fault-tolerant and dynamic scale-up/scale-down on Cloud TPUs with Pathways.
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index d0d3bd7028..d13e1e343c 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -700,6 +700,9 @@ def _should_save_checkpoint_at_step(checkpoint_manager, step, config, force):
   """Returns whether MaxText should build and dispatch checkpoint args."""
   if force:
     return True
+  if step == 0 and not config.save_checkpoint_on_start:
+    # if step = 0, `step % config.checkpoint_period == 0` is always true, force skip
+    return False
   if config.enable_continuous_checkpointing:
     base_checkpoint_due = bool(checkpoint_manager.should_save(step))
   else:
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 0f2960effb..64d23b187f 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -51,6 +51,7 @@ load_full_state_path: ""
 # async_checkpointing is true, else a synchronous one is used. If you have
 # problems with the checkpointer we recommend trying the synchronous one.
 enable_checkpointing: true
+save_checkpoint_on_start: true
 save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
@@ -141,7 +142,7 @@ kv_quant_dtype: "int8"
 checkpoint_is_quantized: false # set to true if reading from a saved aqt quantized checkpoint
 # saves params quantized on fly at following path
 save_quantized_params_path: ""
-#used to configure the mode in which model is called
+# used to configure the mode in which model is called
 # when left as is, corresponds to training
 # accepted values are "inference"
 model_call_mode: ""
@@ -267,6 +268,11 @@ wo_tile_drhs_mlp_dim: 1024
 
 merge_gating_gmm: false
 
+# Use tokamax library for gmm kernel implementation
+use_tokamax_gmm: false
+# Whether to use Tokamax GMM v2 for MoE kernel. Requires use_tokamax_gmm=true.
+use_gmm_v2: false
+
 norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights.
 
 # when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
@@ -1267,11 +1273,6 @@ use_qk_norm_in_gdn: true
 # The ratio of dimension to apply ROPE on
 partial_rotary_factor: 1.0
 
-# Use tokamax library for gmm kernel implementation
-use_tokamax_gmm: false
-# Whether to use GMM v2 for MoE.
-# Requires use_tokamax_gmm: true. Currently incompatible with quantization.
-use_gmm_v2: false
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.
 use_jax_splash: false
diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py
index f398244ca8..f784f27a64 100644
--- a/src/maxtext/configs/pyconfig.py
+++ b/src/maxtext/configs/pyconfig.py
@@ -38,7 +38,6 @@
 from maxtext.utils import max_logging
 
 logger = logging.getLogger(__name__)
-logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))
 
 _BASE_CONFIG_ATTR = "base_config"
 _MAX_PREFIX = "M_"
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index c4301fa93a..1d3661db73 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -358,6 +358,7 @@ class Checkpointing(BaseModel):
   source_checkpoint_layout: Literal["orbax", "safetensors", "safetensors_dynamic"] = Field(
       "orbax", description="The layout of the source checkpoint to load."
   )
+  save_checkpoint_on_start: bool = Field(True, description="If True, saves an initial checkpoint upon training start.")
   save_checkpoint_on_completion: bool = Field(
       True, description="If True, saves a final checkpoint upon training completion."
   )
@@ -450,7 +451,7 @@ class Quantization(BaseModel):
   use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.")
   use_manual_quantization: bool = Field(
       False,
-      description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.",
+      description="Whether to use manual quantization for batch split. Only used if `use_batch_split_schedule=True`.",
   )
   weight_quantization_calibration_method: str = Field(
       "absmax",
@@ -626,20 +627,6 @@ class Attention(BaseModel):
   use_post_attn_norm: bool = Field(False, description="Apply LayerNorm after the attention block.")
   use_post_ffw_norm: bool = Field(False, description="Apply LayerNorm after the feed-forward block.")
   use_ragged_attention: bool = Field(False, description="Whether to use ragged attention kernels.")
-  use_tokamax_gmm: bool = Field(
-      False,
-      description="Whether to use the Tokamax library for GMM kernel implementation.",
-  )
-  use_gmm_v2: bool = Field(
-      False,
-      description=(
-          "Whether to use GMM v2 (with bf16 activations and weights) for MoE."
-          " Requires use_tokamax_gmm: true. Currently incompatible with quantization."
-      ),
-  )
-  num_moe_emb_chunks: int = Field(
-      0, description="Number of chunks for overlapping token all-gather and GMM computation along embedding dimension."
-  )
   ragged_block_size: int = Field(256, description="Block size for ragged attention.")
   enable_padding_causal_mask: bool = Field(True, description="Temporary flag for Transformer Engine padding.")
   use_tokamax_splash: bool = Field(False, description="Whether to use tokamax splash attention.")
@@ -991,6 +978,21 @@ class MoEKernels(BaseModel):
 
   merge_gating_gmm: bool = Field(False, description="whether to merge the two gating gmm kernels into one.")
 
+  num_moe_emb_chunks: int = Field(
+      0, description="Number of chunks for overlapping token all-gather and GMM computation along embedding dimension."
+  )
+
+  # tokamax gmm
+  use_tokamax_gmm: bool = Field(
+      False,
+      description="Whether to use the Tokamax library for GMM kernel implementation.",
+  )
+
+  use_gmm_v2: bool = Field(
+      False,
+      description="Whether to use Tokamax GMM v2 for MoE kernel.",
+  )
+
 
 class DeepSeekMoE(BaseModel):
   """Configuration specific to DeepSeek-style MoE layers."""
@@ -3664,10 +3666,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
     if self.share_kv_projections and self.attention_type == "mla":
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
 
-    if self.use_gmm_v2 and (self.quantization or self.use_qwix_quantization):
-      raise ValueError("Quantization with GMM v2 is not supported yet.")
-    if self.use_gmm_v2 and not self.use_tokamax_gmm:
-      raise ValueError("GMM v2 requires `use_tokamax_gmm=true`.")
+    if self.use_manual_quantization and not self.use_batch_split_schedule:
+      raise ValueError("manual quantization is only used when `use_batch_split_schedule=True`.")
+
+    # Validation for GMM v2
+    if self.use_gmm_v2:
+      if not self.use_tokamax_gmm:
+        raise ValueError("GMM v2 requires `use_tokamax_gmm=True`.")
+      if self.use_batch_split_schedule:
+        raise ValueError("GMM v2 is not supported with a batch split schedule.")
 
     for val in self.compress_ratios:
       if val != 0 and val < 4:
diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py
index dab4eef70a..e4f147a274 100644
--- a/src/maxtext/kernels/megablox/ops.py
+++ b/src/maxtext/kernels/megablox/ops.py
@@ -86,10 +86,12 @@ def gmm(
     interpret = False
   quantization_rule = None
   if use_qwix_quantization:
-    # get_current_rule has to be called outside of the _gmm_fwd function.
+    # 1. for non-batchsplit, retrieve rule ("gmm") via qwix interception
+    #   get_current_rule has to be called outside of the _gmm_fwd function.
+    # 2. for batchsplit, explicitly pass the rule
     quantization_rule = qwix_rule if qwix_rule else qpl.get_current_rule("gmm")
-    if quantization_rule and not isinstance(quantization_rule, qwix.QtRule):
-      raise ValueError("Expect a QtRule for quantized training.")
+    if not quantization_rule or not isinstance(quantization_rule, qwix.QtRule):
+      raise ValueError("Expect a QtRule for quantized training. " f"But get quantization_rule={quantization_rule}")
   else:
     # Handcraft a rule that matches the AQT's behavior.
     if lhs_quantize_dtype or rhs_quantize_dtype:
@@ -101,7 +103,10 @@ 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))
+  gmm_fwd_bwd = jax.custom_vjp(
+      gmm_fwd_bwd,
+      nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15),
+  )
   gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype))
   return gmm_fwd_bwd(
       lhs,
@@ -124,6 +129,11 @@ def gmm(
   )
 
 
+# ==============================================================================
+# Forward: FWD GMM
+# ==============================================================================
+
+
 def _gmm_fwd(
     lhs: jnp.ndarray,
     rhs: jnp.ndarray,
@@ -159,96 +169,228 @@ def _gmm_fwd(
         jnp.ndarray | qpl.QArray,
         jnp.ndarray,
         jnp.ndarray | None,
+        jnp.ndarray | None,
     ],
 ]:
-  """Forward function for GMM VJP."""
+  """Forward function for GMM VJP.
+
+  - lhs: [m, k]
+  - rhs: [g, k, n] if transpose_rhs=False. [g, n, k] if transpose_rhs=True
+  """
+
+  # Quantize activation and weight
   if quantization_rule:
-    if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray):
-      lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
-          lhs,
-          quantization_rule.act_qtype,
-          channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
-          calibration_method=quantization_rule.act_calibration_method,  # pyrefly: ignore[bad-argument-type]
-      )
-    if quantization_rule.weight_qtype and not isinstance(rhs, qpl.QArray):
-      if not use_manual_quantization:
-        rhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
-            rhs,
-            quantization_rule.weight_qtype,
-            # If only considering the fwd pass, we could also enable channelwise
-            # axes for the group axis, i.e., [0, 1 or 2]. However, this makes the
-            # bwd pass unable to reuse the scale easily.
-            channelwise_axes=([] if quantization_rule.disable_channelwise_axes else ([1] if transpose_rhs else [2])),
-            calibration_method=quantization_rule.weight_calibration_method,
-        )
-      else:
-        rhs = quantizations.manual_quantize(  # pyrefly: ignore[bad-assignment]
-            rhs,
-            quantization_rule.weight_qtype,
-            calibration_method=quantization_rule.weight_calibration_method,
-        )
-  # QAG is only supported for following conditions
-  if use_tokamax_backend:
-    if quantization_rule and quantization_rule.bwd_qtype:
-      if quantization_rule.weight_calibration_method.startswith("fixed") and isinstance(rhs, qpl.QArray):
-        if weight_gather_axes:
-          for axis_name, axis_idx in weight_gather_axes:
-            rhs_qvalue = jax.lax.all_gather(rhs.qvalue, axis_name, axis=axis_idx, tiled=True)
-            rhs = dataclasses.replace(rhs, qvalue=rhs_qvalue)
-    # Handle transpose_rhs manually as ragged_dot assumes (G, K, N)
-    if transpose_rhs:
-      rhs = rhs.swapaxes(1, 2)
-
-  if use_gmm_v2:
-    out = gmm_v2.gmm_v2(
-        lhs=lhs,
-        rhs=rhs,
-        group_sizes=group_sizes,
-        tile_info=gmm_v2.TileSizes(
-            tile_m=tiling[0],
-            tile_k=tiling[1],
-            tile_n=tiling[2],
-        ),
-        preferred_element_type=preferred_element_type,
-        partial_sum=partial_sum,
-        group_offset=group_offset,
+    lhs, rhs = _fwd_quantize_activation_and_weight(
+        lhs, rhs, quantization_rule, use_gmm_v2, use_manual_quantization, transpose_rhs
     )
-  elif use_tokamax_backend:
-    # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
-    out_kwargs = {}
-    if use_manual_quantization:
-      # used in batchsplit
-      out_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(varying=frozenset(["data", "fsdp", "expert"]))
-
-    out = tokamax.ragged_dot(
-        lhs=lhs,
-        rhs=rhs,
-        group_sizes=group_sizes,
-        precision=jax.lax.Precision.DEFAULT,
-        preferred_element_type=preferred_element_type,
-        # `group_offset` is not yet supported
-        group_offset=None,
-        implementation="mosaic",
-        **out_kwargs,
+
+  # Quantization All-Gather (QAG) for weight: only supported for following conditions
+  if (
+      use_tokamax_backend
+      and quantization_rule
+      and quantization_rule.bwd_qtype
+      and quantization_rule.weight_calibration_method.startswith("fixed")
+      and isinstance(rhs, qpl.QArray)
+      and weight_gather_axes
+  ):
+    rhs = _fwd_gather_weight(rhs, weight_gather_axes)
+
+  # 3. Backend Execution Routing
+  if use_tokamax_backend and not use_gmm_v2:
+    out = _fwd_run_tokamax_v1(lhs, rhs, group_sizes, preferred_element_type, transpose_rhs, use_manual_quantization)
+  elif use_tokamax_backend and use_gmm_v2:
+    out = _fwd_run_tokamax_v2(
+        lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs
     )
   else:
-    out = backend.gmm(
+    out = _fwd_run_megablox(
         lhs,
         rhs,
         group_sizes,
         preferred_element_type,
-        tiling[:3],
+        tiling,
         group_offset,
         existing_out,
-        transpose_rhs=transpose_rhs,
-        interpret=interpret,
+        transpose_rhs,
+        interpret,
+        lhs_vma_axes,
     )
-    for axis in lhs_vma_axes:
-      out = jax.lax.pcast(out, axis_name=axis, to="varying")
 
   return out, (lhs, rhs, group_sizes, group_offset, partial_sum)
 
 
+def _fwd_quantize_activation_and_weight(
+    lhs: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray | qpl.QArray,
+    quantization_rule: qwix.QtRule,
+    use_gmm_v2: bool,
+    use_manual_quantization: bool,
+    transpose_rhs: bool,
+) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]:
+  """Handles act and weight quantization for GMM forward inputs."""
+  if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not use_gmm_v2:
+    lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
+        lhs,
+        quantization_rule.act_qtype,
+        channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
+        calibration_method=quantization_rule.act_calibration_method,
+    )
+
+  if quantization_rule.weight_qtype and not isinstance(rhs, qpl.QArray):
+    if not use_manual_quantization:
+      rhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
+          rhs,
+          quantization_rule.weight_qtype,
+          # If only considering the fwd pass, we could also enable channelwise
+          # axes for the group axis, i.e., [0, 1 or 2]. However, this makes the
+          # bwd pass unable to reuse the scale easily.
+          channelwise_axes=([] if quantization_rule.disable_channelwise_axes else ([1] if transpose_rhs else [2])),
+          calibration_method=quantization_rule.weight_calibration_method,
+      )
+    else:
+      rhs = quantizations.manual_quantize(  # pyrefly: ignore[bad-assignment]
+          rhs,
+          quantization_rule.weight_qtype,
+          calibration_method=quantization_rule.weight_calibration_method,
+      )
+  return lhs, rhs
+
+
+def _fwd_gather_weight(rhs: qpl.QArray, weight_gather_axes: List[Tuple[str, int]]) -> qpl.QArray:
+  """Applies QAG (Quantization All-Gather) to RHS weights during forward pass."""
+  for axis_name, axis_idx in weight_gather_axes:
+    rhs_qvalue = jax.lax.all_gather(rhs.qvalue, axis_name, axis=axis_idx, tiled=True)
+    # replace the qvalue with the gathered qvalue in the QArray
+    rhs = dataclasses.replace(rhs, qvalue=rhs_qvalue)
+  return rhs
+
+
+def _fwd_run_tokamax_v1(
+    lhs: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray | qpl.QArray,
+    group_sizes: jnp.ndarray,
+    preferred_element_type: jnp.dtype,
+    transpose_rhs: bool,
+    use_manual_quantization: bool,
+) -> jnp.ndarray:
+  """Executes the standard Tokamax GMM V1 for forward pass."""
+  # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
+  out_kwargs = {}
+  if use_manual_quantization:
+    # used in batchsplit
+    out_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(varying=frozenset(["data", "fsdp", "expert"]))
+
+  if transpose_rhs:
+    rhs = rhs.swapaxes(1, 2)
+
+  return tokamax.ragged_dot(
+      lhs=lhs,
+      rhs=rhs,
+      group_sizes=group_sizes,
+      precision=jax.lax.Precision.DEFAULT,
+      preferred_element_type=preferred_element_type,
+      # `group_offset` is not yet supported
+      group_offset=None,
+      implementation="mosaic",
+      **out_kwargs,
+  )
+
+
+def _fwd_prepare_rhs_scale(rhs: qpl.QArray, transpose_rhs: bool = False) -> jnp.ndarray:
+  """Formats and broadcasts rhs scale for the V2 GMM forward kernel."""
+  # Target shape: (size_group, num_quant_blocks, 1, size_n)
+  if transpose_rhs:
+    G, N, _ = rhs.qvalue.shape
+    scale = rhs.scale
+    if scale.ndim == 3:
+      scale = scale.swapaxes(1, 2)
+  else:
+    G, _, N = rhs.qvalue.shape
+    scale = rhs.scale
+
+  if scale.ndim == 2:  # Per-Channel quantization
+    rhs_scale = jnp.expand_dims(scale, axis=(1, 2))
+  elif scale.ndim == 3:  # Block-wise quantization
+    rhs_scale = jnp.expand_dims(scale, axis=2)
+  else:  # Per-tensor quantization, (1, 1, 1, 1)
+    rhs_scale = scale
+
+  num_quant_blocks = rhs_scale.shape[1] if rhs_scale.ndim > 1 else 1
+  return jnp.broadcast_to(rhs_scale, (G, num_quant_blocks, 1, N))
+
+
+def _fwd_run_tokamax_v2(
+    lhs: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray | qpl.QArray,
+    group_sizes: jnp.ndarray,
+    preferred_element_type: jnp.dtype,
+    tiling: tuple,
+    group_offset: jnp.ndarray | None,
+    partial_sum: jnp.ndarray | None,
+    transpose_rhs: bool,
+) -> jnp.ndarray:
+  """Executes the Tokamax GMM V2 backend for forward pass OUT = LHS @ RHS."""
+  # if transpose_rhs=False, rhs is [g, k, n], remain unchanged
+  # if transpose_rhs=True, rhs [g, n, k], explicit transpose to [g, k, n]
+  rhs_operand = rhs if not transpose_rhs else rhs.swapaxes(1, 2)
+  rhs_scale = None
+
+  if isinstance(rhs, qpl.QArray):
+    rhs_operand = rhs_operand.qvalue
+    rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs)
+
+  custom_fwd_tiling = gmm_v2.TileSizes(
+      tile_m=tiling[0],
+      tile_k=tiling[1],
+      tile_n=tiling[2],
+  )
+
+  return gmm_v2.gmm_v2(
+      lhs=lhs,
+      rhs=rhs_operand,
+      group_sizes=group_sizes,
+      rhs_scale=rhs_scale,
+      tile_info=custom_fwd_tiling,
+      preferred_element_type=preferred_element_type,
+      partial_sum=partial_sum,
+      group_offset=group_offset,
+  )
+
+
+def _fwd_run_megablox(
+    lhs: jnp.ndarray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    preferred_element_type: jnp.dtype,
+    tiling: tuple,
+    group_offset: jnp.ndarray | None,
+    existing_out: jnp.ndarray | None,
+    transpose_rhs: bool,
+    interpret: bool,
+    lhs_vma_axes: tuple,
+) -> jnp.ndarray:
+  """Executes the Megablox backend fallback for forward pass."""
+  out = backend.gmm(
+      lhs,
+      rhs,
+      group_sizes,
+      preferred_element_type,
+      tiling[:3],
+      group_offset,
+      existing_out,
+      transpose_rhs=transpose_rhs,
+      interpret=interpret,
+  )
+  for axis in lhs_vma_axes:
+    out = jax.lax.pcast(out, axis_name=axis, to="varying")
+  return out
+
+
+# ==============================================================================
+# Backward: BWD DLHS GMM + BWD DRHS TGMM
+# ==============================================================================
+
+
 def _gmm_bwd(
     lhs_dtype: jax.typing.DTypeLike,
     rhs_dtype: jax.typing.DTypeLike,
@@ -271,7 +413,7 @@ def _gmm_bwd(
         jnp.ndarray | None,
     ],
     grad: jnp.ndarray,
-) -> tuple[jnp.ndarray, jnp.ndarray, None, None, jnp.ndarray, jnp.ndarray | None]:
+) -> tuple[jnp.ndarray, jnp.ndarray, None, None, jnp.ndarray | None, jnp.ndarray | None]:
   """Backward function for throughput GMM VJP."""
   del preferred_element_type
   lhs, rhs, group_sizes, group_offset, partial_sum_fwd = residual
@@ -286,20 +428,120 @@ def _gmm_bwd(
   #  - dlhs_dout: the incoming gradient used to calculate dlhs.
   #  - drhs_dout: the incoming gradient used to calculate drhs.
 
+  # 1. Scale Application & QArray Unwrapping
+  dlhs_dout, drhs_dout, lhs, rhs = _bwd_prepare_inputs(
+      grad, lhs, rhs, group_sizes, use_gmm_v2, transpose_rhs, quantization_rule
+  )
+
+  # 2. Backward Pass Quantization
+  if quantization_rule:
+    dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule)
+
+  # 3. DLHS Gradient Execution
+  dlhs = _compute_dlhs(
+      dlhs_dout,
+      rhs,
+      group_sizes,
+      group_offset,
+      lhs_dtype,
+      tiling,
+      transpose_rhs,
+      use_tokamax_backend,
+      use_gmm_v2,
+      use_manual_quantization,
+      interpret,
+      lhs_vma_axes,
+  )
+
+  # 4. DRHS Gradient Execution
+  drhs = _compute_drhs(
+      drhs_dout,
+      lhs,
+      group_sizes,
+      group_offset,
+      num_actual_groups,
+      rhs_dtype,
+      tiling,
+      use_tokamax_backend,
+      use_gmm_v2,
+      use_manual_quantization,
+      weight_gather_axes,
+      interpret,
+      rhs_vma_axes,
+      quantization_rule,
+  )
+
+  # 5. Output Formatting
+  # NOTE: If the rhs transposition is fused into the forward pass we need to
+  # return the transpose of the rhs gradient that we calculated above.
+  #
+  # TODO(tgale, enriqueps, apaske): Fuse this transposition into the tgmm.
+  drhs = drhs.swapaxes(1, 2) if transpose_rhs else drhs
+  dpartial_sum = grad if partial_sum_fwd is not None else None
+  d_existing_out = None if use_tokamax_backend else grad
+
+  return dlhs, drhs, None, None, d_existing_out, dpartial_sum
+
+
+def _bwd_prepare_inputs(
+    grad: jnp.ndarray,
+    lhs: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray | qpl.QArray,
+    group_sizes: jnp.ndarray,
+    use_gmm_v2: bool,
+    transpose_rhs: bool,
+    quantization_rule: qwix.QtRule | None,
+) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, jnp.ndarray, jnp.ndarray]:
+  """Prepares backward operands."""
+
   # dlhs_dout and drhs_dout can be different when quantization is enabled.
   dlhs_dout = grad
   drhs_dout = grad
-  if isinstance(rhs, qpl.QArray):  # qvalue: [g, k, n] scale: [1, 1, n]
-    # Apply rhs.scale to dlhs_dout to avoid dequantizing or requantizing rhs.
-    # We cannot apply the scale to dlhs because axis n will disappear there.
-    dlhs_dout *= rhs.scale.astype(grad.dtype).reshape(1, -1)  # [1, n]
-    rhs = rhs.qvalue
-  if isinstance(lhs, qpl.QArray):  # qvalue: [m, k] scale: [m, 1]
-    # Apply lhs.scale to drhs_dout, as axis m will disappear in drhs.
+
+  # Apply rhs.scale to dlhs_dout, dlhs_dout[m, n] @ rhs_tranpose[g, n, k] = dlhs[m, k]
+  # Assume channelwise scale on rhs n.
+  # Apply rhs.scale to dlhs_dout to avoid dequantizing or requantizing rhs.
+  # We cannot apply the scale to dlhs because axis n will disappear there.
+  if isinstance(rhs, qpl.QArray):
+    # rhs - qvalue: [g, k, n] scale: [1, 1, n], assume transpose_rhs=False
+    if not use_gmm_v2:
+      dlhs_dout *= rhs.scale.astype(grad.dtype).reshape(1, -1)
+      rhs = rhs.qvalue
+    else:
+      # NOTE: rhs.scale is for the contracting dimension (N) in DLHS, but gmm_v2
+      # only supports scaling the output dimension. Thus, we must scale dlhs_dout
+      # beforehand. If the kernel supports transposing RHS internally, we can fuse
+      # this scale inside the kernel.
+      dlhs_dout = _dlhs_scale_grad_by_rhs_scale(dlhs_dout, rhs, group_sizes, transpose_rhs)
+      rhs = rhs.qvalue
+
+  # 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):
+    lhs = qpl.quantize(  # pyrefly: ignore[bad-assignment]
+        lhs,
+        quantization_rule.act_qtype,
+        channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
+        calibration_method=quantization_rule.act_calibration_method,
+    )
+
+  # Assume channelwise scale on lhs m, lhs_transpose[k, m] @ drhs_out[m, n] = drhs[g, k, n]
+  # Apply lhs.scale to drhs_dout, as axis m will disappear in drhs.
+  if isinstance(lhs, qpl.QArray):
+    # lhs - qvalue: [m, k] scale: [m, 1]
     drhs_dout *= lhs.scale.astype(grad.dtype)
     lhs = lhs.qvalue
-  if quantization_rule and quantization_rule.bwd_qtype:
-    # Enable backward pass quantization
+
+  return dlhs_dout, drhs_dout, lhs, rhs
+
+
+def _bwd_quantize_gradient(
+    dlhs_dout: jnp.ndarray | qpl.QArray,
+    drhs_dout: jnp.ndarray | qpl.QArray,
+    quantization_rule: qwix.QtRule,
+) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]:
+  """Applies backward quantization to incoming gradients."""
+  if quantization_rule.bwd_qtype:
     dlhs_dout = qpl.quantize(
         dlhs_dout,
         quantization_rule.bwd_qtype,
@@ -312,113 +554,301 @@ def _gmm_bwd(
         channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [1],
         calibration_method=quantization_rule.bwd_calibration_method,
     )
-  if use_tokamax_backend or use_gmm_v2:
-    # Handle transpose_rhs manually
-    dlhs_rhs = rhs
-    if transpose_rhs:
-      dlhs_rhs = dlhs_rhs.swapaxes(1, 2)
-
-    # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
-    dlhs_kwargs = {}
-    drhs_kwargs = {}
-    if use_manual_quantization:
-      # used in batchsplit
-      dlhs_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(varying=frozenset(["data", "fsdp", "expert"]))
-      drhs_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(
-          varying=frozenset(["expert"]), unreduced=frozenset(["data", "fsdp"])
-      )
+  return dlhs_dout, drhs_dout
 
-    if use_gmm_v2:
-      dlhs = gmm_v2.gmm_v2(
-          lhs=dlhs_dout,
-          rhs=dlhs_rhs.swapaxes(1, 2),  # requires rhs to be [g, n, k]
-          group_sizes=group_sizes,
-          tile_info=gmm_v2.TileSizes(
-              tile_m=tiling[3],
-              tile_k=tiling[4],
-              tile_n=tiling[5],
-          ),
-          preferred_element_type=lhs_dtype,
-          group_offset=group_offset,
-      )
 
-      # tgmm_v2_op requires lhs and rhs to have the same dtype.
-      if lhs.dtype != drhs_dout.dtype:
-        drhs_dout = drhs_dout.astype(lhs.dtype)
-
-      drhs = tgmm_v2.tgmm_v2(
-          lhs=lhs,
-          rhs=drhs_dout,
-          group_sizes=group_sizes,
-          num_actual_groups=num_actual_groups,
-          precision=jax.lax.Precision.DEFAULT,
-          preferred_element_type=rhs_dtype,
-          group_offset=group_offset,
-          tile_info=gmm_v2.TileSizes(
-              tile_m=tiling[6],
-              tile_k=tiling[7],
-              tile_n=tiling[8],
-          ),
-      )
-    else:
-      dlhs = tokamax.ragged_dot_general(
-          lhs=dlhs_dout,
-          rhs=dlhs_rhs,
-          group_sizes=group_sizes,
-          ragged_dot_dimension_numbers=DLHS_RAGGED_DOT_DIM_NUMS,
-          precision=jax.lax.Precision.DEFAULT,
-          preferred_element_type=lhs_dtype,
-          # `group_offset` is not yet supported
-          group_offset=None,
-          implementation="mosaic",
-          **dlhs_kwargs,
-      )
+# ==============================================================================
+# BWD DLHS GMM
+# ==============================================================================
 
-      drhs = tokamax.ragged_dot_general(
-          lhs=lhs,
-          rhs=drhs_dout,
-          group_sizes=group_sizes,
-          ragged_dot_dimension_numbers=DRHS_RAGGED_DOT_DIM_NUMS,
-          precision=jax.lax.Precision.DEFAULT,
-          preferred_element_type=rhs_dtype,
-          # `group_offset` is not yet supported
-          group_offset=None,
-          implementation="mosaic",
-          **drhs_kwargs,
-      )
-    if quantization_rule and quantization_rule.bwd_qtype and weight_gather_axes:
-      # Scatter back in reverse order of gather
-      for axis_name, axis_idx in reversed(weight_gather_axes):
-        drhs = jax.lax.psum_scatter(drhs, axis_name, scatter_dimension=axis_idx, tiled=True)
-  else:
-    dlhs = backend.gmm(
+
+def _compute_dlhs(
+    dlhs_dout: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    lhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+    transpose_rhs: bool,
+    use_tokamax_backend: bool,
+    use_gmm_v2: bool,
+    use_manual_quantization: bool,
+    interpret: bool,
+    lhs_vma_axes: tuple,
+) -> jnp.ndarray:
+  """Routes execution of DLHS based on backend choices."""
+  if use_tokamax_backend and not use_gmm_v2:
+    return _dlhs_run_tokamax_v1(
         dlhs_dout,
         rhs,
         group_sizes,
         lhs_dtype,
-        tiling[3:6],
-        group_offset,
-        transpose_rhs=not transpose_rhs,
-        interpret=interpret,
-        varying_axes=lhs_vma_axes,
+        transpose_rhs,
+        use_manual_quantization,
+    )
+  elif use_tokamax_backend and use_gmm_v2:
+    return _dlhs_run_tokamax_v2(dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, transpose_rhs)
+  else:
+    return _dlhs_run_megablox(
+        dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, transpose_rhs, interpret, lhs_vma_axes
     )
 
-    drhs = backend.tgmm(
-        lhs.swapaxes(0, 1),
-        drhs_dout,
+
+def _dlhs_run_tokamax_v1(
+    dlhs_dout: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    lhs_dtype: jax.typing.DTypeLike,
+    transpose_rhs: bool,
+    use_manual_quantization: bool,
+) -> jnp.ndarray:
+  """Executes DLHS using GMM 1"""
+  dlhs_kwargs = {}
+  if use_manual_quantization:
+    dlhs_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(varying=frozenset(["data", "fsdp", "expert"]))
+
+  dlhs_rhs = rhs.swapaxes(1, 2) if transpose_rhs else rhs
+  return tokamax.ragged_dot_general(
+      lhs=dlhs_dout,
+      rhs=dlhs_rhs,
+      group_sizes=group_sizes,
+      ragged_dot_dimension_numbers=DLHS_RAGGED_DOT_DIM_NUMS,
+      precision=jax.lax.Precision.DEFAULT,
+      preferred_element_type=lhs_dtype,
+      # `group_offset` is not yet supported
+      group_offset=None,
+      implementation="mosaic",
+      **dlhs_kwargs,
+  )
+
+
+def _dlhs_scale_grad_by_rhs_scale(
+    grad: jnp.ndarray,
+    rhs: qpl.QArray,
+    group_sizes: jnp.ndarray,
+    transpose_rhs: bool = False,
+) -> jnp.ndarray:
+  """Squeezes the rhs scale and multiplies it with the incoming gradient.
+
+  Scaling is applied before the V2 GMM DLHS kernel.
+  """
+  rhs_scale = rhs.scale
+
+  # 1. Squeeze the scale to 2D [g, n] based on transpose_rhs
+  if rhs_scale.ndim == 3:
+    squeeze_axis = 2 if transpose_rhs else 1
+    if rhs_scale.shape[squeeze_axis] == 1:
+      rhs_scale = rhs_scale.squeeze(axis=squeeze_axis)
+
+  # 2. Apply scale (handle shared vs per-expert scales)
+  if rhs_scale.shape[0] == 1:
+    return grad * rhs_scale.astype(grad.dtype)
+  else:
+    repeated_scale = jnp.repeat(
+        rhs_scale.astype(grad.dtype),
         group_sizes,
-        rhs_dtype,
-        tiling[-3:],
-        group_offset,
-        num_actual_groups,
-        interpret=interpret,
-        varying_axes=rhs_vma_axes,
+        axis=0,
+        total_repeat_length=grad.shape[0],
     )
+    return grad * repeated_scale
 
-  # NOTE: If the rhs transposition is fused into the forward pass we need to
-  # return the transpose of the rhs gradient that we calculated above.
-  #
-  # TODO(tgale, enriqueps, apaske): Fuse this transposition into the tgmm.
-  drhs = drhs.swapaxes(1, 2) if transpose_rhs else drhs
-  dpartial_sum = grad if partial_sum_fwd is not None else None
-  return dlhs, drhs, None, None, grad, dpartial_sum
+
+def _dlhs_run_tokamax_v2(
+    dlhs_dout: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    lhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+    transpose_rhs: bool,
+) -> jnp.ndarray:
+  """Executes Tokamax GMM V2 backend for DLHS = DLHS_dout @ RHS^T."""
+  # NOTE: We manually transpose RHS here because gmm_v2 lacks native transpose_rhs
+  # support. Fusing this transpose into the kernel would also allow us to fuse
+  # the rhs_scale application.
+  dlhs_rhs = rhs if transpose_rhs else rhs.swapaxes(1, 2)
+  dlhs_lhs = dlhs_dout.qvalue if isinstance(dlhs_dout, qpl.QArray) else dlhs_dout
+
+  custom_dlhs_tiling = gmm_v2.TileSizes(tile_m=tiling[3], tile_k=tiling[4], tile_n=tiling[5])
+
+  dlhs = gmm_v2.gmm_v2(
+      lhs=dlhs_lhs,
+      rhs=dlhs_rhs,
+      group_sizes=group_sizes,
+      # rhs scale is already applied to dlhs_lhs
+      rhs_scale=None,
+      tile_info=custom_dlhs_tiling,
+      preferred_element_type=lhs_dtype,
+      group_offset=group_offset,
+  )
+
+  if isinstance(dlhs_dout, qpl.QArray):
+    dlhs *= dlhs_dout.scale.astype(dlhs.dtype)
+
+  return dlhs
+
+
+def _dlhs_run_megablox(
+    dlhs_dout: jnp.ndarray | qpl.QArray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    lhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+    transpose_rhs: bool,
+    interpret: bool,
+    lhs_vma_axes: tuple,
+) -> jnp.ndarray:
+  """Executes Megablox fallback for DLHS."""
+  return backend.gmm(
+      dlhs_dout,
+      rhs,
+      group_sizes,
+      lhs_dtype,
+      tiling[3:6],
+      group_offset,
+      transpose_rhs=not transpose_rhs,
+      interpret=interpret,
+      varying_axes=lhs_vma_axes,
+  )
+
+
+# ==============================================================================
+# BWD DRHS TGMM
+# ==============================================================================
+
+
+def _compute_drhs(
+    drhs_dout: jnp.ndarray | qpl.QArray,
+    lhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    num_actual_groups: int,
+    rhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+    use_tokamax_backend: bool,
+    use_gmm_v2: bool,
+    use_manual_quantization: bool,
+    weight_gather_axes: List[Tuple[str, int]] | None,
+    interpret: bool,
+    rhs_vma_axes: tuple,
+    quantization_rule: qwix.QtRule | None,
+) -> 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)
+  else:
+    drhs = _drhs_run_megablox(
+        drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, interpret, rhs_vma_axes
+    )
+
+  if use_tokamax_backend and quantization_rule and quantization_rule.bwd_qtype and weight_gather_axes:
+    drhs = _drhs_scatter_weight(drhs, weight_gather_axes)
+
+  return drhs
+
+
+def _drhs_scatter_weight(drhs: jnp.ndarray, weight_gather_axes: List[Tuple[str, int]]) -> jnp.ndarray:
+  """Scatters the DRHS output back in the reverse order of the forward gather."""
+  for axis_name, axis_idx in reversed(weight_gather_axes):
+    drhs = jax.lax.psum_scatter(drhs, axis_name, scatter_dimension=axis_idx, tiled=True)
+  return drhs
+
+
+def _drhs_run_tokamax_v1(
+    drhs_dout: jnp.ndarray | qpl.QArray,
+    lhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    rhs_dtype: jax.typing.DTypeLike,
+    use_manual_quantization: bool,
+) -> jnp.ndarray:
+  """Executes standard Tokamax ragged_dot for DRHS."""
+  drhs_kwargs = {}
+  if use_manual_quantization:
+    drhs_kwargs["manual_axis_type"] = jax.sharding.ManualAxisType(
+        varying=frozenset(["expert"]), unreduced=frozenset(["data", "fsdp"])
+    )
+  return tokamax.ragged_dot_general(
+      lhs=lhs,
+      rhs=drhs_dout,
+      group_sizes=group_sizes,
+      ragged_dot_dimension_numbers=DRHS_RAGGED_DOT_DIM_NUMS,
+      precision=jax.lax.Precision.DEFAULT,
+      preferred_element_type=rhs_dtype,
+      # `group_offset` is not yet supported
+      group_offset=None,
+      implementation="mosaic",
+      **drhs_kwargs,
+  )
+
+
+def _drhs_prepare_bwd_scale(drhs_dout: qpl.QArray) -> jnp.ndarray:
+  """Formats and broadcasts drhs_dout scale to (1, 1, size_n) for V2 TGMM kernel."""
+  scale = drhs_dout.scale
+  size_n = drhs_dout.shape[1]
+  # per channel: (1, n) -> (1, 1, n)
+  # per tensor: (1, 1) -> (1, 1, 1)
+  rhs_scale = jnp.expand_dims(scale, axis=1)
+  # per-tensor quantization: broadcast (1, 1, 1) to (1, 1, size_n)
+  if rhs_scale.shape[2] == 1:
+    rhs_scale = jnp.broadcast_to(rhs_scale, (1, 1, size_n))
+  return rhs_scale
+
+
+def _drhs_run_tokamax_v2(
+    drhs_dout: jnp.ndarray | qpl.QArray,
+    lhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    num_actual_groups: int,
+    rhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+) -> 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
+  drhs_lhs = lhs
+
+  rhs_scale = None
+  if isinstance(drhs_dout, qpl.QArray):
+    rhs_scale = _drhs_prepare_bwd_scale(drhs_dout)
+
+  custom_drhs_tiling = gmm_v2.TileSizes(tile_m=tiling[6], tile_k=tiling[7], tile_n=tiling[8])
+
+  return tgmm_v2.tgmm_v2(
+      lhs=drhs_lhs,
+      rhs=drhs_rhs,
+      group_sizes=group_sizes,
+      num_actual_groups=num_actual_groups,
+      rhs_scale=rhs_scale,
+      precision=jax.lax.Precision.DEFAULT,
+      preferred_element_type=rhs_dtype,
+      group_offset=group_offset,
+      tile_info=custom_drhs_tiling,
+  )
+
+
+def _drhs_run_megablox(
+    drhs_dout: jnp.ndarray | qpl.QArray,
+    lhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    group_offset: jnp.ndarray | None,
+    num_actual_groups: int,
+    rhs_dtype: jax.typing.DTypeLike,
+    tiling: tuple,
+    interpret: bool,
+    rhs_vma_axes: tuple,
+) -> jnp.ndarray:
+  """Executes Megablox fallback for DRHS."""
+  return backend.tgmm(
+      lhs.swapaxes(0, 1),
+      drhs_dout,
+      group_sizes,
+      rhs_dtype,
+      tiling[-3:],
+      group_offset,
+      num_actual_groups,
+      interpret=interpret,
+      varying_axes=rhs_vma_axes,
+  )
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
index 7ba897025f..46fcfbeafb 100644
--- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
@@ -13,8 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # Forked from:
-# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/
-# tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py
+# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py
 """GMM kernel implemented using Pallas."""
 
 from abc import ABC, abstractmethod
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py
index 52d2a94a2b..cdef961def 100644
--- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py
@@ -13,8 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # Forked from:
-# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/
-# tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_tgmm_kernel.py
+# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_tgmm_kernel.py
 """TGMM kernel"""
 
 import dataclasses
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index d722270c4c..6c5d3ffcca 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -996,7 +996,7 @@ def __call__(
             # scan with initialized parameters.
             if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
               # old version of batch-split that fully uses qwix quantization.
-              if cfg.use_qwix_quantization and not cfg.use_manual_quantization:
+              if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization:
                 y = deepseek_batchsplit_fp8.scan_batch_split_layers(
                     y,
                     self.variables["params"]["moe_layers"],
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 915193eb0f..6ae9abf2cc 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -1391,6 +1391,7 @@ def sparse_matmul(
     def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, padding_amount):
       """Execute jax.lax.ragged_dot, with potential quantization"""
       m, k, n = inputs.shape[0], inputs.shape[1], kernel.shape[2]
+      # Clamps the tile size using the minimum
       tiling = (
           min(tiling[0], m),
           min(tiling[1], k),
@@ -1401,7 +1402,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments,
         if kernel.bias or kernel.sparsity_mask or len(kernel.scale) > 1:
           raise ValueError("Unsupported usecase for ragged_dot with quantized kernel.")
         rhs_inputs = kernel.qvalue
-      if self.config.use_qwix_quantization:
+      if self.config.quantization and self.config.use_qwix_quantization:
         # Use full contraction for QWIX quantization to allow quantization
         # fusion (max reduce over contracting dimension).
         tiling = (tiling[0], k, tiling[2])
@@ -1432,7 +1433,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments,
       return output
 
     def get_tokamax_group_sizes(group_sizes, inputs, _kernel):
-      if self.config.use_qwix_quantization:
+      if self.config.quantization and self.config.use_qwix_quantization:
         return group_sizes
       elif self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
         return group_sizes
@@ -1492,42 +1493,28 @@ def extract_vma(tensor):
       # mode on a TPU target breaks check_vma and bloats HBM temporaries.
       megablox_interpret = self.mesh.devices.flat[0].platform != "tpu"
 
-      # We support three implementations for gmm - tokamax, older forked kernel, or jax.lax.ragged_dot
-      # For quantized tokamax we call a forked version that supports our quantization recipes.
-      if self.config.use_tokamax_gmm:
-        # tokamax gmm v1 (quantized) or tokamax gmm v2 (unquantized)
-        # tokamax gmm v2 (quantized) not supported yet
-        if self.config.quantization or self.config.use_gmm_v2:
-          output = mblx.gmm(
-              lhs=inputs,
-              rhs=kernel,
-              group_sizes=group_sizes,
-              preferred_element_type=self.dtype,
-              tiling=tiling,
-              group_offset=group_offset,
-              lhs_quantize_dtype=lhs_quantize_dtype,
-              rhs_quantize_dtype=rhs_quantize_dtype,
-              use_qwix_quantization=self.config.use_qwix_quantization,
-              use_tokamax_backend=self.config.use_tokamax_gmm,
-              weight_gather_axes=weight_gather_axes,
-              lhs_vma_axes=lhs_vma_axes,
-              rhs_vma_axes=rhs_vma_axes,
-              use_gmm_v2=self.config.use_gmm_v2,
-              partial_sum=partial_sum,
-              interpret=megablox_interpret,
-          )
-        else:  # tokamax (unquantized)
-          output = tokamax.ragged_dot(
-              lhs=inputs,
-              rhs=kernel,
-              group_sizes=tokamax_group_sizes,
-              precision=jax.lax.Precision.DEFAULT,
-              preferred_element_type=self.dtype,
-              implementation="mosaic",
-              # `group_offset` is not yet supported
-              group_offset=None,
-          )
-      elif self.config.megablox:  # Older forked megablox
+      # We support various implementations for gmm - tokamax gmm (v1, v2), older forked megablox, or jax.lax.ragged_dot
+      # Determine whether we can use: tokamax gmm v1 (quantized)
+      is_tokamax_v1_unquantized = (
+          self.config.use_tokamax_gmm and not self.config.quantization and not self.config.use_gmm_v2
+      )
+      # Use custom vjp: tokamax gmm v1 (quantized), tokamax gmm v2 (quantized, unquantized), older forked megablox
+      use_custom_vjp_gmm = self.config.use_tokamax_gmm or self.config.megablox
+
+      if is_tokamax_v1_unquantized:
+        # tokamax v1 (unquantized)
+        output = tokamax.ragged_dot(
+            lhs=inputs,
+            rhs=kernel,
+            group_sizes=tokamax_group_sizes,
+            precision=jax.lax.Precision.DEFAULT,
+            preferred_element_type=self.dtype,
+            implementation="mosaic",
+            # `group_offset` is not yet supported
+            group_offset=None,
+        )
+      elif use_custom_vjp_gmm:
+        # tokamax gmm v1 (quantized), tokamax gmm v2 (quantized, unquantized), older forked megablox
         output = mblx.gmm(
             lhs=inputs,
             rhs=kernel,
@@ -1537,14 +1524,17 @@ def extract_vma(tensor):
             group_offset=group_offset,
             lhs_quantize_dtype=lhs_quantize_dtype,
             rhs_quantize_dtype=rhs_quantize_dtype,
-            use_qwix_quantization=self.config.use_qwix_quantization,
+            use_qwix_quantization=bool(self.config.quantization) and self.config.use_qwix_quantization,
             use_tokamax_backend=self.config.use_tokamax_gmm,
             weight_gather_axes=weight_gather_axes,
             lhs_vma_axes=lhs_vma_axes,
             rhs_vma_axes=rhs_vma_axes,
+            use_gmm_v2=self.config.use_gmm_v2,
+            partial_sum=partial_sum,
             interpret=megablox_interpret,
         )
-      else:  # jax.lax.ragged_dot
+      else:
+        # jax.lax.ragged_dot
         output = jax_ragged_dot_gmm(
             inputs,
             kernel,
@@ -1553,6 +1543,7 @@ def extract_vma(tensor):
             expert_assignments,
             padding_amount,
         )
+
       if padding_amount > 0:
         output = output[: orig_inputs_shape[0]]
       return output
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 9d076abceb..5ff713f861 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -1742,7 +1742,7 @@ def __call__(
               policy = self.get_remat_policy()
               mock_params = self._build_linen_params(self.moe_layers)
 
-              if cfg.use_qwix_quantization and not cfg.use_manual_quantization:
+              if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization:
                 y = deepseek_batchsplit_fp8.scan_batch_split_layers(
                     y,
                     mock_params,
diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py
index e610fe49b8..c2c3aba722 100644
--- a/src/maxtext/layers/quantizations.py
+++ b/src/maxtext/layers/quantizations.py
@@ -658,7 +658,7 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"):
     # The pure JAX version of batch-split that uses manual quantization for dot general.
     return None
 
-  if config.use_qwix_quantization:
+  if config.quantization and config.use_qwix_quantization:
     return None
   quant_cfg = _get_quant_config(config)
   if quant_cfg:
@@ -813,6 +813,7 @@ def make_qt_rule(dtype) -> list[qwix.QtRule]:
             act_qtype=dtype,
             bwd_qtype=dtype,
             bwd_weight_grad_tile_size=1 / config.quantization_local_shard_count,
+            disable_channelwise_axes=False,
             op_names=("dot_general",),
         )
     ]
@@ -856,7 +857,7 @@ def get_qt_provider(config):
 def maybe_quantize_model(model, config):
   """Quantize the model if quantization is enabled."""
   # Batch split is not using Qwix's interception feature but manual plumbing
-  if config.use_qwix_quantization and not config.use_batch_split_schedule:
+  if config.quantization and config.use_qwix_quantization and not config.use_batch_split_schedule:
     quantization_provider = get_qt_provider(config)
     if quantization_provider:
       if config.pure_nnx:
@@ -912,13 +913,13 @@ def _get_max_min(target_dtype):
 
 
 def manual_quantize(tensor: jax.Array, dtype: jax.typing.DTypeLike, calibration_method: str) -> qwix.QArray:
-  """Manually quantizes a tensor based on a fixed calibration method.
+  """Manually quantizes a tensor based on per-tensor scaling with symmetric fixed range calibration.
 
   Args:
     tensor: The tensor to quantize.
     dtype: The logical type of the quantized value, e.g. jnp.float8_e4m3fn
-    calibration_method: A string specifying the calibration method. Expected
-      format is "fixed,{scale},{max_val}". e.g., "fixed,-224,224"
+    calibration_method: A string specifying the calibration method. Currently only support
+    symmetric fixed range calibration: Expected format is "fixed,{-max_val},{max_val}".
 
   Returns:
     A qwix.QArray containing the quantized value and the scale.
@@ -932,13 +933,16 @@ def manual_quantize(tensor: jax.Array, dtype: jax.typing.DTypeLike, calibration_
     raise ValueError("calibration_method cannot be None for manual quantization")
   if not calib_method.startswith("fixed"):
     # we can use static scale for weight/activation, but grad usually needs dynamic
-    raise ValueError("Only static scale quantization is supported, but got" f" {calib_method}")
+    raise ValueError(f"Only static scale quantization is supported, but got {calib_method}")
   parts = calib_method.split(",")
   if len(parts) != 3:
-    raise ValueError(f"Unexpected format for weight calibration method: {calib_method}")
+    raise ValueError(f"Unexpected format for calibration method: {calib_method}")
 
   dtype_max, dtype_min = _get_max_min(dtype)
-  max_val = float(parts[2])
+  min_val, max_val = float(parts[1]), float(parts[2])
+  if max_val <= 0 or min_val != -max_val:
+    raise ValueError(f"Unexpected format for calibration method: {calib_method}")
+
   scale = max_val / dtype_max
   scale = jnp.where(scale == 0, 1.0, scale)
   # scale must be converted to a tensor because grad has reduced axes.
@@ -947,7 +951,8 @@ def manual_quantize(tensor: jax.Array, dtype: jax.typing.DTypeLike, calibration_
   max_bound = _make_scale_tensor(dtype_max, tensor)
   q_tensor = jnp.clip(tensor / scale_tensor, min_bound, max_bound).astype(dtype)
 
-  # get scale for QArray
+  # get scale for QArray.
+  # per-tensor scaling: same scale for each axis
   scale_shape = [1] * tensor.ndim
   # It must stay fully replicated for the backward pass and Pallas.
   scale_tensor_qpl = jnp.full(scale_shape, scale, dtype=tensor.dtype)
diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py
index e0e09f4175..8d4f6d5ca0 100644
--- a/src/maxtext/models/deepseek.py
+++ b/src/maxtext/models/deepseek.py
@@ -450,7 +450,7 @@ def __call__(
     # in `Decoder`, since they will never be executed together.
     if self.config.use_batch_split_schedule:
       # The older version of batch-split that fully uses qwix quantization.
-      if self.config.use_qwix_quantization and not self.config.use_manual_quantization:
+      if self.config.quantization and self.config.use_qwix_quantization and not self.config.use_manual_quantization:
         activation_pspec = jax.sharding.PartitionSpec(
             ("data", "fsdp", "fsdp_transpose", "expert", "context"),
             None,
diff --git a/src/maxtext/models/deepseek_batchsplit_fp8.py b/src/maxtext/models/deepseek_batchsplit_fp8.py
index 0b6410c9e5..0f86667861 100644
--- a/src/maxtext/models/deepseek_batchsplit_fp8.py
+++ b/src/maxtext/models/deepseek_batchsplit_fp8.py
@@ -949,7 +949,7 @@ def gmm(
       preferred_element_type,
       weight_gather_axes,
   ):
-    if config.use_qwix_quantization:
+    if config.quantization and config.use_qwix_quantization:
       output = megablox.gmm(
           lhs=inputs,
           rhs=kernel,
@@ -1000,7 +1000,7 @@ def gmm(
       config.wo_tile_drhs_embed_dim,  # Called n in megablox, and indeed is the RHS batch dim
   )
 
-  if config.use_qwix_quantization:
+  if config.quantization and config.use_qwix_quantization:
     gating_pspec, linear_pspec = moe_lib.get_batchsplit_init_kernel_axes()
     w0_pspec = nn.logical_to_mesh_axes(gating_pspec)
     wo_pspec = nn.logical_to_mesh_axes(linear_pspec)
@@ -1152,7 +1152,7 @@ def process_activations(
       None,
       None,
   )
-  if config.use_qwix_quantization:
+  if config.quantization and config.use_qwix_quantization:
     gating_pspec, linear_pspec = moe_lib.get_batchsplit_init_kernel_axes()
     gating_pspec = nn.logical_to_mesh_axes(gating_pspec)
     linear_pspec = nn.logical_to_mesh_axes(linear_pspec)
diff --git a/tests/integration/sparsity_test.py b/tests/integration/sparsity_test.py
index 7154ed00f8..685d45afc4 100644
--- a/tests/integration/sparsity_test.py
+++ b/tests/integration/sparsity_test.py
@@ -69,7 +69,7 @@ def test_different_quant_sparsity_configs(self, quantization: str, use_sparsity:
         "shared_experts=1",
         "sparse_matmul=True",
         "megablox=False",
-        f'quantization="{quantization}"',
+        f"quantization={quantization}",
         "use_qwix_quantization=True",
         "per_device_batch_size=2",
         "max_target_length=128",
diff --git a/tests/integration/tokamax_test.py b/tests/integration/tokamax_test.py
index d2a4d291ad..e7d41e82e8 100644
--- a/tests/integration/tokamax_test.py
+++ b/tests/integration/tokamax_test.py
@@ -11,8 +11,7 @@
 # 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.
-
-"""Test for tokamax gmm and splash."""
+"""Test for tokamax gmm."""
 
 import os
 import tempfile
@@ -25,38 +24,51 @@
 train_main = train.main
 gettempdir = tempfile.gettempdir
 
+# TODO: revert before merge
+from absl import logging
+
+logging.set_verbosity(logging.INFO)
+
 
 @pytest.mark.integration_test
 class Train(parameterized.TestCase):
-  """Test for tokamax gmm and splash."""
+  """Smoke test for tokamax gmm.
+
+  Similar to `train_using_ragged_dot_smoke_train.py`
+  """
 
   @parameterized.named_parameters(
       {
-          "testcase_name": "gmm bf16",
-          "quantization": "",
-          "use_gmm_v2": False,
-      },
-      {
-          "testcase_name": "gmm fp8",
-          "quantization": "fp8_full",
-          "use_gmm_v2": False,
-      },
-      {
-          "testcase_name": "gmm v2 bf16",
-          "quantization": "",
-          "use_gmm_v2": True,
-      },
+          "testcase_name": f"{base_name}_ep{ici_expert_parallelism}",
+          "quantization": quantization,
+          "use_gmm_v2": use_gmm_v2,
+          "ici_expert_parallelism": ici_expert_parallelism,
+      }
+      for base_name, quantization, use_gmm_v2, ici_expert_parallelism in [
+          ("tokamax_v1_bf16", "", False, 1),
+          ("tokamax_v1_fp8", "fp8_full", False, 1),
+          ("tokamax_v2_bf16", "", True, 1),
+          ("tokamax_v2_fp8", "fp8_full", True, 1),
+          ("tokamax_v2_bf16", "", True, 2),
+          ("tokamax_v2_fp8", "fp8_full", True, 2),
+      ]
   )
   @pytest.mark.tpu_only
-  def test_different_configs(self, quantization: str, use_gmm_v2: bool):
+  def test_smoke_train(
+      self,
+      quantization: str,
+      use_gmm_v2: bool,
+      ici_expert_parallelism: int,
+  ):
     """Smoke train with small config."""
+    sharding_tolerance = 0.22 if ici_expert_parallelism > 1 else 2e-2
     test_tmpdir = os.environ.get("TEST_TMPDIR", gettempdir())
     outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", test_tmpdir)
     args = [
         None,
         get_test_config_path(),
         f"base_output_directory={test_tmpdir}",
-        "run_name=tokamax_test",
+        "run_name=test_smoke_train",
         # model
         "base_emb_dim=256",
         "base_num_query_heads=1",
@@ -69,6 +81,8 @@ def test_different_configs(self, quantization: str, use_gmm_v2: bool):
         "attention_type=mla",
         "num_experts=2",
         "shared_experts=1",
+        f"ici_expert_parallelism={ici_expert_parallelism}",
+        f"sharding_tolerance={sharding_tolerance}",
         # tokamax gmm
         "sparse_matmul=True",
         "megablox=False",
@@ -96,16 +110,17 @@ def test_different_configs(self, quantization: str, use_gmm_v2: bool):
         # tokamax splash
         "max_target_length=128",
         "attention=flash",
-        "use_tokamax_splash=True",
+        "use_tokamax_splash=False",
         # quantization
         f"quantization={quantization}",
-        f"use_qwix_quantization={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",
         # train
         "per_device_batch_size=1",
         "dataset_type=synthetic",
-        "steps=3",
+        "steps=2",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
         "enable_checkpoint_cloud_logger=False",
diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py
index 01c6640963..c42fabb0ef 100644
--- a/tests/unit/configs_value_test.py
+++ b/tests/unit/configs_value_test.py
@@ -308,40 +308,6 @@ def test_initialize_pydantic_bad_keys(self):
           ]
       )
 
-  def test_gmm_v2_quantization_disallowed(self):
-    """Tests that use_gmm_v2=True with quantization enabled is disallowed."""
-    argv = [
-        "",
-        _BASE_CONFIG_PATH,
-        "run_name=test",
-        "use_gmm_v2=true",
-        "quantization=fp8_full",
-    ]
-    with self.assertRaises(pydantic.ValidationError):
-      pyconfig.initialize(argv)
-
-    argv_qwix = [
-        "",
-        _BASE_CONFIG_PATH,
-        "run_name=test",
-        "use_gmm_v2=true",
-        "use_qwix_quantization=true",
-    ]
-    with self.assertRaises(pydantic.ValidationError):
-      pyconfig.initialize(argv_qwix)
-
-  def test_gmm_v2_requires_tokamax_gmm(self):
-    """Tests that use_gmm_v2=True requires use_tokamax_gmm=True."""
-    argv = [
-        "",
-        _BASE_CONFIG_PATH,
-        "run_name=test",
-        "use_gmm_v2=true",
-        "use_tokamax_gmm=false",
-    ]
-    with self.assertRaises(pydantic.ValidationError):
-      pyconfig.initialize(argv)
-
   def test_safetensors_dynamic_disallows_single_controller(self):
     """Tests that source_checkpoint_layout=safetensors_dynamic disallows enable_single_controller=True."""
     argv = [
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index 96f76106ed..af27486d43 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -14,6 +14,8 @@
 """Mixture of Experts (MoE) tests."""
 
 import unittest
+from absl.testing import parameterized
+import pytest
 
 from flax import nnx
 import flax.linen as nn
@@ -21,6 +23,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import qwix
 from jax.sharding import Mesh
 from maxtext.configs import pyconfig
 from maxtext.common.common_types import Config, DType
@@ -32,7 +35,38 @@
 from maxtext.utils import max_logging, maxtext_utils
 from maxtext.utils.sharding import remove_expert_from_partition_spec
 from tests.utils.test_helpers import get_test_config_path
-import pytest
+
+
+def compare_tree(a, b, relative_norm_diff_threshold=1e-02):
+  """
+  Compute the relative norm difference between two pytrees and enforce threshold.
+  """
+  # the first key is customized prefix define by user.
+  # the rest of the keys are extracted from the path tuple.
+  leaves_a, tree_def_a = jax.tree_util.tree_flatten_with_path(a)
+  leaves_b, tree_def_b = jax.tree_util.tree_flatten_with_path(b)
+
+  if tree_def_a != tree_def_b:
+    raise ValueError("Reference and actual pytrees must have the same structure.")
+
+  log_lines = ["CALCULATE DIFF"]
+  for (path, val_a), (_, val_b) in zip(leaves_a, leaves_b):
+    path = "-".join([k.key for k in path if hasattr(k, "key")]) or "root"
+
+    assert val_a.dtype in (jnp.float32, jnp.bfloat16, jnp.float16)
+    val_a = val_a.astype(jnp.float32)
+    val_b = val_b.astype(jnp.float32)
+
+    max_abs_diff = jnp.max(jnp.abs(val_a - val_b))
+    relative_norm_diff = jnp.linalg.norm(val_a - val_b) / jnp.linalg.norm(val_a)
+
+    log_line = f"{path}" f" | max_abs_diff: {max_abs_diff}" f" | relative_norm_diff: | {relative_norm_diff}"
+    log_lines.append(log_line)
+    assert (
+        relative_norm_diff < relative_norm_diff_threshold
+    ), f"relative_norm_diff exceeds {relative_norm_diff_threshold=}"
+  diff_summary = "\n".join(log_lines)
+  return diff_summary
 
 
 def assert_moe_close(actual, expected, dtype):
@@ -403,7 +437,7 @@ def get_moe_loop(
   return module
 
 
-class RoutedMoeTest(unittest.TestCase):
+class RoutedMoeTest(parameterized.TestCase):
   """Routed Mixture of Experts test."""
 
   def get_expected_output(self, rng, hidden_states, cfg, mesh):
@@ -1390,6 +1424,184 @@ def test_large_ep_worst_case(self):
     )
     self.assertEqual(expected_ragged_buffer, actual_ragged_buffer)
 
+  @parameterized.named_parameters(
+      {
+          "testcase_name": f"{base_name}_ep{ici_expert_parallelism}",
+          "quantization": quantization,
+          "use_gmm_v2": use_gmm_v2,
+          "ici_expert_parallelism": ici_expert_parallelism,
+      }
+      for base_name, quantization, use_gmm_v2, ici_expert_parallelism in [
+          ("tokamax_v1_bf16", "", False, 1),
+          ("tokamax_v1_fp8", "fp8_full", False, 1),
+          ("tokamax_v2_bf16", "", True, 1),
+          ("tokamax_v2_fp8", "fp8_full", True, 1),
+          ("tokamax_v2_bf16", "", True, 4),
+          ("tokamax_v2_fp8", "fp8_full", True, 4),
+      ]
+  )
+  @pytest.mark.tpu_only
+  def test_gmm_grad_equivalence(
+      self,
+      quantization: str,
+      use_gmm_v2: bool,
+      ici_expert_parallelism: int,
+      **kwargs,
+  ):
+    megablox = False
+    use_tokamax_gmm = True
+
+    def _build_cfg(
+        sparse_matmul,
+        quantization,
+        megablox,
+        use_tokamax_gmm,
+        use_gmm_v2,
+        ici_expert_parallelism,
+    ):
+      return pyconfig.initialize(
+          [None, get_test_config_path()],
+          run_name="gmm_grad_equivalence_test",
+          enable_checkpointing=False,
+          model_name="mixtral-8x7b",
+          weight_dtype="float32",
+          dtype="bfloat16",
+          per_device_batch_size=1,
+          max_target_length=256,
+          float32_gate_logits=True,
+          ici_expert_parallelism=ici_expert_parallelism,
+          sparse_matmul=sparse_matmul,
+          megablox=megablox,
+          use_tokamax_gmm=use_tokamax_gmm,
+          use_gmm_v2=use_gmm_v2,
+          quantization=quantization,
+          use_qwix_quantization=True,
+          weight_quantization_calibration_method="fixed,-224,224",
+          act_quantization_calibration_method="fixed,-224,224",
+          bwd_quantization_calibration_method="absmax",
+          wi_tile_fwd_batch_seq=128,
+          wi_tile_dlhs_batch_seq=128,
+          wi_tile_dlhs_embed_dim=256,
+          wi_tile_drhs_batch_seq=128,
+          wo_tile_fwd_batch_seq=128,
+          wo_tile_fwd_embed_dim=256,
+          wo_tile_dlhs_batch_seq=128,
+          wo_tile_dlhs_mlp_dim=256,
+          wo_tile_drhs_batch_seq=128,
+      )
+
+    def _build_model(cfg, mesh):
+      model = moe.get_routed_moe(
+          name="MoeBlock",
+          config=cfg,
+          num_experts=cfg.num_experts,
+          num_experts_per_tok=cfg.num_experts_per_tok,
+          mesh=mesh,
+          kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+          kernel_axes=("embed", "mlp"),
+          intermediate_dim=cfg.mlp_dim,
+          dtype=cfg.dtype,
+      )
+
+      # quantize the model with qwix rule for gmm
+      if cfg.quantization:
+        if not (cfg.quantization == "fp8_full" and cfg.use_qwix_quantization):
+          raise ValueError("Only fp8_full with qwix quantization is supported for MoE testing.")
+
+        # Similar to `quantizations.get_fp8_full_qwix_rule_w_sparsity`
+        def get_fp8_full_qwix_rule_for_test(config: Config):
+          return [
+              qwix.QtRule(
+                  module_path=".*",
+                  weight_qtype=jnp.float8_e4m3fn,
+                  act_qtype=jnp.float8_e4m3fn,
+                  bwd_qtype=jnp.float8_e5m2,
+                  weight_calibration_method=config.weight_quantization_calibration_method,
+                  act_calibration_method=config.act_quantization_calibration_method,
+                  bwd_calibration_method=config.bwd_quantization_calibration_method,
+                  op_names=("gmm", "ragged_dot"),
+              ),
+          ]
+
+        quantization_rule = get_fp8_full_qwix_rule_for_test(cfg)
+        quantization_provider = qwix.QtProvider(quantization_rule)
+        model = qwix.quantize_model(model, quantization_provider)
+
+      return model
+
+    def _loss_and_grad(model, variables, hidden_states):
+      def loss_fn(params, x):
+        out, lb_loss, _ = model.apply({"params": params}, x)
+        loss = jnp.mean(out.astype(jnp.float32) ** 2)
+        if lb_loss is not None:
+          loss = loss + lb_loss.astype(jnp.float32)
+        return loss, out
+
+      return jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(variables["params"], hidden_states)
+
+    rng = jax.random.PRNGKey(2345)
+    rng_model, rng_hidden_states = jax.random.split(rng)
+    device_count = jax.device_count()
+
+    # Reference run: dense matmul, no quantization, EP=1
+    cfg_ref = _build_cfg(
+        sparse_matmul=False,
+        quantization="",
+        megablox=False,
+        use_tokamax_gmm=False,
+        use_gmm_v2=False,
+        ici_expert_parallelism=1,
+    )
+    # Use normal distribution to generate realistic variances and negative values
+    # to guarantee the quantization scale != 1.0, which catches scale-dropping bugs.
+    hidden_states = jax.random.normal(
+        rng_hidden_states,
+        (int(cfg_ref.per_device_batch_size) * device_count, cfg_ref.max_target_length, cfg_ref.base_emb_dim),
+        dtype=cfg_ref.dtype,
+    )
+    devices_array_ref = maxtext_utils.create_device_mesh(cfg_ref)
+    mesh_ref = Mesh(devices_array_ref, cfg_ref.mesh_axes)
+    model_ref = _build_model(cfg_ref, mesh_ref)
+
+    with jax.set_mesh(mesh_ref), nn_partitioning.axis_rules(cfg_ref.logical_axis_rules):
+      variables = model_ref.init({"params": rng_model, "dropout": rng_model}, hidden_states)
+      (_, output_ref), (grads_ref, x_grad_ref) = _loss_and_grad(model_ref, variables, hidden_states)
+
+    # Target run: custom configuration (sparse, GMM, EP, quantization)
+    cfg_tgt = _build_cfg(
+        sparse_matmul=True,
+        quantization=quantization,
+        megablox=megablox,
+        use_tokamax_gmm=use_tokamax_gmm,
+        use_gmm_v2=use_gmm_v2,
+        ici_expert_parallelism=ici_expert_parallelism,
+    )
+    devices_array_tgt = maxtext_utils.create_device_mesh(cfg_tgt)
+    mesh_tgt = Mesh(devices_array_tgt, cfg_tgt.mesh_axes)
+    model_tgt = _build_model(cfg_tgt, mesh_tgt)
+
+    with jax.set_mesh(mesh_tgt), nn_partitioning.axis_rules(cfg_tgt.logical_axis_rules):
+      # Re-initialize variables for the target run to ensure they are sharded correctly
+      # for the target mesh, but use the same RNG to get the same initial values.
+      variables_tgt = model_tgt.init({"params": rng_model, "dropout": rng_model}, hidden_states)
+      (_, output_tgt), (grads_tgt, x_grad_tgt) = _loss_and_grad(model_tgt, variables_tgt, hidden_states)
+
+    # Compare results
+    tree_ref = {
+        "output": output_ref,
+        "state_grad": x_grad_ref,
+        "var_grad": grads_ref,
+    }
+    tree_tgt = {
+        "output": output_tgt,
+        "state_grad": x_grad_tgt,
+        "var_grad": grads_tgt,
+    }
+
+    relative_norm_diff_threshold = 0.22 if quantization else 0.012
+    diff_summary = compare_tree(tree_ref, tree_tgt, relative_norm_diff_threshold)
+    max_logging.log("\n" + diff_summary)
+
 
 def make_moe(cfg, mesh):
   return moe.RoutedMoE(