Fix EP vllm keyerror

Author: dipannita08

src/maxtext/eval/runner/common.py

@@ -23,13 +23,6 @@
 if TYPE_CHECKING:
   from maxtext.eval.runner.server_manager import VllmServerManager
 
-ENABLE_EXPERT_PARALLEL_HELP = (
-    "Enable expert parallelism in vLLM. Required for MoE models such as "
-    "qwen3-30b-a3b, qwen3-235b-a22b, deepseek-v3, etc. Without this flag "
-    "tpu-inference omits the 'expert' mesh axis and MaxText's MoE sharding "
-    "raises KeyError."
-)
-
 
 def resolve_token(cfg: dict, hf_token: str | None) -> str | None:
   """Return HF token from explicit arg or HF_TOKEN env var."""
@@ -71,10 +64,9 @@ def build_server_manager(cfg: dict, token: str | None) -> "VllmServerManager":
   if max_num_seqs is not None:
     max_num_seqs = int(max_num_seqs)
 
+  expert_parallel_size = int(cfg.get("expert_parallel_size") or 1)
+
   server_env = {"HF_TOKEN": token} if token else None
-  additional_vllm_kwargs: dict = {}
-  if cfg.get("enable_expert_parallel"):
-    additional_vllm_kwargs["enable_expert_parallel"] = True
 
   return VllmServerManager(
       model_path=hf_path,
@@ -83,11 +75,11 @@ def build_server_manager(cfg: dict, token: str | None) -> "VllmServerManager":
       host=server_host,
       port=server_port,
       tensor_parallel_size=tensor_parallel_size,
+      expert_parallel_size=expert_parallel_size,
       max_model_len=max_model_len,
       max_num_batched_tokens=max_num_batched_tokens,
       max_num_seqs=max_num_seqs,
       env=server_env,
-      additional_vllm_kwargs=additional_vllm_kwargs or None,
   )
 
 
@@ -121,7 +113,12 @@ def add_server_args(parser: argparse.ArgumentParser) -> None:
   parser.add_argument("--max_num_seqs", type=int, help="vLLM max concurrent sequences.")
   parser.add_argument("--hf_mode", action="store_true", help="HF safetensors mode.")
   parser.add_argument(
-      "--enable_expert_parallel", action="store_true", help=ENABLE_EXPERT_PARALLEL_HELP
+      "--expert_parallel_size",
+      type=int,
+      default=0,
+      help=(
+          "Chips allocated to the expert mesh axis (EP). "
+      ),
   )
   parser.add_argument("--hf_token", help="HuggingFace token for gated models.")
   parser.add_argument(

src/maxtext/eval/runner/server_manager.py

@@ -188,7 +188,8 @@ class VllmServerManager:
     maxtext_model_name: MaxText model name (e.g. "llama3.1-8b").
     host: Hostname the HTTP server binds to (rank-0 only).
     port: Port the HTTP server listens on.
-    tensor_parallel_size: Tensor parallelism.
+    tensor_parallel_size: Total number of chips.
+    expert_parallel_size: Chips allocated to the expert mesh axis (EP).
     max_model_len: Maximum sequence length.
     dtype: Activation dtype string passed to vLLM (e.g. "bfloat16").
     max_num_batched_tokens: Tokens per scheduler step (None = vLLM default).
@@ -206,6 +207,7 @@ def __init__(
       host: str = "localhost",
       port: int = 8000,
       tensor_parallel_size: int = 4,
+      expert_parallel_size: int = 1,
       max_model_len: int = 4096,
       dtype: str = "bfloat16",
       max_num_batched_tokens: int | None = None,
@@ -216,12 +218,18 @@ def __init__(
   ):
     if checkpoint_path and not maxtext_model_name:
       raise ValueError("maxtext_model_name is required when checkpoint_path is set.")
+    if tensor_parallel_size % expert_parallel_size != 0:
+      raise ValueError(
+          f"tensor_parallel_size ({tensor_parallel_size}) is not divisible by "
+          f"expert_parallel_size ({expert_parallel_size})."
+      )
     self.model_path = model_path
     self.checkpoint_path = checkpoint_path
     self.maxtext_model_name = maxtext_model_name
     self.host = host
     self.port = port
     self.tensor_parallel_size = tensor_parallel_size
+    self.expert_parallel_size = expert_parallel_size
     self.max_model_len = max_model_len
     self.dtype = dtype
     self.max_num_batched_tokens = max_num_batched_tokens
@@ -251,9 +259,13 @@ def start(self) -> None:
     if self.env:
       os.environ.update(self.env)
 
+    # total chips = ici_tensor_parallelism x ici_expert_parallelism.
+    ici_tp = self.tensor_parallel_size // self.expert_parallel_size
+    ici_ep = self.expert_parallel_size
+
     vllm_kwargs: dict = {
         "model": self.model_path,
-        "tensor_parallel_size": self.tensor_parallel_size,
+        "tensor_parallel_size": ici_tp,
         "max_model_len": self.max_model_len,
         "dtype": self.dtype,
     }
@@ -269,14 +281,15 @@ def start(self) -> None:
               "model_name": self.maxtext_model_name,
               "load_parameters_path": self.checkpoint_path,
               "log_config": False,
-          }
+              "ici_tensor_parallelism": ici_tp,
+              "ici_expert_parallelism": ici_ep,
+          },
+          "sharding": {
+              "sharding_strategy": {},
+          },
       }
-      if self.additional_vllm_kwargs.get("enable_expert_parallel"):
-        vllm_kwargs["additional_config"]["sharding"] = {
-            "sharding_strategy": {
-                "expert_parallelism": self.tensor_parallel_size,
-            }
-        }
+      if ici_ep > 1:
+        vllm_kwargs["additional_config"]["sharding"]["sharding_strategy"]["expert_parallelism"] = ici_ep
     else:
       vllm_kwargs["load_format"] = "auto"
 
@@ -298,8 +311,9 @@ def start(self) -> None:
           vllm_kwargs[_k] = _v
 
     logger.info(
-        "Initializing in-process vLLM (tp=%d, max_len=%d)...",
-        self.tensor_parallel_size,
+        "Initializing in-process vLLM (tp=%d, ep=%d, max_len=%d)...",
+        ici_tp,
+        ici_ep,
         self.max_model_len,
     )
     self._llm = LLM(**vllm_kwargs)

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

@@ -20,11 +20,12 @@
 from flax import nnx
 import flax.linen as nn
 from jax import numpy as jnp
-from jax.sharding import Mesh
+from jax.sharding import AxisType, Mesh
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE
 from maxtext.utils import max_logging
+from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
 
 
@@ -98,8 +99,9 @@ def __init__(self, vllm_config: VllmConfig, rng_key: jax.Array, mesh: Mesh):
     self.cfg = vllm_config.model_config
     self.maxtext_config = generate_maxtext_config(vllm_config)
 
-    # Model configuration
-    self.mesh = mesh
+    devices_array = maxtext_utils.create_device_mesh(self.maxtext_config)
+    axis_types = tuple([AxisType.Auto] * len(self.maxtext_config.mesh_axes))
+    self.mesh = Mesh(devices_array, self.maxtext_config.mesh_axes, axis_types=axis_types)
     self.model_mode = MODEL_MODE_AUTOREGRESSIVE
     self.is_text_generation_model = True