From 823f416fa3627042b85f37782812ff93b0497a1d Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Mon, 10 Aug 2026 22:54:36 +0000 Subject: [PATCH] Support extracting router replay expert indices from MaxText inference - 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 --- src/maxtext/configs/inference/vllm.yml | 2 +- src/maxtext/configs/types.py | 4 + .../vllm/maxtext_vllm_adapter/adapter.py | 9 +- src/maxtext/layers/moe.py | 3 + src/maxtext/layers/nnx_decoders.py | 15 ++- src/maxtext/models/models.py | 16 ++- tests/test_inference_router_replay.py | 123 ++++++++++++++++++ 7 files changed, 167 insertions(+), 5 deletions(-) create mode 100644 tests/test_inference_router_replay.py diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index 82046b7ee8..c9232d6983 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 522a7cf621..004f28a034 100644 --- a/src/maxtext/configs/types.py +++ b/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] diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index bacd28e551..61713cf096 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/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])) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f6e89f204d..1daa768e7f 100644 --- a/src/maxtext/layers/moe.py +++ b/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: diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ee02407502..5dba75e721 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/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, diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index a70ad780a3..eafbd8f4fc 100644 --- a/src/maxtext/models/models.py +++ b/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 diff --git a/tests/test_inference_router_replay.py b/tests/test_inference_router_replay.py new file mode 100644 index 0000000000..07a33d5e88 --- /dev/null +++ b/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()