From 58e4ddc6148118f7f791688e3645a31190967aba Mon Sep 17 00:00:00 2001
From: Xuefeng Gu <xfgu@google.com>
Date: Wed, 3 Jun 2026 23:36:37 +0000
Subject: [PATCH] [MaxText] Implement MoE routing mismatch measurement

Calculate the mismatch rate between model-calculated routing and externally
supplied forced_routed_experts. This is limited to Qwen3 and Qwen3.5 models
in training mode.

The mismatch rate is logged as 'learning/routing_mismatch_rate'.
---
 .../integration/tunix/tunix_adapter.py        |   2 +
 src/maxtext/layers/decoders.py                |  59 ++++++-
 src/maxtext/layers/moe.py                     | 150 ++++++++++++++++--
 src/maxtext/layers/nnx_decoders.py            |  89 +++++++++--
 src/maxtext/models/models.py                  |   5 +
 src/maxtext/models/qwen3.py                   |  25 ++-
 src/maxtext/models/qwen3_5.py                 |  18 ++-
 src/maxtext/trainers/pre_train/train.py       |  17 ++
 .../post_training/unit/tunix_adapter_test.py  |  37 ++++-
 tests/unit/forced_routing_test.py             | 129 +++++++++++++++
 10 files changed, 496 insertions(+), 35 deletions(-)
 create mode 100644 tests/unit/forced_routing_test.py

diff --git a/src/maxtext/integration/tunix/tunix_adapter.py b/src/maxtext/integration/tunix/tunix_adapter.py
index 68da164112..89cbb13e8b 100644
--- a/src/maxtext/integration/tunix/tunix_adapter.py
+++ b/src/maxtext/integration/tunix/tunix_adapter.py
@@ -70,6 +70,7 @@ def __call__(
       attention_mask: Optional[Array],  # [B, L, L] or None
       decoder_segment_ids: Optional[Array] = None,
       output_hidden_states: bool = False,  # ignored
+      forced_routed_experts: Optional[Array] = None,
   ) -> Tuple[Array, None]:
     """Forward compatible with Tunix Trainers default loss.
     Returns logits, None.
@@ -80,6 +81,7 @@ def __call__(
         decoder_input_tokens=input_tokens,
         decoder_positions=positions,
         decoder_segment_ids=decoder_segment_ids,
+        forced_routed_experts=forced_routed_experts,
     )
     return logits, None
 
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index b28b6dcb7a..c149860324 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -856,6 +856,7 @@ def __call__(
       kv_caches: list[jax.Array] | None = None,
       attention_metadata=None,
       deepstack_visual_embeds: None | list[jnp.ndarray] = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     cfg = self.config
     mesh = self.mesh
@@ -877,6 +878,11 @@ def __call__(
       y = mhc_expand(y)
 
     policy = self.get_remat_policy()
+    supports_forced_routing = cfg.decoder_block in (
+        DecoderBlockType.QWEN3_MOE,
+        DecoderBlockType.QWEN3_NEXT,
+        DecoderBlockType.QWEN3_5,
+    )
     RemattedBlockLayers = self.set_remat_policy(self.decoder_layer, policy)
     # scan does not support kwargs in layer call, passing broadcast_args as positional arg
     broadcast_args = (
@@ -1075,6 +1081,17 @@ def __call__(
           current_broadcast_args = list(broadcast_args)
           current_in_axes_tuple = list(in_axes_tuple)
 
+          if supports_forced_routing and forced_routed_experts is not None:
+            # Transpose [B, L, N, E] -> [N, B, L, E] for scan
+            forced_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+            cycle_interval = cfg.inhomogeneous_layer_cycle_interval
+            forced_routed_experts_t = jnp.reshape(
+                forced_experts_t,
+                (scan_length, cycle_interval) + forced_experts_t.shape[1:],
+            )
+          else:
+            forced_routed_experts_t = None
+
           if kv_caches is not None:
             # Stack kv_caches for scan: [num_layers, ...]
             stacked_kv_cache = jnp.stack(kv_caches, axis=0)
@@ -1088,6 +1105,14 @@ def __call__(
             current_broadcast_args.extend([None, None, None, attention_metadata])
             current_in_axes_tuple.extend([nn.broadcast] * 4)
 
+            if supports_forced_routing:
+              if forced_routed_experts_t is not None:
+                current_broadcast_args.append(forced_routed_experts_t)
+                current_in_axes_tuple.append(0)
+              else:
+                current_broadcast_args.append(None)
+                current_in_axes_tuple.append(nn.broadcast)
+
             max_logging.info(f"DEBUG: len(current_broadcast_args)={len(current_broadcast_args)}")
             max_logging.info(f"DEBUG: current_broadcast_args={[type(a) for a in current_broadcast_args]}")
 
@@ -1109,8 +1134,19 @@ def __call__(
               kv_caches[i] = returned_kv_cache[i]
           else:
             # Fallback to old behavior if kv_caches is None (not vLLM RPA)
-            current_broadcast_args.append(None)
-            current_in_axes_tuple.append(nn.broadcast)
+            if supports_forced_routing:
+              current_broadcast_args.extend([None, None, None, None])
+              current_in_axes_tuple.extend([nn.broadcast] * 4)
+
+              if forced_routed_experts_t is not None:
+                current_broadcast_args.append(forced_routed_experts_t)
+                current_in_axes_tuple.append(0)
+              else:
+                current_broadcast_args.append(None)
+                current_in_axes_tuple.append(nn.broadcast)
+            else:
+              current_broadcast_args.append(None)
+              current_in_axes_tuple.append(nn.broadcast)
 
             y, _ = self.scan_decoder_layers(
                 cfg,
@@ -1139,6 +1175,11 @@ def __call__(
               global_layer_idx = global_layer_idx_offset + index
               kv_cache = kv_caches[index] if kv_caches is not None else None
               input_tokens = decoder_input_tokens if cfg.engram_layers else None
+              extra_kwargs = {}
+              if layer_prefix == "moe_layers" and forced_routed_experts is not None:
+                forced_experts = forced_routed_experts[:, :, index, :]
+                extra_kwargs["forced_routed_experts"] = forced_experts
+
               y, kv_cache = layer(
                   config=cfg,
                   mesh=mesh,
@@ -1157,6 +1198,7 @@ def __call__(
                   kv_cache=kv_cache,
                   attention_metadata=attention_metadata,
                   decoder_input_tokens=input_tokens,
+                  **extra_kwargs,
               )
               if kv_caches is not None and kv_cache is not None:
                 kv_caches[index] = kv_cache
@@ -1176,6 +1218,7 @@ def __call__(
               slot=slot,
           )
         else:
+          moe_lyr_idx = 0
           for lyr in range(cfg.num_decoder_layers):
             RemattedBlockLayer = RemattedBlockLayers[0]
             layer_kwargs = {}
@@ -1212,6 +1255,17 @@ def __call__(
             layer = RemattedBlockLayer(
                 config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs
             )
+            current_forced_routed_experts = None
+            if supports_forced_routing and forced_routed_experts is not None:
+              current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :]
+              moe_lyr_idx += 1
+            elif supports_forced_routing:
+              moe_lyr_idx += 1
+
+            extra_kwargs = {}
+            if supports_forced_routing and current_forced_routed_experts is not None:
+              extra_kwargs["forced_routed_experts"] = current_forced_routed_experts
+
             y, returned_cache = layer(
                 y,
                 decoder_segment_ids,
@@ -1222,6 +1276,7 @@ def __call__(
                 slot=slot,
                 kv_cache=kv_cache,
                 attention_metadata=attention_metadata,
+                **extra_kwargs,
                 **layer_call_kwargs,
             )
             if kv_caches is not None and returned_cache is not None:
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 020956098c..b13519b04e 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -90,6 +90,11 @@ class RouteOutput:
   bias_updates: Optional[jax.Array]
   # Shape [local experts], tracks number of local tokens routed to every local expert.
   local_group_sizes: Optional[jax.Array] = None
+  mismatch_rate: Optional[jax.Array] = None
+
+
+class RoutingMismatchRate(nnx.Intermediate):
+  """Variable type for storing routing mismatch rate."""
 
 
 def _sort_activations(
@@ -600,6 +605,29 @@ def __init__(
     ):
       self.wo.value = self.wo.value * self.per_expert_scale.value[:, None, None]
 
+  def _calculate_mismatch(self, top_k_indices_model, forced_experts):
+    """Calculates the mismatch rate between model routing and forced routing."""
+    # Mask out padding tokens (where forced routing is -1)
+    valid_mask = forced_experts[:, :, 0] != -1
+    # Count matching experts for each token.
+    # top_k_indices_model: [B, L, K]
+    # forced_experts: [B, L, K]
+    # We want to check if the set of experts matches, regardless of order.
+    # Since K is small (usually 1 or 2), we can do an all-to-all comparison.
+    # [B, L, K, 1] == [B, L, 1, K] -> [B, L, K, K] -> any over last axis -> [B, L, K] -> sum over K -> matches
+    matches = jnp.any(
+        jnp.expand_dims(top_k_indices_model, -1) == jnp.expand_dims(forced_experts, -2),
+        axis=-1,
+    )
+    # Sum matches for each token and normalize by K to get fraction of correct experts
+    token_match_fraction = jnp.sum(matches, axis=-1) / self.num_experts_per_tok
+    # Mismatch rate per token
+    token_mismatch = 1.0 - token_match_fraction
+    # Average mismatch over valid tokens
+    total_valid = jnp.sum(valid_mask)
+    mismatch_rate = jnp.sum(token_mismatch * valid_mask) / (total_valid + 1e-8)
+    return mismatch_rate
+
   def _maybe_shard_with_logical(self, inputs, logical_name):
     return maybe_shard_with_logical(
         inputs,
@@ -791,6 +819,7 @@ def permute(
       rngs=None,
       roll_to_expert_id=None,
       input_ids=None,
+      forced_experts: jax.Array | None = None,
   ):
     """Permute tokens to group by expert to fit gmm call."""
     # reshape inputs (batch, sequence, emb) to (batch * sequence, emb)
@@ -798,6 +827,9 @@ def permute(
     bsz_times_seq_len = inputs_shape[0] * inputs_shape[1]
     inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2]))
     weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids)
+    mismatch_rate = None
+    if forced_experts is not None:
+      mismatch_rate = self._calculate_mismatch(selected_experts, forced_experts)
     lb_loss = None
     if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing:
       softmax_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1).astype(self.dtype)
@@ -900,6 +932,7 @@ def permute(
         lb_loss,
         bias_updates,
         local_group_size,
+        mismatch_rate,
     )
 
   def unpermute(
@@ -1200,6 +1233,7 @@ def sparse_matmul(
       w1_bias,
       wo_bias,
       input_ids=None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     """Perform sparse matrix multiplication of inputs and Experts."""
 
@@ -1449,7 +1483,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids):
     ) = get_routed_moe_shardings(is_batch_sharded_by_expert, input_ids is not None)
     w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec)
 
-    def route(x, logits, pre_bias_logits, rngs, input_ids=None):
+    def route(x, logits, pre_bias_logits, rngs, input_ids=None, forced_experts=None):
       """Performs both across device and within device token routing/sorting"""
       num_ep = self.get_expert_parallelism_size()
       expert_shard_id = jax.lax.axis_index(self._expert_parallelism_name) if num_ep > 1 else 0
@@ -1479,6 +1513,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None):
             lb_loss,
             bias_updates,
             local_group_sizes,
+            mismatch_rate,
         ) = self.permute(
             x,
             logits,
@@ -1487,6 +1522,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None):
             roll_to_expert_id=num_experts_per_shard * expert_shard_id,
             rngs=rngs,
             input_ids=input_ids,
+            forced_experts=forced_experts,
         )
 
       else:
@@ -1499,7 +1535,16 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None):
             lb_loss,
             bias_updates,
             local_group_sizes,
-        ) = self.permute(x, logits, pre_bias_logits, self.config.use_custom_sort_vjp, rngs, input_ids=input_ids)
+            mismatch_rate,
+        ) = self.permute(
+            x,
+            logits,
+            pre_bias_logits,
+            self.config.use_custom_sort_vjp,
+            rngs,
+            input_ids=input_ids,
+            forced_experts=forced_experts,
+        )
 
         if num_ep > 1:
           batch_axis = self._expert_parallelism_name if is_batch_sharded_by_expert else "data"
@@ -1565,6 +1610,7 @@ def route(x, logits, pre_bias_logits, rngs, input_ids=None):
               lb_loss=lb_loss,
               bias_updates=bias_updates,
               local_group_sizes=local_group_sizes,
+              mismatch_rate=mismatch_rate,
           ),
           RouteMetadata(
               expert_shard_id=expert_shard_id,
@@ -1655,17 +1701,39 @@ def gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather):
             wo_bias_pspec,
             decoder_tokens_pspec,
             P(),  # Replicate the input key
+            gate_logits_pspec if forced_routed_experts is not None else None,
         ),
         out_specs=(
             self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", "activation_embed")),
             P(),  # Handle None or replicate the output
             P(),  # Handle None or replicate the output
+            P(),  # mismatch_rate
         ),
         check_vma=self.config.check_vma,
     )
-    def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, sharded_input_ids, rngs):
+    def wrapper(
+        x,
+        logits,
+        pre_bias_logits,
+        w0,
+        w1,
+        wo,
+        w0_bias,
+        w1_bias,
+        wo_bias,
+        sharded_input_ids,
+        rngs,
+        forced_experts,
+    ):
       batch_size, sequence_length, _ = x.shape
-      x, routing, route_metadata = route(x, logits, pre_bias_logits, rngs, input_ids=sharded_input_ids)
+      x, routing, route_metadata = route(
+          x,
+          logits,
+          pre_bias_logits,
+          rngs,
+          input_ids=sharded_input_ids,
+          forced_experts=forced_experts,
+      )
 
       if self.config.mlp_bias:
         w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias)
@@ -1804,7 +1872,7 @@ def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, s
             group_sizes=routing.group_sizes,
         )
 
-      return output, routing.lb_loss, routing.bias_updates
+      return output, routing.lb_loss, routing.bias_updates, routing.mismatch_rate
 
     if self.config.moe_fsdp_use_two_stage_all_gather:
       # Unshard on fsdp axis
@@ -1841,6 +1909,12 @@ def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, s
     gate_logits = self._maybe_shard_with_logical(gate_logits, gate_logits_axes)
     pre_bias_logits = self._maybe_shard_with_logical(pre_bias_logits, pre_bias_logits_axes)
 
+    forced_experts = (
+        self._maybe_shard_with_logical(forced_routed_experts, gate_logits_axes)
+        if forced_routed_experts is not None
+        else None
+    )
+
     w0_kernel = self._maybe_shard_with_pspec(w0_kernel, w0_pspec)
     w1_kernel = self._maybe_shard_with_pspec(w1_kernel, w1_pspec)
     wo_kernel = self._maybe_shard_with_pspec(wo_kernel, wo_pspec)
@@ -1863,6 +1937,7 @@ def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, s
         wo_bias,
         input_ids,
         self.rngs,
+        forced_experts,
     )
 
   def reshape_and_update_weights(self, weights, indices):
@@ -2123,7 +2198,8 @@ def dense_matmul(
       w1_bias,
       wo_bias,
       input_ids=None,
-  ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]:
+      forced_routed_experts: jax.Array | None = None,
+  ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array], Optional[jax.Array]]:
     """Dense matrix multiplication."""
     # gate_logits: batch, length, expert
     gate_logits = self._maybe_shard_with_logical(gate_logits, ("activation_batch_moe", "activation_length_moe", None))
@@ -2133,7 +2209,17 @@ def dense_matmul(
       pre_bias_logits = self._maybe_shard_with_logical(
           pre_bias_logits, ("activation_batch_moe", "activation_length_moe", None)
       )
+    forced_experts = None
+    if forced_routed_experts is not None:
+      forced_experts = self._maybe_shard_with_logical(
+          forced_routed_experts,
+          ("activation_batch_moe", "activation_length_moe", None),
+      )
     top_k_weights, top_k_indices = self.get_topk(gate_logits, pre_bias_logits, self.rngs, input_ids=input_ids)
+    mismatch_rate = None
+    if forced_routed_experts is not None:
+      mismatch_rate = self._calculate_mismatch(top_k_indices, forced_experts)
+
     is_llama4_decoder_layer = self.config.decoder_block == ctypes.DecoderBlockType.LLAMA4
     if is_llama4_decoder_layer:
       router_scores = jax.nn.sigmoid(top_k_weights.astype(jnp.float32)).astype(self.dtype)
@@ -2360,7 +2446,7 @@ def dense_matmul(
                   output.shape[3],
               ),
           )
-      return output, lb_loss, bias_updates
+      return output, lb_loss, bias_updates, mismatch_rate
     else:
       inputs = self._maybe_shard_with_logical(
           inputs, ("activation_batch_moe", "activation_norm_length_moe", "activation_embed_moe")
@@ -2410,7 +2496,7 @@ def dense_matmul(
             weights,
             precision=matmul_precision,
         ).astype(self.dtype)
-      return output, lb_loss, bias_updates
+      return output, lb_loss, bias_updates, mismatch_rate
 
   def fused_moe_matmul(
       self,
@@ -2420,7 +2506,7 @@ def fused_moe_matmul(
       w0_kernel=None,
       w1_kernel=None,
       fused_kernel=None,
-  ) -> tuple[jax.Array, None, None]:
+  ) -> tuple[jax.Array, None, None, None]:
     """Fused MoE via tpu_inference fused_moe_func (vllm_rpa path only).
 
     fused_moe_func handles routing, GMM, and weighted combination internally.
@@ -2474,7 +2560,7 @@ def fused_moe_matmul(
 
     # Reshape output 2D [T, D] -> 3D [B, S, D]
     output = jnp.reshape(output_2d, (batch_size, seq_len, emb_dim))
-    return output, None, None
+    return output, None, None, None
 
   def retrieve_quantized_weight(
       self,
@@ -2511,6 +2597,7 @@ def __call__(
       input_ids: jax.Array | None = None,
       gate_inputs: jax.Array | None = None,
       out_sharding: NamedSharding | None = None,
+      forced_routed_experts: jax.Array | None = None,
   ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]:
     """Executes the routed MoE block.
 
@@ -2520,6 +2607,7 @@ def __call__(
         in DeepSeek V4's early MoE layers. If None, routing relies purely on `gate_inputs` or `inputs`.
       gate_inputs: Optional alternate inputs to feed into the routing gate.
       out_sharding: Optional sharding specification for the output.
+      forced_routing_config: Optional configuration for forced routing and mismatch measurement.
 
     Returns:
       A tuple containing the MoE output, the load balance loss (if applicable),
@@ -2562,7 +2650,7 @@ def __call__(
     # weights. Hash routed layers bypass this kernel and fall back
     # to the sparse matmul implementation.
     if cfg.attention == "vllm_rpa" and not self.is_hash_routing:
-      output, lb_loss, bias_updates = self.fused_moe_matmul(
+      output, lb_loss, bias_updates, mismatch_rate = self.fused_moe_matmul(
           inputs, gate_logits, wo_kernel, w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel
       )
     elif cfg.sparse_matmul:
@@ -2578,13 +2666,37 @@ def __call__(
             w1_bias,
             wo_bias,
         )
-      output, lb_loss, bias_updates = self.sparse_matmul(
-          inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias, input_ids
+      output, lb_loss, bias_updates, mismatch_rate = self.sparse_matmul(
+          inputs,
+          gate_logits,
+          pre_bias_logits,
+          w0_kernel,
+          w1_kernel,
+          wo_kernel,
+          w0_bias,
+          w1_bias,
+          wo_bias,
+          input_ids=input_ids,
+          forced_routed_experts=forced_routed_experts,
       )
     else:
-      output, lb_loss, bias_updates = self.dense_matmul(
-          inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias, input_ids
+      output, lb_loss, bias_updates, mismatch_rate = self.dense_matmul(
+          inputs,
+          gate_logits,
+          pre_bias_logits,
+          w0_kernel,
+          w1_kernel,
+          wo_kernel,
+          w0_bias,
+          w1_bias,
+          wo_bias,
+          input_ids=input_ids,
+          forced_routed_experts=forced_routed_experts,
       )
+    if mismatch_rate is not None:
+      self.mismatch_rate = RoutingMismatchRate(mismatch_rate)
+    else:
+      self.mismatch_rate = None
     return output, lb_loss, bias_updates
 
 
@@ -2673,6 +2785,7 @@ def __call__(
       intermediate_sharding: NamedSharding | None = None,
       out_sharding: NamedSharding | None = None,
       input_ids: jax.Array | None = None,
+      forced_routed_experts: jax.Array | None = None,
   ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]:
     """Executes both the routed experts and the shared expert block.
 
@@ -2684,13 +2797,18 @@ def __call__(
       out_sharding: Optional sharding spec for the final combined output.
       input_ids: Optional token IDs corresponding to the inputs, used strictly for Hash Routing
         in DeepSeek V4's early MoE layers.
+      forced_routing_config: Optional configuration for forced routing and mismatch measurement.
 
     Returns:
       A tuple containing the combined MoE output (routed + shared),
       the load balance loss, and any routed bias updates.
     """
     routed_experts, load_balance_loss, moe_bias_updates = self.routed_moe(
-        inputs, gate_inputs=gate_inputs, out_sharding=out_sharding, input_ids=input_ids
+        inputs,
+        gate_inputs=gate_inputs,
+        out_sharding=out_sharding,
+        input_ids=input_ids,
+        forced_routed_experts=forced_routed_experts,
     )
     shared_experts = self.shared_experts(inputs, intermediate_sharding=intermediate_sharding, out_sharding=out_sharding)
     return routed_experts + shared_experts, load_balance_loss, moe_bias_updates
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 5d7262e5a4..fa4169afac 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -945,7 +945,16 @@ def pure_layer_fn(state_in, y_in):
 
     return out
 
-  def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs):
+  def _apply_layers_sequentially(
+      self,
+      layers,
+      x_in,
+      *args,
+      length: int,
+      kv_caches_stacked=None,
+      forced_routed_experts: jax.Array | None = None,
+      **kwargs,
+  ):
     """Runs the layer stack using nnx.scan.
 
     Args:
@@ -956,6 +965,7 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches
       kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...].
         When provided, the i-th slice is passed as `kv_cache=` to layer i and the
         updated caches are returned as a third element of the tuple.
+      forced_routed_experts: Optional jax.Array containing experts for mismatch measurement.
       **kwargs: Keyword args forwarded to the layer (filtered by the layer signature).
 
     Returns:
@@ -991,14 +1001,23 @@ def _extract_matching_state(template, full):
     updated_graphdef = [graphdef]
 
     use_kv = kv_caches_stacked is not None
+    use_forced = forced_routed_experts is not None
 
     def layer_fn(carry, scanned_vars):
 
       # Unpack the sliced variables for THIS layer
       if use_kv:
-        current_params, current_state, kv_cache_layer = scanned_vars
+        if use_forced:
+          current_params, current_state, kv_cache_layer, forced_routed_experts_layer = scanned_vars
+        else:
+          current_params, current_state, kv_cache_layer = scanned_vars
+          forced_routed_experts_layer = None
       else:
-        current_params, current_state = scanned_vars
+        if use_forced:
+          current_params, current_state, forced_routed_experts_layer = scanned_vars
+        else:
+          current_params, current_state = scanned_vars
+          forced_routed_experts_layer = None
         kv_cache_layer = None
 
       if self.config.parameter_memory_host_offload:
@@ -1013,6 +1032,8 @@ def layer_fn(carry, scanned_vars):
       call_kwargs = dict(valid_kwargs)
       if kv_cache_layer is not None:
         call_kwargs["kv_cache"] = kv_cache_layer
+      if forced_routed_experts_layer is not None:
+        call_kwargs["forced_routed_experts"] = forced_routed_experts_layer
 
       layer_out = layer(carry, *args, **call_kwargs)
 
@@ -1052,11 +1073,20 @@ def layer_fn(carry, scanned_vars):
         # Statically slice the parameters and state for this layer
         current_params = jax.tree.map(lambda x, i=i: x[i], params)
         current_state = jax.tree.map(lambda x, i=i: x[i], state)
+        current_forced = jax.tree.map(lambda x, i=i: x[i], forced_routed_experts) if use_forced else None
+
+        if use_forced:
+          scanned_vars = (
+              current_params,
+              current_state,
+              kv_caches_list[i],
+              current_forced,
+          )
+        else:
+          scanned_vars = (current_params, current_state, kv_caches_list[i])
 
         # Call the layer
-        current_carry, (_, updated_kv) = layer_fn_wrapped(
-            current_carry, (current_params, current_state, kv_caches_list[i])
-        )
+        current_carry, (_, updated_kv) = layer_fn_wrapped(current_carry, scanned_vars)
 
         # Update the list in-place (mutates the list passed by reference)
         kv_caches_list[i] = updated_kv
@@ -1068,7 +1098,11 @@ def layer_fn(carry, scanned_vars):
       params = nnx_ensure_scan_leading_axis(params, length)
       state = nnx_ensure_scan_leading_axis(state, length)
 
-      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state))
+      if use_forced:
+        scan_xs = (params, state, forced_routed_experts)
+      else:
+        scan_xs = (params, state)
+      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, scan_xs)
       returned_kv_stacked = None
 
     if scan_axis != 0:
@@ -1556,6 +1590,7 @@ def __call__(
       attention_metadata=None,
       deepstack_visual_embeds: None | list[jnp.ndarray] = None,
       multimodal_input: None | MultimodalInput = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     cfg = self.config
     assert decoder_input_tokens.ndim == 2  # [batch, len]
@@ -1762,11 +1797,16 @@ def __call__(
                     num_layers=num_moe,
                 )
             else:
+              if forced_routed_experts is not None:
+                forced_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+              else:
+                forced_experts_t = None
               y, self.moe_layers, _ = self._apply_layers_sequentially(
                   self.moe_layers,
                   y,
                   *layer_args,
                   length=num_moe,
+                  forced_routed_experts=forced_experts_t,
                   **layer_kwargs,
               )
         elif self.is_gemma3:
@@ -1793,6 +1833,16 @@ def __call__(
           )
         else:
           scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval)
+          if forced_routed_experts is not None:
+            forced_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+            cycle_interval = cfg.inhomogeneous_layer_cycle_interval
+            forced_routed_experts_t = jnp.reshape(
+                forced_experts_t,
+                (scan_length, cycle_interval) + forced_experts_t.shape[1:],
+            )
+          else:
+            forced_routed_experts_t = None
+
           if kv_caches is not None:
             # Pass the kv_caches list directly to avoid copying in jnp.stack,
             # which breaks vLLM PagedAttention in-place memory updates.
@@ -1803,6 +1853,7 @@ def __call__(
                 *layer_args,
                 length=scan_length,
                 kv_caches_stacked=kv_caches,
+                forced_routed_experts=forced_routed_experts_t,
                 **layer_kwargs,
             )
             # kv_caches list is updated in-place inside _apply_layers_sequentially
@@ -1812,13 +1863,14 @@ def __call__(
                 y,
                 *layer_args,
                 length=scan_length,
+                forced_routed_experts=forced_routed_experts_t,
                 **layer_kwargs,
             )
       else:
         prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg)
 
         # Hoisted function to preserve XLA cache ID
-        def pure_layer_fn(graphdef, state_in, y_in, kv_in):
+        def pure_layer_fn(graphdef, state_in, y_in, kv_in, forced_routed_experts_in):
 
           if cfg.parameter_memory_host_offload:
             state_in = jax.tree.map(
@@ -1827,11 +1879,15 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in):
             )
 
           merged_layer = nnx.merge(graphdef, state_in)
-          out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs)
+          call_kwargs = dict(layer_kwargs)
+          if forced_routed_experts_in is not None:
+            call_kwargs["forced_routed_experts"] = forced_routed_experts_in
+          out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **call_kwargs)
           return out_y, out_kv, nnx.state(merged_layer)
 
         checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
 
+        moe_lyr_idx = 0
         for lyr, layer in enumerate(self.layers):
           graphdef, state = nnx.split(layer)
           if kv_caches is not None:
@@ -1852,7 +1908,20 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in):
           if input_tokens is not None:
             layer_kwargs["decoder_input_tokens"] = input_tokens
 
-          y, kv_cache, new_state = checkpointed_fn(graphdef, state, y, kv_cache)
+          supports_forced_routing = cfg.decoder_block in (
+              DecoderBlockType.QWEN3_MOE,
+              DecoderBlockType.QWEN3_NEXT,
+              DecoderBlockType.QWEN3_5,
+          )
+
+          current_forced_routed_experts = None
+          if supports_forced_routing and forced_routed_experts is not None:
+            current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :]
+            moe_lyr_idx += 1
+          elif supports_forced_routing:
+            moe_lyr_idx += 1
+
+          y, kv_cache, new_state = checkpointed_fn(graphdef, state, y, kv_cache, current_forced_routed_experts)
           nnx.update(layer, new_state)
 
           if kv_caches is not None and kv_cache is not None:
diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py
index ac908c0f96..211921394d 100644
--- a/src/maxtext/models/models.py
+++ b/src/maxtext/models/models.py
@@ -140,6 +140,7 @@ def __call__(
       nnx_method=None,
       kv_caches: list[jax.Array] | None = None,
       attention_metadata: dict[str, Any] | None = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     """Applies Transformer decoder-branch on encoded-input and target.
 
@@ -212,6 +213,7 @@ def __call__(
         kv_caches=kv_caches,
         attention_metadata=attention_metadata,
         deepstack_visual_embeds=deepstack_visual_embeds,
+        forced_routed_experts=forced_routed_experts,
     )  # pytype: disable=wrong-keyword-args
 
     # If we are initializing the model AND MTP is enabled, we must create
@@ -452,6 +454,7 @@ def __call__(
       decoder_target_mask: jax.Array | None = None,
       kv_caches: list[jax.Array] | None = None,
       attention_metadata: dict[str, Any] | None = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     """Applies the Zero-1 FSDP wrapped Transformer model.
 
@@ -548,6 +551,7 @@ def __call__(
           kv_caches=kv_caches,
           attention_metadata=attention_metadata,
           deepstack_visual_embeds=deepstack_visual_embeds,
+          forced_routed_experts=forced_routed_experts,
       )  # pytype: disable=wrong-keyword-args
     else:
       logits, hidden_state, kv_caches = self.decoder(
@@ -564,6 +568,7 @@ def __call__(
           attention_metadata=attention_metadata,
           deepstack_visual_embeds=deepstack_visual_embeds,
           mutable=mutable_collections,
+          forced_routed_experts=forced_routed_experts,
       )  # pytype: disable=wrong-keyword-args
 
     # If we are initializing the model AND MTP is enabled, we must create
diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py
index 3ddcb7ed12..e62d6c92b2 100644
--- a/src/maxtext/models/qwen3.py
+++ b/src/maxtext/models/qwen3.py
@@ -1115,13 +1115,19 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn
         rngs=rngs,
     )
 
-  def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Array | None]:
+  def __call__(
+      self,
+      hidden_states: Array,
+      deterministic: bool,
+      forced_routed_experts: jax.Array | None = None,
+  ) -> tuple[Array, Array | None]:
     """
     Applies the sparse MoE block to the input hidden states.
 
     Args:
       hidden_states: The input array from the previous layer. Shape: (batch, seq, embed_dim)
       deterministic: If True, disables dropout.
+      forced_routed_experts: Optional array containing experts for mismatch measurement.
 
     Returns:
       A tuple containing:
@@ -1129,7 +1135,7 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar
         - The load balancing loss from the routed experts, if applicable during training.
     """
     # 1. Apply the routed experts block.
-    routed_output, load_balance_loss, _ = self.routed_experts(hidden_states)
+    routed_output, load_balance_loss, _ = self.routed_experts(hidden_states, forced_routed_experts=forced_routed_experts)
 
     # 2. Apply the shared expert.
     shared_expert_output = self.shared_expert(hidden_states, deterministic=deterministic)
@@ -1191,6 +1197,7 @@ def __call__(
       slot: None | int = None,
       kv_cache=None,
       attention_metadata=None,
+      forced_routed_experts: jax.Array | None = None,
   ) -> tuple[Array, None]:
     """Applies the block of decoder layers to the input carry.
 
@@ -1210,6 +1217,9 @@ def __call__(
       layer = getattr(self, f"layer_{i}")
       # The second return value is kv_cache, which we ignore here because
       # it is not passed as a carry in scannable layers.
+      current_forced_routed_experts = None
+      if forced_routed_experts is not None:
+        current_forced_routed_experts = forced_routed_experts[i]
       x, _ = layer(
           x,
           decoder_segment_ids,
@@ -1220,6 +1230,7 @@ def __call__(
           slot,
           kv_cache=kv_cache,
           attention_metadata=attention_metadata,
+          forced_routed_experts=current_forced_routed_experts,
       )
 
     # The output of the block is the carry for the next scan iteration.
@@ -1307,6 +1318,7 @@ def __call__(
       slot: None | int = None,
       kv_cache: None | dict[str, Array] = None,
       attention_metadata: None | dict[str, Any] = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
     if isinstance(inputs, tuple):
@@ -1349,7 +1361,11 @@ def __call__(
     hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names)
 
     # Instantiate and call our `Qwen3NextSparseMoeBlock`.
-    mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic)
+    mlp_output, load_balance_loss = self.mlp(
+        hidden_states,
+        deterministic=deterministic,
+        forced_routed_experts=forced_routed_experts,
+    )
 
     # We sow the load balancing loss so it can be collected and added to the total loss
     # during training.
@@ -1576,6 +1592,7 @@ def __call__(
       slot: None | int = None,
       kv_cache: None | jnp.ndarray = None,
       attention_metadata: None | dict[str, Any] = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
     is_scan_carry = False
@@ -1598,7 +1615,7 @@ def __call__(
         attention_metadata=attention_metadata,
     )
 
-    mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states)
+    mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states, forced_routed_experts=forced_routed_experts)
     mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names)
     if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
       self.moe_lb_loss = nnx.Intermediate(load_balance_loss)
diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py
index de8bd1cf18..ce1d3b1006 100644
--- a/src/maxtext/models/qwen3_5.py
+++ b/src/maxtext/models/qwen3_5.py
@@ -18,6 +18,7 @@
 
 from typing import Any, cast
 
+import jax
 from jax.sharding import Mesh
 import jax.numpy as jnp
 
@@ -31,6 +32,7 @@
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_utils
 
+
 from maxtext.models.qwen3 import (
     Qwen3NextGatedDeltaNet,
     Qwen3NextFullAttention,
@@ -89,12 +91,18 @@ def __call__(
       model_mode: str,
       previous_chunk=None,
       slot: None | int = None,
+      kv_cache=None,
+      attention_metadata=None,
+      forced_routed_experts: jax.Array | None = None,
   ) -> tuple[Array, None]:
     cfg = self.config
     x = carry
 
     for i in range(cfg.inhomogeneous_layer_cycle_interval):
       layer = getattr(self, f"layer_{i}")
+      current_forced_routed_experts = None
+      if forced_routed_experts is not None:
+        current_forced_routed_experts = forced_routed_experts[i]
       x, _ = layer(
           x,
           decoder_segment_ids,
@@ -103,6 +111,9 @@ def __call__(
           model_mode,
           previous_chunk,
           slot,
+          kv_cache=kv_cache,
+          attention_metadata=attention_metadata,
+          forced_routed_experts=current_forced_routed_experts,
       )
 
     return x, None
@@ -185,6 +196,7 @@ def __call__(
       slot: None | int = None,
       kv_cache: None | dict[str, Array] = None,
       attention_metadata: None | dict[str, Any] = None,
+      forced_routed_experts: jax.Array | None = None,
   ):
     # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
     if isinstance(inputs, tuple):
@@ -227,7 +239,11 @@ def __call__(
     hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names)
 
     # Instantiate and call our `Qwen3_5SparseMoEBlock`.
-    mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic)
+    mlp_output, load_balance_loss = self.mlp(
+        hidden_states,
+        deterministic=deterministic,
+        forced_routed_experts=forced_routed_experts,
+    )
 
     # We sow the load balancing loss so it can be collected and added to the total loss
     # during training.
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 047ddb97a8..844f401206 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -278,6 +278,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
     else:
       max_logging.debug("\nNo MoE load balance loss found. Defaulting to 0.0.")
 
+  # get MoE routing mismatch rate
+  moe_mismatch_rate = None
+  if config.num_experts > 1:
+    moe_mismatch_rates = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "mismatch_rate")
+    if moe_mismatch_rates:
+      moe_mismatch_rate = jnp.mean(jnp.stack(moe_mismatch_rates))
+
   # get MoE routed bias term updates
   moe_bias_updates = None
   if config.routed_bias and config.routed_bias_update_rate > 0.0:
@@ -299,6 +306,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
       "mtp_loss": mtp_loss,
       "batch_stats": (intermediate_outputs.get("batch_stats", None) if hasattr(intermediate_outputs, "get") else None),
   }
+  if moe_mismatch_rate is not None:
+    aux["moe_mismatch_rate"] = moe_mismatch_rate
   return loss, aux
 
 
@@ -409,6 +418,12 @@ def diff_wrapper(curr_params, custom_params, rest, config, data):
   xent_sum = aux["xent_sum"]
   total_weights = aux["total_weights"]
   moe_lb_loss = aux["moe_lb_loss"]
+  if config.num_experts > 1 and "moe_mismatch_rate" in aux:
+    moe_mismatch_rate = aux["moe_mismatch_rate"]
+    if config.gradient_accumulation_steps > 1:
+      moe_mismatch_rate = moe_mismatch_rate / config.gradient_accumulation_steps
+  else:
+    moe_mismatch_rate = None
   indexer_loss = aux.get("indexer_loss", 0.0)
   z_loss = aux.get("z_loss", 0.0)
   moe_bias_updates = aux.get("moe_bias_updates")
@@ -511,6 +526,8 @@ def move(path, value):
       "learning/mtp_loss": mtp_loss,
       "learning/total_weights": total_weights,
   }
+  if moe_mismatch_rate is not None:
+    scalar_metrics["learning/routing_mismatch_rate"] = moe_mismatch_rate
   if config.use_qk_clip:
     if isinstance(model, nn.Module):
       new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config)
diff --git a/tests/post_training/unit/tunix_adapter_test.py b/tests/post_training/unit/tunix_adapter_test.py
index 0ec0f69f1d..a858ca0037 100644
--- a/tests/post_training/unit/tunix_adapter_test.py
+++ b/tests/post_training/unit/tunix_adapter_test.py
@@ -40,10 +40,18 @@ def __init__(self):
     self.config = SimpleNamespace(model_name=_STUB_MODEL_NAME)
     self.captured = {}
 
-  def __call__(self, *, decoder_input_tokens, decoder_positions, decoder_segment_ids):
+  def __call__(
+      self,
+      *,
+      decoder_input_tokens,
+      decoder_positions,
+      decoder_segment_ids,
+      forced_routed_experts=None,
+  ):
     self.captured["decoder_input_tokens"] = decoder_input_tokens
     self.captured["decoder_positions"] = decoder_positions
     self.captured["decoder_segment_ids"] = decoder_segment_ids
+    self.captured["forced_routed_experts"] = forced_routed_experts
     # Return dummy logits shaped [B, L, V=2] so the adapter has something to forward.
     b, l = decoder_input_tokens.shape
     return jnp.zeros((b, l, 2), dtype=jnp.float32)
@@ -116,7 +124,10 @@ def test_passes_through_explicit_segment_ids_unchanged(self):
 
     adapter(input_tokens, positions, None, None, decoder_segment_ids=explicit_seg)
 
-    np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(explicit_seg))
+    np.testing.assert_array_equal(
+        np.asarray(self.base.captured["decoder_segment_ids"]),
+        np.asarray(explicit_seg),
+    )
 
   @pytest.mark.cpu_only
   def test_returns_logits_and_none_tuple(self):
@@ -135,6 +146,28 @@ def test_returns_logits_and_none_tuple(self):
     self.assertIsNone(second)
     self.assertEqual(logits.shape, (1, 5, 2))
 
+  @pytest.mark.cpu_only
+  def test_passes_forced_routed_experts(self):
+    """Caller passes forced_routed_experts -> adapter passes it directly to base."""
+    adapter = TunixMaxTextAdapter(base_model=self.base, pad_id=99)
+
+    input_tokens = jnp.array([[10, 11, 12, 99, 99]], dtype=jnp.int32)
+    positions = jnp.arange(5, dtype=jnp.int32)[None, :]
+    forced_experts = jnp.array([[[0, 1], [2, 3]]])
+
+    adapter(
+        input_tokens,
+        positions,
+        None,
+        None,
+        decoder_segment_ids=None,
+        forced_routed_experts=forced_experts,
+    )
+
+    captured_experts = self.base.captured["forced_routed_experts"]
+    self.assertIsNotNone(captured_experts)
+    np.testing.assert_array_equal(np.asarray(captured_experts), np.asarray(forced_experts))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/forced_routing_test.py b/tests/unit/forced_routing_test.py
new file mode 100644
index 0000000000..dd1d14abac
--- /dev/null
+++ b/tests/unit/forced_routing_test.py
@@ -0,0 +1,129 @@
+# Copyright 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 forced routing mismatch calculation."""
+
+# pylint: disable=protected-access,missing-function-docstring
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from maxtext.configs import pyconfig
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+from absl.testing import absltest
+
+
+class ForcedRoutingTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.cfg = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="forced_routing_test",
+        enable_checkpointing=False,
+        model_name="qwen3-30b-a3b",
+        dtype="bfloat16",
+        megablox=False,
+        sparse_matmul=False,
+        max_target_length=80,
+        per_device_batch_size=1,
+        num_experts=8,
+        num_experts_per_tok=2,
+        override_model_config=True,
+    )
+    self.rngs = nnx.Rngs(params=0)
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = Mesh(devices_array, self.cfg.mesh_axes)
+    self.model = moe.RoutedMoE(
+        config=self.cfg,
+        num_experts=self.cfg.num_experts,
+        num_experts_per_tok=self.cfg.num_experts_per_tok,
+        mesh=self.mesh,
+        kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+        kernel_axes=("embed", "mlp"),
+        dtype=self.cfg.dtype,
+        rngs=self.rngs,
+    )
+
+  def test_calculate_mismatch_perfect_match(self):
+    # Shape: [batch, length, num_experts_per_tok] -> [1, 2, 2]
+    model_indices = jnp.array([[[0, 1], [2, 3]]])
+    forced_indices = jnp.array([[[0, 1], [2, 3]]])
+    mismatch = self.model._calculate_mismatch(model_indices, forced_indices)
+    self.assertAlmostEqual(mismatch.item(), 0.0, places=6)
+
+  def test_calculate_mismatch_complete_mismatch(self):
+    model_indices = jnp.array([[[0, 1], [2, 3]]])
+    forced_indices = jnp.array([[[4, 5], [6, 7]]])
+    mismatch = self.model._calculate_mismatch(model_indices, forced_indices)
+    self.assertAlmostEqual(mismatch.item(), 1.0, places=6)
+
+  def test_calculate_mismatch_partial_match(self):
+    model_indices = jnp.array([[[0, 1], [2, 3]]])
+    forced_indices = jnp.array([[[0, 2], [2, 4]]])
+    # Token 0: model={0,1}, forced={0,2}. Intersection={0} (size 1). Mismatch = 1 - 1/2 = 0.5
+    # Token 1: model={2,3}, forced={2,4}. Intersection={2} (size 1). Mismatch = 1 - 1/2 = 0.5
+    # Average = 0.5
+    mismatch = self.model._calculate_mismatch(model_indices, forced_indices)
+    self.assertAlmostEqual(mismatch.item(), 0.5, places=6)
+
+  def test_calculate_mismatch_with_padding(self):
+    model_indices = jnp.array([[[0, 1], [2, 3]]])
+    # Token 1 is padded (starts with -1)
+    forced_indices = jnp.array([[[0, 2], [-1, -1]]])
+    # Token 0: model={0,1}, forced={0,2}. Mismatch = 0.5
+    # Token 1: ignored
+    # Average = 0.5
+    mismatch = self.model._calculate_mismatch(model_indices, forced_indices)
+    self.assertAlmostEqual(mismatch.item(), 0.5, places=6)
+
+  def test_calculate_mismatch_all_padded(self):
+    model_indices = jnp.array([[[0, 1], [2, 3]]])
+    forced_indices = jnp.array([[[-1, -1], [-1, -1]]])
+    mismatch = self.model._calculate_mismatch(model_indices, forced_indices)
+    # Should handle all padded case without division by zero error, returning 0.0 due to 1e-8 eps
+    self.assertAlmostEqual(mismatch.item(), 0.0, places=6)
+
+  def test_routed_moe_calculates_and_sows_mismatch(self):
+    batch_size = 1
+    seq_len = 2
+    emb_dim = self.cfg.emb_dim
+    inputs = jax.random.normal(jax.random.key(1), (batch_size, seq_len, emb_dim))
+    forced_experts = jnp.array([[[0, 1], [2, 3]]])
+
+    _ = self.model(inputs, forced_routed_experts=forced_experts)
+
+    self.assertTrue(hasattr(self.model, "mismatch_rate"))
+    mismatch_rate = self.model.mismatch_rate
+    self.assertIsInstance(mismatch_rate, moe.RoutingMismatchRate)
+    self.assertIsInstance(mismatch_rate[...], jax.Array)
+    self.assertTrue(0.0 <= mismatch_rate[...].item() <= 1.0)
+
+    state = nnx.state(self.model, nnx.Intermediate)
+    flat_state = list(nnx.to_flat_state(state))
+    has_mismatch_rate = False
+    for path, leaf in flat_state:
+      if path[-1] == "mismatch_rate":
+        has_mismatch_rate = True
+        self.assertIsInstance(leaf, moe.RoutingMismatchRate)
+        break
+    self.assertTrue(has_mismatch_rate)
+
+
+if __name__ == "__main__":
+  absltest.main()