fix "enable_ep_moe" bug

awex/config.py

@@ -33,7 +33,7 @@ class InferenceConfig:
     dp_size: Optional[int] = None
     load_balance_method: Optional[str] = None
     # Expert parallelism
-    ep_size: Optional[int] = None
+    ep_size: Optional[int] = 1
     enable_dp_attention: Optional[bool] = None
     enable_dp_lm_head: Optional[bool] = None
     enable_ep_moe: Optional[bool] = None

awex/meta/infer_meta_resolver.py

@@ -147,6 +147,7 @@ def _get_model_param_info(
         """
         model = kwargs["model"]
         model_context = kwargs["model_context"]
+        model_context["infer_engine_config"] = infer_engine_config
         params_meta = []
         rank_info = get_rank_info_extractor(engine_name)(model_context, engine_rank)
         model_arch_name = type(model).__name__

awex/sharding/sglang_sharding.py

@@ -42,7 +42,7 @@ def get_sglang_sharding_strategy(
 
 def get_sglang_rank_info(model_context, engine_rank) -> RankInfo:
     scheduler = model_context["scheduler"]
-    infer_engine_config = scheduler.server_args
+    infer_engine_config = model_context["infer_engine_config"]
     if infer_engine_config.dp_size != 1 and not infer_engine_config.enable_dp_attention:
         raise ValueError(
             f"DP size is not 1, but {infer_engine_config.dp_size}. This is not supported yet."