[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'.

Author: xuefgu

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
 

src/maxtext/layers/decoders.py

@@ -794,6 +794,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
@@ -1013,19 +1014,34 @@ def __call__(
           current_broadcast_args = list(broadcast_args)
           current_in_axes_tuple = list(in_axes_tuple)
 
+          if forced_routed_experts is not None:
+            # Transpose [B, L, N, E] -> [N, B, L, E] for scan
+            forced_routed_experts_t = jnp.transpose(forced_routed_experts, (2, 0, 1, 3))
+            cycle_interval = cfg.inhomogeneous_layer_cycle_interval
+            reshaped_forced = jnp.reshape(
+                forced_routed_experts_t, (scan_length, cycle_interval) + forced_routed_experts_t.shape[1:]
+            )
+          else:
+            reshaped_forced = None
+
           if kv_caches is not None:
             # Stack kv_caches for scan: [num_layers, ...]
             stacked_kv_cache = jnp.stack(kv_caches, axis=0)
 
             # We pass (y, stacked_kv_cache, 0) as the carry
             carry = (y, stacked_kv_cache, 0)
 
-            # We don't pass kv_cache as a scanned argument anymore
-
             # Pass None for previous_chunk, slot, kv_cache to align with __call__ signature
             current_broadcast_args.extend([None, None, None, attention_metadata])
             current_in_axes_tuple.extend([nn.broadcast] * 4)
 
+            if reshaped_forced is not None:
+              current_broadcast_args.append(reshaped_forced)
+              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]}")
 
@@ -1047,8 +1063,15 @@ 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)
+            current_broadcast_args.extend([None, None, None, None])
+            current_in_axes_tuple.extend([nn.broadcast] * 4)
+
+            if reshaped_forced is not None:
+              current_broadcast_args.append(reshaped_forced)
+              current_in_axes_tuple.append(0)
+            else:
+              current_broadcast_args.append(None)
+              current_in_axes_tuple.append(nn.broadcast)
 
             y, _ = self.scan_decoder_layers(
                 cfg,
@@ -1077,6 +1100,10 @@ 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:
+                extra_kwargs["forced_routed_experts"] = forced_routed_experts[:, :, index, :]
+
               y, kv_cache = layer(
                   config=cfg,
                   mesh=mesh,
@@ -1095,6 +1122,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
@@ -1114,6 +1142,7 @@ def __call__(
               slot=slot,
           )
         else:
+          moe_lyr_idx = 0
           for lyr in range(cfg.num_decoder_layers):
             RemattedBlockLayer = RemattedBlockLayers[0]
             layer_kwargs = {}
@@ -1150,6 +1179,26 @@ def __call__(
             layer = RemattedBlockLayer(
                 config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs
             )
+
+            is_moe = False
+            if cfg.decoder_block in (
+                DecoderBlockType.QWEN3_MOE,
+                DecoderBlockType.QWEN3_NEXT,
+                DecoderBlockType.QWEN3_5,
+            ):
+              is_moe = True
+
+            current_forced_routed_experts = None
+            if is_moe and forced_routed_experts is not None:
+              current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :]
+              moe_lyr_idx += 1
+            elif is_moe:
+              moe_lyr_idx += 1
+
+            extra_kwargs = {}
+            if is_moe 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,
@@ -1160,6 +1209,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: