Support extracting router replay expert indices from MaxText inference

src/maxtext/configs/inference/vllm.yml

@@ -33,7 +33,7 @@ vllm_additional_config: {}
 
 
 # -------------- Logical Axis Rules --------------
-mesh_axes: ['data', 'attn_dp', 'model', 'expert', 'attn_dp_expert']
+mesh_axes: ['data', 'attn_dp', 'model', 'expert', 'attn_dp_expert', 'dcp', 'pcp']
 logical_axis_rules: [
                       # ==========================================
                       # Vocabulary Embedding

src/maxtext/configs/types.py

@@ -3917,6 +3917,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         "autoregressive": self.ici_autoregressive_parallelism,
         "attn_dp": (1),  # initialized to 1, vLLM will auto calculate this value based on TP and num_kv_heads
         "attn_dp_expert": (1),  # initialized to 1, vLLM will auto calculate this value based on EP
+        "dcp": (1),
+        "pcp": (1),
     }
     self.ici_parallelism = [ici_map[axis] for axis in self.mesh_axes]
 
@@ -3936,6 +3938,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         "autoregressive": self.dcn_autoregressive_parallelism,
         "attn_dp": (1),  # initialized to 1, vLLM will auto calculate this value based on TP and num_kv_heads
         "attn_dp_expert": (1),  # initialized to 1, vLLM will auto calculate this value based on EP
+        "dcp": (1),
+        "pcp": (1),
     }
     self.dcn_parallelism = [dcn_map[axis] for axis in self.mesh_axes]
 

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -263,7 +263,7 @@ def __call__(
     with self.mesh, nn.logical_axis_rules(self.maxtext_config.logical_axis_rules):
       aux_hidden_states = []
       expert_indices = None
-      hidden, kv_caches = self.model(
+      res = self.model(
           decoder_input_tokens=input_ids,
           decoder_positions=input_positions,
           kv_caches=kv_caches,
@@ -272,6 +272,13 @@ def __call__(
           **kwargs,
       )
 
+      if isinstance(res, tuple) and len(res) == 3:
+        hidden, kv_caches, expert_indices = res
+      elif isinstance(res, tuple) and len(res) == 4:
+        _, hidden, kv_caches, expert_indices = res
+      else:
+        hidden, kv_caches = res
+
       # To be compatible with vLLM, we reshape to (batch * seq, dim).
       hidden = hidden.reshape((-1, hidden.shape[-1]))
 

src/maxtext/layers/moe.py

@@ -2990,6 +2990,9 @@ def fused_moe_matmul(
     hidden_states = jnp.reshape(inputs, (batch_size * seq_len, emb_dim))
     gating_output = jnp.reshape(gate_logits, (batch_size * seq_len, self.num_experts))
 
+    _, top_k_indices = jax.lax.top_k(gating_output, self.num_experts_per_tok)
+    self.selected_experts = nnx.Intermediate(top_k_indices)
+
     # Concatenate gate and up projections: [E, D, H] + [E, D, H] -> [E, D, 2H]
     # fused_moe_func splits this internally: gate=w1[..., :H], up=w1[..., H:]
     if fused_kernel is None:

src/maxtext/layers/nnx_decoders.py

@@ -2011,7 +2011,20 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
     else:
       logits = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode)
 
-    return logits, hidden_state, kv_caches
+    expert_indices = None
+    try:
+      expert_indices_list = []
+      intermediates = nnx.state(self, nnx.Intermediate)
+      for path, val in intermediates.flat_state():
+        if path and str(path[-1]) == "selected_experts":
+          v = val.value if hasattr(val, "value") else val
+          expert_indices_list.append(v)
+      if expert_indices_list:
+        expert_indices = jnp.stack(expert_indices_list, axis=0)
+    except Exception:
+      expert_indices = None
+
+    return logits, hidden_state, kv_caches, expert_indices
 
   def _apply_deepseek4_scanned_blocks(
       self,

src/maxtext/models/models.py

@@ -549,7 +549,7 @@ def __call__(
       mutable_collections.append("intermediates")
 
     if self.config.pure_nnx_decoder:
-      logits, hidden_state, kv_caches = self.decoder(
+      res = self.decoder(
           shared_embedding=self.token_embedder,
           decoder_input_tokens=decoder_input_tokens,
           decoder_positions=decoder_positions,
@@ -563,8 +563,13 @@ def __call__(
           attention_metadata=attention_metadata,
           deepstack_visual_embeds=deepstack_visual_embeds,
       )  # pytype: disable=wrong-keyword-args
+      if isinstance(res, tuple) and len(res) == 4:
+        logits, hidden_state, kv_caches, expert_indices = res
+      else:
+        logits, hidden_state, kv_caches = res
+        expert_indices = None
     else:
-      logits, hidden_state, kv_caches = self.decoder(
+      res = self.decoder(
           shared_embedding=self.token_embedder,
           decoder_input_tokens=decoder_input_tokens,
           decoder_positions=decoder_positions,
@@ -579,6 +584,11 @@ def __call__(
           deepstack_visual_embeds=deepstack_visual_embeds,
           mutable=mutable_collections,  # pyrefly: ignore[unexpected-keyword]
       )  # pytype: disable=wrong-keyword-args
+      if isinstance(res, tuple) and len(res) == 4:
+        logits, hidden_state, kv_caches, expert_indices = res
+      else:
+        logits, hidden_state, kv_caches = res
+        expert_indices = None
 
     # If we are initializing the model AND MTP is enabled, we must create
     # dummy target tensors. This allows Flax to trace the MTPBlock and create
@@ -614,6 +624,8 @@ def __call__(
 
     if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
       # In vLLM, logits are computed separately after updating the KV cache.
+      if expert_indices is not None:
+        return hidden_state, kv_caches, expert_indices
       return hidden_state, kv_caches
 
     return logits

tests/test_inference_router_replay.py

@@ -0,0 +1,123 @@
+"""Integration Test: Extract Router Replay Data from MaxText MoE Inference.
+
+Validates:
+1. MaxText MoE inference execution (attention="vllm_rpa", fused_moe_matmul) capturing `selected_experts` / `expert_indices`.
+2. Router replay extraction from intermediate state and decoder outputs with shape (batch_size, seq_len, top_k).
+"""
+
+import os
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+os.environ["NEW_MODEL_DESIGN"] = "1"
+os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+import sys
+import unittest
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+
+from maxtext.configs import pyconfig
+from maxtext.models import models
+from maxtext.utils import maxtext_utils
+from maxtext.common.common_types import (
+    DECODING_ACTIVE_SEQUENCE_INDICATOR,
+    MODEL_MODE_PREFILL,
+)
+from tests.utils.test_helpers import get_test_config_path
+
+
+class InferenceRouterReplayExtractionTest(unittest.TestCase):
+
+  def setUp(self):
+    os.environ["NEW_MODEL_DESIGN"] = "1"
+    os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+
+  def test_extract_router_replay_from_inference(self):
+    seq_len = 16
+    batch_size = 2
+    num_layers = 1
+    num_experts = 4
+    top_k = 2
+
+    test_tokens = [791, 7155, 315, 9342, 374, 9897, 323, 374] * 4
+    raw_tokens = test_tokens[:seq_len]
+
+    base_kwargs = {
+        "run_name": "test_router_replay_extraction_fast",
+        "enable_checkpointing": False,
+        "override_model_config": True,
+        "base_num_decoder_layers": num_layers,
+        "num_decoder_layers": num_layers,
+        "model_name": "qwen3.5-35b-a3b",
+        "num_experts": num_experts,
+        "num_experts_per_tok": top_k,
+        "base_emb_dim": 512,
+        "base_num_query_heads": 2,
+        "base_num_kv_heads": 2,
+        "head_dim": 256,
+        "partial_rotary_factor": 0.25,
+        "base_mlp_dim": 1024,
+        "base_moe_mlp_dim": 1024,
+        "vocab_size": 1000,
+        "max_target_length": seq_len,
+        "max_prefill_predict_length": seq_len,
+        "per_device_batch_size": float(batch_size),
+        "scan_layers": False,
+        "weight_dtype": "bfloat16",
+        "dtype": "bfloat16",
+        "log_config": False,
+        "skip_jax_distributed_system": True,
+        "ici_tensor_parallelism": 4,
+        "ici_data_parallelism": 1,
+        "ici_expert_parallelism": 1,
+        "enable_nnx": True,
+        "pure_nnx": True,
+        "pure_nnx_decoder": True,
+    }
+
+    cfg_infer = pyconfig.initialize(
+        [sys.argv[0], get_test_config_path("inference/vllm.yml"), "attention=vllm_rpa"],
+        **base_kwargs,
+    )
+
+    devices_array = maxtext_utils.create_device_mesh(cfg_infer)
+    mesh = Mesh(devices_array, cfg_infer.mesh_axes)
+    rng = jax.random.PRNGKey(42)
+
+    ids = jnp.tile(jnp.expand_dims(jnp.array(raw_tokens, dtype=jnp.int32), axis=0), (batch_size, 1))
+    decoder_positions = jnp.tile(jnp.expand_dims(jnp.arange(seq_len, dtype=jnp.int32), axis=0), (batch_size, 1))
+    segment_ids = jnp.zeros((batch_size, seq_len), dtype=jnp.int32) + DECODING_ACTIVE_SEQUENCE_INDICATOR
+
+    model_infer = models.transformer_as_linen(config=cfg_infer, mesh=mesh, quant=None, model_mode=MODEL_MODE_PREFILL)
+    init_params_rng, init_dropout_rng = jax.random.split(rng)
+    vars_dict_init = model_infer.init(
+        {"params": init_params_rng, "dropout": init_dropout_rng},
+        ids,
+        decoder_positions,
+        segment_ids,
+        enable_dropout=False,
+    )
+    vars_infer = dict(vars_dict_init)
+
+    # Run inference prefill
+    out, cache_infer = model_infer.apply(
+        vars_infer,
+        ids,
+        decoder_positions,
+        segment_ids,
+        enable_dropout=False,
+        model_mode=MODEL_MODE_PREFILL,
+        mutable=["cache", "intermediates"],
+    )
+
+    self.assertIsInstance(out, tuple, "Inference output must be a tuple")
+    self.assertEqual(len(out), 3, "Inference output must be (hidden_state, kv_caches, expert_indices)")
+    hidden_state, kv_caches, expert_indices = out
+    self.assertIsNotNone(expert_indices, "expert_indices must not be None")
+    self.assertEqual(expert_indices.shape[1:], (batch_size * seq_len, top_k))
+    print(f"\n[Inference Router Extraction] Extracted routing data shape: {expert_indices.shape}")
+
+
+if __name__ == "__main__":
+  unittest.main()