Support extracting router replay expert indices from MaxText inference

[khatwanimohit]

Description

• In RoutedMoE, record selected_experts as nnx.Intermediate in fused_moe_matmul
• In NNXDecoder, extract selected_experts from intermediates and return expert_indices
• In Transformer, forward expert_indices from decoder to vLLM caller
• In MaxTextForCausalLM vLLM adapter, return expert_indices in forward outputs
• Added dcp/pcp axes to ici_map, dcn_map and vllm.yml for TPU vLLM shard_map compatibility
• Added fast lightweight test_inference_router_replay.py integration test

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for extracting router replay data (expert_indices) during MaxText MoE inference. It captures selected experts as intermediates in the MoE layer, propagates them through the decoder and model outputs, and adds an integration test to validate this extraction. The review feedback highlights two critical issues in nnx_decoders.py: an alphabetical sorting bug of layers during flattening (due to unpadded names like moe_layers_10 sorting before moe_layers_2) and a shape mismatch when scan_layers is enabled. Additionally, an unreachable branch was identified in adapter.py since the model call never returns a 4-tuple.

[gemini-code-assist]

Alphabetical Sorting and Scan Layers Shape Mismatch Bug

There are two significant issues with the current implementation of expert_indices extraction:

1. Alphabetical Sorting Bug: intermediates.flat_state() returns the flat state dictionary. During flattening in JAX/Flax, dictionary keys are sorted alphabetically to ensure deterministic flattening. Since the sequential layer names are not zero-padded (e.g., moe_layers_0, moe_layers_1, ..., moe_layers_10), alphabetical sorting will place moe_layers_10 before moe_layers_2. This results in the returned expert_indices tensor having out-of-order layer indices, which will break downstream router replay analysis.
2. Scan Layers Shape Mismatch: When scan_layers is enabled, the scanned block's selected_experts already has a leading layer dimension of shape (num_layers, batch_size * seq_len, top_k). Stacking it with jnp.stack will result in an incorrect 4D shape (1, num_layers, batch_size * seq_len, top_k), whereas sequential layers (which are 2D) will be stacked into a 3D shape (num_layers, batch_size * seq_len, top_k).

Solution

We can resolve both issues by naturally sorting the paths (extracting integers from string keys) and dynamically handling both 2D (sequential) and 3D (scanned) tensors before concatenating them along the layer dimension.

    expert_indices = None
    try:
      import re
      def path_key(path_item):
        key = []
        for part in path_item[0]:
          if isinstance(part, int):
            key.append(part)
          else:
            nums = [int(x) for x in re.findall(r'\d+', str(part))]
            key.append(nums[0] if nums else str(part))
        return tuple(key)

      intermediates = nnx.state(self, nnx.Intermediate)
      flat_items = []
      for path, val in intermediates.flat_state():
        if path and str(path[-1]) == "selected_experts":
          v = val.value if hasattr(val, "value") else val
          flat_items.append((path, v))
      
      flat_items.sort(key=path_key)
      expert_indices_list = [v for _, v in flat_items]
      if expert_indices_list:
        processed_list = []
        for v in expert_indices_list:
          if v.ndim == 2:
            processed_list.append(jnp.expand_dims(v, axis=0))
          elif v.ndim == 3:
            processed_list.append(v)
        if processed_list:
          expert_indices = jnp.concatenate(processed_list, axis=0)
    except Exception:
      expert_indices = None

    return logits, hidden_state, kv_caches, expert_indices

[gemini-code-assist]

Dead Code / Unreachable Branch

Transformer.__call__ never returns a 4-tuple. When attention is in ("vllm_rpa", "vllm_batched_rpa"), it returns either a 3-tuple (hidden_state, kv_caches, expert_indices) or a 2-tuple (hidden_state, kv_caches). Otherwise, it returns a single logits array. Therefore, the len(res) == 4 check is dead code and can be safely removed to simplify the implementation.

      if isinstance(res, tuple) and len(res) == 3:
        hidden, kv_caches, expert_indices = res
      else:
        hidden, kv_caches = res