fixes

Signed-off-by: Rui Qiao <ruisearch42@gmail.com>

Author: ruisearch42

vllm/entrypoints/openai/api_server.py

@@ -1080,16 +1080,16 @@ async def scale(raw_request: Request):
         await client.scale(new_data_parallel_size, drain_timeout)
         return JSONResponse({
             "message":
-            f"Scaled up to {new_data_parallel_size} "
+            f"Scaled to {new_data_parallel_size} "
             "data parallel engines",
         })
     except TimeoutError as e:
-        raise HTTPException(
-            status_code=408,
-            detail="Scale up failed due to request drain timeout "
-            f"after {drain_timeout} seconds") from e
+        raise HTTPException(status_code=408,
+                            detail="Scale failed due to request drain timeout "
+                            f"after {drain_timeout} seconds") from e
     except Exception as e:
-        raise HTTPException(status_code=500, detail="Scale up failed") from e
+        logger.error("Scale failed: %s", e)
+        raise HTTPException(status_code=500, detail="Scale failed") from e
     finally:
         raw_request.app.state.scaling = False
 

vllm/model_executor/layers/fused_moe/layer.py

@@ -223,7 +223,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
     def __init__(self, moe: FusedMoEConfig):
         super().__init__()
         self.fused_experts = fused_experts  # type: ignore
-        self.topk_indices_dtype = None
+        self.topk_indices_dtype = torch.uint32
         self.moe = moe
 
         self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

@@ -78,7 +78,9 @@ def max_num_tokens_per_rank(self) -> Optional[int]:
         return self.max_num_tokens
 
     def topk_indices_dtype(self) -> Optional[torch.dtype]:
-        return torch.int32
+        # FIXME(rui): this needs to be int32,
+        # see https://github.com/vllm-project/vllm/pull/20166
+        return torch.uint32
 
     def num_dispatchers(self) -> int:
         return self.num_dispatchers_
@@ -100,9 +102,10 @@ def prepare(
         hidden_dim = a1.size(-1)  # K
 
         assert topk_ids.size(0) == num_tokens
-        assert expert_map is None, """with expert map, -1 id is used for
-            non-local token; this causes error when casting ids to the
-            topk_indices_dtype() uint32"""
+        # FIXME(rui)
+        # assert expert_map is None, """with expert map, -1 id is used for
+        #     non-local token; this causes error when casting ids to the
+        #     topk_indices_dtype() uint32"""
 
         # Is this always going to be a1.device?
         device = a1.device

vllm/v1/engine/async_llm.py

@@ -638,7 +638,7 @@ async def scale(self,
                 Maximum time to wait for requests to drain (seconds)
         """
         from vllm.v1.engine.core_client import RayDPClient
-        
+
         if not isinstance(self.engine_core, RayDPClient):
             raise NotImplementedError(
                 "Scale up/down only supported by RayDPClient")

vllm/v1/engine/coordinator.py

@@ -223,14 +223,12 @@ def process_input_socket(self, front_publish_address: str,
                             self.engines_running = False
                             logger.info(
                                 "DPCoordinator scaled up from %s to %s "
-                                "engines",
-                                current_count, new_engine_count)
+                                "engines", current_count, new_engine_count)
                         else:
                             self.engines = self.engines[:new_engine_count]
                             logger.info(
                                 "DPCoordinator scaled down from %s to %s "
-                                "engines",
-                                current_count, new_engine_count)
+                                "engines", current_count, new_engine_count)
                         continue  # Skip normal engine notification processing
 
                     # We received a message on the front-end XPUB socket,

vllm/v1/engine/core.py

@@ -146,8 +146,7 @@ def _initialize_kv_caches(
         if os.environ.get("VLLM_EEP_RECONFIGURE_LAUNCH") == "1":
             dp_group = getattr(self, "dp_group", None)
             assert dp_group is not None
-            kv_cache_memory = ParallelConfig.sync_kv_cache_memory(
-                dp_group, -1)
+            kv_cache_memory = ParallelConfig.sync_kv_cache_memory(dp_group, -1)
             available_gpu_memory = [kv_cache_memory] * len(kv_cache_specs)
         else:
             # Profiles the peak memory usage of the model to determine how much

vllm/v1/engine/core_client.py

@@ -28,7 +28,7 @@
 from vllm.v1.engine.coordinator import DPCoordinator
 from vllm.v1.engine.core import EngineCore, EngineCoreProc
 from vllm.v1.engine.exceptions import EngineDeadError
-from vllm.v1.engine.utils import (CoreEngine, CoreEngineActorManager,
+from vllm.v1.engine.utils import (CoreEngineActorManager,
                                   CoreEngineProcManager, EngineZmqAddresses,
                                   launch_core_engines)
 from vllm.v1.executor.abstract import Executor
@@ -94,6 +94,8 @@ def make_async_mp_client(
                 # External load balancer - client per DP rank.
                 return DPAsyncMPClient(*client_args)
             # Internal load balancer - client balances to all DP ranks.
+            if parallel_config.data_parallel_backend == "ray":
+                return RayDPClient(*client_args)
             return DPLBAsyncMPClient(*client_args)
         return AsyncMPClient(*client_args)
 
@@ -1115,7 +1117,7 @@ def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,
 
     async def _send_reconfig_message(
             self, reconfig_request: ReconfigureDistributedRequest,
-            engine: CoreEngine) -> asyncio.Future:
+            engine: EngineIdentity) -> asyncio.Future:
         """Send reconfiguration message and return the result future without
         waiting for completion."""
         call_id = uuid.uuid1().int >> 64
@@ -1160,17 +1162,17 @@ async def scale_up(self, new_data_parallel_size: int) -> None:
         # Phase 2: Create new engines now that reconfig messages have been sent
         # self.resources.engine_manager is guaranteed to be
         # CoreEngineActorManager for RayDPClient
-        assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
+        assert isinstance(self.resources.engine_manager,
+                          CoreEngineActorManager)
         self.resources.engine_manager.scale_up(self.vllm_config,
                                                new_data_parallel_size)
 
         # Create new CoreEngine objects for the new engines
         new_engine_identities = set()
         for i in range(current_dp_size, new_data_parallel_size):
-            # TODO(yongji): check if the engine is local
-            new_engine = CoreEngine(index=i, local=False)
+            new_engine = i.to_bytes(2, "little")
             self.core_engines.append(new_engine)
-            new_engine_identities.add(new_engine.identity)
+            new_engine_identities.add(new_engine)
 
         # Wait for ready messages from new engines on the input socket
         sync_input_socket = zmq.Socket.shadow(self.input_socket)
@@ -1233,7 +1235,8 @@ async def scale_down(self, new_data_parallel_size: int) -> None:
 
         await asyncio.gather(*reconfig_futures)
 
-        assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
+        assert isinstance(self.resources.engine_manager,
+                          CoreEngineActorManager)
         self.resources.engine_manager.scale_down(current_dp_size,
                                                  new_data_parallel_size)
 

vllm/v1/engine/utils.py

@@ -221,7 +221,7 @@ def __init__(
             dp_vllm_config = copy.deepcopy(vllm_config)
             pg = placement_groups[index]
             dp_vllm_config.parallel_config.placement_group = pg
-            on_head_node = index < local_engine_count
+            local_client = index < local_engine_count
             actor = ray.remote(DPEngineCoreActor).options(
                 scheduling_strategy=PlacementGroupSchedulingStrategy(
                     placement_group=pg,
@@ -230,15 +230,15 @@ def __init__(
                 runtime_env=runtime_env).remote(vllm_config=dp_vllm_config,
                                                 executor_class=executor_class,
                                                 log_stats=log_stats,
-                                                on_head_node=on_head_node,
+                                                local_client=local_client,
                                                 addresses=addresses,
                                                 dp_rank=index,
                                                 local_dp_rank=local_index)
-            if on_head_node:
+            if local_client:
                 self.local_engine_actors.append(actor)
             else:
                 self.remote_engine_actors.append(actor)
-            self.placement_group_is_local.append(on_head_node)
+            self.placement_group_is_local.append(local_client)
             refs.append(actor.wait_for_init.remote())
 
         ray.get(refs)
@@ -435,11 +435,11 @@ def scale_up(self, old_vllm_config: VllmConfig,
             dp_vllm_config.parallel_config.placement_group = pg
 
             # Check if this placement group is on the head node
-            on_head_node = any(
+            local_client = any(
                 bundle.get("node:" + dp_master_ip, 0) > 0
                 for bundle in pg.bundle_specs)
 
-            if on_head_node:
+            if local_client:
                 new_local_engines += 1
                 # Update data_parallel_size_local
                 dp_vllm_config.parallel_config.data_parallel_size_local = (
@@ -455,17 +455,17 @@ def scale_up(self, old_vllm_config: VllmConfig,
                     vllm_config=dp_vllm_config,
                     executor_class=self.executor_class,
                     log_stats=self.log_stats,
-                    on_head_node=on_head_node,
+                    local_client=local_client,
                     addresses=self.addresses,
                     dp_rank=rank,
                     local_dp_rank=local_rank)
 
-            if on_head_node:
+            if local_client:
                 self.local_engine_actors.append(actor)
             else:
                 self.remote_engine_actors.append(actor)
             self.created_placement_groups.append(pg)
-            self.placement_group_is_local.append(on_head_node)
+            self.placement_group_is_local.append(local_client)
 
         ray.get([
             actor.wait_for_init.remote()

vllm/v1/worker/cpu_model_runner.py

@@ -50,7 +50,7 @@ def replace_tensor(obj: Any, cpu_attr_name: str,
             if k.endswith("_cpu") and isinstance(v, torch.Tensor):
                 replace_tensor(self.input_batch.block_table, k, k[:-4])
 
-    def load_model(self) -> None:
+    def load_model(self, reconfigure: bool = False) -> None:
         logger.info("Starting to load model %s...", self.model_config.model)
         self.model = get_model(vllm_config=self.vllm_config)
 

vllm/v1/worker/gpu_worker.py

@@ -372,6 +372,7 @@ def reinitialize_distributed(
                 old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
                 for old_ep_rank in range(old_ep_size)
             }
+            assert self.model_runner.eplb_state is not None
             self.model_runner.eplb_state.rearrange(self.model_runner.model,
                                                    execute_shuffle=True,
                                                    global_expert_load=None,
@@ -427,6 +428,7 @@ def reinitialize_distributed(
             module.moe_config.moe_parallel_config = module.moe_parallel_config
         if new_ep_size < old_ep_size:
             num_local_physical_experts = num_local_experts
+            assert self.model_runner.eplb_state is not None
             new_physical_experts = \
                 self.model_runner.eplb_state.physical_to_logical_map.shape[1]
             parallel_config.num_redundant_experts = (
@@ -441,6 +443,7 @@ def reinitialize_distributed(
                                         group_src=0)
             num_local_physical_experts = num_local_physical_experts.item()
             new_physical_experts = num_local_physical_experts * new_ep_size
+            assert self.model_runner.eplb_state is not None
             global_expert_load = self.model_runner.eplb_state.rearrange(
                 self.model_runner.model, execute_shuffle=False)
             parallel_config.num_redundant_experts = (
@@ -457,14 +460,14 @@ def reinitialize_distributed(
                 old_ep_rank: old_ep_rank
                 for old_ep_rank in range(old_ep_size)
             }
+            assert self.model_runner.eplb_state is not None
             self.model_runner.eplb_state.rearrange(
                 self.model_runner.model,
                 execute_shuffle=True,
                 global_expert_load=global_expert_load,
                 rank_mapping=rank_mapping)
             if get_ep_group().rank == 0:
                 logger.info("[Elastic EP] Expert resharding completed!")
-        self.model_runner.eplb_state.expert_rearrangement_step = 0
 
     def save_sharded_state(
         self,