@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 217937221b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 01:00 UTC

Elastic Expert Parallel Initial Support (#20775)

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

Rui Qiao committed Jul 18, 2025 at 17:46 UTC 217937221b6845913502371aba554a3357fbccfb
24 files changed +1659 -68
examples/online_serving/elastic_ep/bench.sh new
+57
@@ -0,0 +1,57 @@
1 +#!/bin/bash
2 +
3 +MODEL_NAME="deepseek-ai/DeepSeek-V2-Lite"
4 +LOCAL_MODEL_PATH="/models/models--deepseek-ai--DeepSeek-V2-Lite/snapshots/604d5664dddd88a0433dbae533b7fe9472482de0"
5 +HOST="localhost"
6 +PORT=8006
7 +NUM_PROMPTS=20
8 +REQUEST_RATE=5
9 +
10 +# Parse command line arguments
11 +while [[ $# -gt 0 ]]; do
12 + case $1 in
13 + --model)
14 + MODEL_NAME="$2"
15 + shift 2
16 + ;;
17 + --local-model)
18 + MODEL_NAME=$LOCAL_MODEL_PATH
19 + shift
20 + ;;
21 + --host)
22 + HOST="$2"
23 + shift 2
24 + ;;
25 + --port)
26 + PORT="$2"
27 + shift 2
28 + ;;
29 + --num-prompts)
30 + NUM_PROMPTS="$2"
31 + shift 2
32 + ;;
33 + --request-rate)
34 + REQUEST_RATE="$2"
35 + shift 2
36 + ;;
37 + -h|--help)
38 + echo "Usage: $0 [OPTIONS]"
39 + echo "Options:"
40 + echo " --model MODEL_NAME Set model name or path (default: deepseek-ai/DeepSeek-V2-Lite)"
41 + echo " --local-model Use local model path (convenience option)"
42 + exit 0
43 + ;;
44 + *)
45 + echo "Unknown option: $1"
46 + echo "Use -h or --help for usage information"
47 + exit 1
48 + ;;
49 + esac
50 +done
51 +
52 +vllm bench serve \
53 + --model $MODEL_NAME \
54 + --host $HOST \
55 + --port $PORT \
56 + --num-prompts $NUM_PROMPTS \
57 + --request-rate $REQUEST_RATE
examples/online_serving/elastic_ep/scale.py new
+53
@@ -0,0 +1,53 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +import argparse
6 +import json
7 +import sys
8 +
9 +import requests
10 +
11 +
12 +def scale(host, port, new_dp_size):
13 + url = f"http://{host}:{port}/scale_elastic_ep"
14 + payload = {"new_data_parallel_size": new_dp_size}
15 + headers = {"Content-Type": "application/json"}
16 +
17 + print(f"Sending scale request to {url}")
18 + print(f"Payload: {json.dumps(payload, indent=2)}")
19 +
20 + try:
21 + response = requests.post(url, json=payload, headers=headers, timeout=300)
22 +
23 + print(f"Status Code: {response.status_code}")
24 + print(f"Response: {response.text}")
25 +
26 + if response.status_code == 200:
27 + print("Scale up/down request successful!")
28 + return True
29 + else:
30 + print("Scale up/down request failed!")
31 + return False
32 +
33 + except requests.exceptions.RequestException as e:
34 + print(f"Request failed: {e}")
35 + return False
36 +
37 +
38 +def main():
39 + parser = argparse.ArgumentParser(description="Test scale up/down functionality")
40 + parser.add_argument("--host", default="localhost", help="API server host")
41 + parser.add_argument("--port", type=int, default=8006, help="API server port")
42 + parser.add_argument(
43 + "--new-dp-size", type=int, default=2, help="New data parallel size"
44 + )
45 +
46 + args = parser.parse_args()
47 +
48 + success = scale(args.host, args.port, args.new_dp_size)
49 + sys.exit(0 if success else 1)
50 +
51 +
52 +if __name__ == "__main__":
53 + main()
examples/online_serving/elastic_ep/serve_deepseek_v2.sh new
+72
@@ -0,0 +1,72 @@
1 +#!/bin/bash
2 +
3 +HOST="0.0.0.0"
4 +PORT=8006
5 +DATA_PARALLEL_SIZE=4
6 +REDUNDANT_EXPERTS=0
7 +LOCAL_MODEL_PATH="/models/models--deepseek-ai--DeepSeek-V2-Lite/snapshots/604d5664dddd88a0433dbae533b7fe9472482de0"
8 +MODEL_NAME="deepseek-ai/DeepSeek-V2-Lite"
9 +
10 +while [[ $# -gt 0 ]]; do
11 + case $1 in
12 + --dp)
13 + DATA_PARALLEL_SIZE="$2"
14 + shift 2
15 + ;;
16 + --re)
17 + REDUNDANT_EXPERTS="$2"
18 + shift 2
19 + ;;
20 + --host)
21 + HOST="$2"
22 + shift 2
23 + ;;
24 + --port)
25 + PORT="$2"
26 + shift 2
27 + ;;
28 + --model)
29 + MODEL_NAME="$2"
30 + shift 2
31 + ;;
32 + --local-model)
33 + MODEL_NAME=$LOCAL_MODEL_PATH
34 + shift
35 + ;;
36 + -h|--help)
37 + echo "Usage: $0 [OPTIONS]"
38 + echo "Options:"
39 + echo " --dp SIZE Set data parallel size (default: 4)"
40 + echo " --re SIZE Set redundant experts (default: 0)"
41 + echo " --host HOST Set host address (default: 0.0.0.0)"
42 + echo " --port PORT Set port number (default: 8006)"
43 + echo " --model MODEL_NAME Set model name or path"
44 + echo " -h, --help Show this help message"
45 + exit 0
46 + ;;
47 + *)
48 + echo "Unknown option: $1"
49 + echo "Use -h or --help for usage information"
50 + exit 1
51 + ;;
52 + esac
53 +done
54 +
55 +echo "Starting vLLM server for $MODEL_NAME with data parallel size: $DATA_PARALLEL_SIZE and redundant experts: $REDUNDANT_EXPERTS"
56 +
57 +export RAY_DEDUP_LOGS=0
58 +export VLLM_USE_V1=1
59 +export VLLM_ALL2ALL_BACKEND="pplx"
60 +export VLLM_USE_DEEP_GEMM=1
61 +
62 +vllm serve $MODEL_NAME \
63 + --data-parallel-size $DATA_PARALLEL_SIZE \
64 + --data-parallel-size-local $DATA_PARALLEL_SIZE \
65 + --data-parallel-backend ray \
66 + --enforce-eager \
67 + --enable-expert-parallel \
68 + --enable-eplb \
69 + --num-redundant-experts $REDUNDANT_EXPERTS \
70 + --trust-remote-code \
71 + --host $HOST \
72 + --port $PORT
tools/ep_kernels/elastic_ep/eep_nvshmem.patch new
+92
@@ -0,0 +1,92 @@
1 +From 18c0599c2f07ec965132efa25961dc8179c2dda3 Mon Sep 17 00:00:00 2001
2 +From: Yongji Wu <wuyongji317@gmail.com>
3 +Date: Tue, 20 May 2025 13:41:12 -0700
4 +Subject: [PATCH] fix reinit issues due to states not cleaned up
5 +
6 +fix double free
7 +---
8 + src/host/init/init.cu | 10 ++++++++++
9 + .../internal/host/nvshmemi_mem_transport.hpp | 15 +++++++++++++++
10 + src/modules/bootstrap/uid/bootstrap_uid.cpp | 5 +++++
11 + 3 files changed, 30 insertions(+)
12 +
13 +diff --git a/src/host/init/init.cu b/src/host/init/init.cu
14 +index b1c5dbf..1fecb4b 100644
15 +--- a/src/host/init/init.cu
16 ++++ b/src/host/init/init.cu
17 +@@ -43,6 +43,8 @@
18 + #include "internal/host/nvshmemi_types.h"
19 + #include "internal/host/shared_memory.h"
20 + #include "internal/host/nvshmemi_symmetric_heap.hpp"
21 ++// eep-dev
22 ++#include "internal/host/nvshmemi_mem_transport.hpp"
23 +
24 + extern __constant__ nvshmemi_device_host_state_t nvshmemi_device_state_d;
25 + static std::map<void *, int> registered_device_states;
26 +@@ -1293,6 +1295,14 @@ void nvshmemid_hostlib_finalize(void *device_ctx, void *transport_device_ctx) {
27 + /* Multi-init Multi-fini*/
28 + nvshmemi_state = NULL;
29 + nvshmemi_device_state.nvshmemi_is_nvshmem_initialized = 0;
30 ++
31 ++ // eep-dev
32 ++ nvshmemi_mem_p2p_transport::destroy_instance();
33 ++ nvshmemi_mem_remote_transport::destroy_instance();
34 ++ free(nvshmemi_default_session);
35 ++ nvshmemi_default_session = nullptr;
36 ++ nvshmemi_device_state.nvshmemi_is_nvshmem_bootstrapped = false;
37 ++
38 + nvshmemi_is_device_state_ready = false;
39 + } else
40 + nvshmemi_boot_handle.barrier(&nvshmemi_boot_handle);
41 +diff --git a/src/include/internal/host/nvshmemi_mem_transport.hpp b/src/include/internal/host/nvshmemi_mem_transport.hpp
42 +index 2495844..e4f408a 100644
43 +--- a/src/include/internal/host/nvshmemi_mem_transport.hpp
44 ++++ b/src/include/internal/host/nvshmemi_mem_transport.hpp
45 +@@ -36,6 +36,13 @@ class nvshmemi_mem_p2p_transport final {
46 + return p2p_objref_;
47 + }
48 + }
49 ++ // eep-dev
50 ++ static void destroy_instance(void) {
51 ++ if (p2p_objref_ != nullptr) {
52 ++ delete p2p_objref_;
53 ++ p2p_objref_ = nullptr;
54 ++ }
55 ++ }
56 +
57 + void print_mem_handle(int pe_id, int transport_idx, nvshmemi_symmetric_heap &obj);
58 +
59 +@@ -87,6 +94,14 @@ class nvshmemi_mem_remote_transport final {
60 + }
61 + }
62 +
63 ++ // eep-dev
64 ++ static void destroy_instance(void) {
65 ++ if (remote_objref_ != nullptr) {
66 ++ delete remote_objref_;
67 ++ remote_objref_ = nullptr;
68 ++ }
69 ++ }
70 ++
71 + int gather_mem_handles(nvshmemi_symmetric_heap &obj, uint64_t heap_offset, size_t size);
72 + /* On-demand registration and release of memory */
73 + int register_mem_handle(nvshmem_mem_handle_t *local_handles, int transport_idx,
74 +diff --git a/src/modules/bootstrap/uid/bootstrap_uid.cpp b/src/modules/bootstrap/uid/bootstrap_uid.cpp
75 +index a1fa748..788fa96 100644
76 +--- a/src/modules/bootstrap/uid/bootstrap_uid.cpp
77 ++++ b/src/modules/bootstrap/uid/bootstrap_uid.cpp
78 +@@ -630,6 +630,11 @@ int nvshmemi_bootstrap_plugin_pre_init(bootstrap_handle_t* handle, const int abi
79 + // Discover the network for bootstrap, if not done previously.
80 + // This code needs to be stateful to be able to be called multiple times by the caller
81 + BOOTSTRAP_CHECK(bootstrap_net_init());
82 ++ // eep-dev
83 ++ if (handle->pre_init_ops != nullptr) {
84 ++ BOOTSTRAP_PTR_FREE(handle->pre_init_ops);
85 ++ handle->pre_init_ops = nullptr;
86 ++ }
87 + if (handle->pre_init_ops == nullptr) {
88 + BOOTSTRAP_CALLOC(&handle->pre_init_ops, 1);
89 + handle->pre_init_ops->get_unique_id = bootstrap_get_unique_id;
90 +--
91 +2.43.0
92 +
tools/ep_kernels/elastic_ep/install_eep_libraries.sh new
+86
@@ -0,0 +1,86 @@
1 +#!/bin/bash
2 +
3 +set -ex
4 +
5 +# Default workspace directory
6 +WORKSPACE=$(pwd)/eep_kernels_workspace
7 +INSTALL_NVSHMEM=true
8 +
9 +# Parse command line arguments
10 +while getopts "w:n" opt; do
11 + case $opt in
12 + w)
13 + WORKSPACE="$OPTARG"
14 + ;;
15 + n)
16 + INSTALL_NVSHMEM=false
17 + ;;
18 + \?)
19 + echo "Invalid option: -$OPTARG" >&2
20 + exit 1
21 + ;;
22 + esac
23 +done
24 +
25 +if [ ! -d "$WORKSPACE" ]; then
26 + mkdir -p $WORKSPACE
27 +fi
28 +
29 +
30 +# install dependencies if not installed
31 +pip3 install cmake torch ninja
32 +
33 +# build nvshmem
34 +pushd $WORKSPACE
35 +# Reset NVSHMEM build if requested
36 +if [ "$INSTALL_NVSHMEM" = true ]; then
37 + mkdir -p nvshmem_src
38 + wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.2.5/source/nvshmem_src_3.2.5-1.txz
39 + tar -xvf nvshmem_src_3.2.5-1.txz -C nvshmem_src --strip-components=1
40 + pushd nvshmem_src
41 + wget https://github.com/deepseek-ai/DeepEP/raw/main/third-party/nvshmem.patch
42 + git init
43 + git apply -vvv nvshmem.patch
44 + git apply --reject --whitespace=fix ../../eep_nvshmem.patch
45 +else
46 + pushd nvshmem_src
47 +fi
48 +
49 +# assume CUDA_HOME is set correctly
50 +if [ -z "$CUDA_HOME" ]; then
51 + echo "CUDA_HOME is not set, please set it to your CUDA installation directory."
52 + exit 1
53 +fi
54 +
55 +# disable all features except IBGDA
56 +export NVSHMEM_IBGDA_SUPPORT=1
57 +
58 +export NVSHMEM_SHMEM_SUPPORT=0
59 +export NVSHMEM_UCX_SUPPORT=0
60 +export NVSHMEM_USE_NCCL=0
61 +export NVSHMEM_PMIX_SUPPORT=0
62 +export NVSHMEM_TIMEOUT_DEVICE_POLLING=0
63 +export NVSHMEM_USE_GDRCOPY=0
64 +export NVSHMEM_IBRC_SUPPORT=0
65 +export NVSHMEM_BUILD_TESTS=0
66 +export NVSHMEM_BUILD_EXAMPLES=0
67 +export NVSHMEM_MPI_SUPPORT=0
68 +export NVSHMEM_BUILD_HYDRA_LAUNCHER=0
69 +export NVSHMEM_BUILD_TXZ_PACKAGE=0
70 +export NVSHMEM_TIMEOUT_DEVICE_POLLING=0
71 +
72 +cmake -G Ninja -S . -B $WORKSPACE/nvshmem_build/ -DCMAKE_INSTALL_PREFIX=$WORKSPACE/nvshmem_install
73 +cmake --build $WORKSPACE/nvshmem_build/ --target install
74 +
75 +popd
76 +
77 +export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH
78 +
79 +# build and install pplx, require pytorch installed
80 +pushd $WORKSPACE
81 +git clone https://github.com/ppl-ai/pplx-kernels
82 +cd pplx-kernels
83 +# see https://github.com/pypa/pip/issues/9955#issuecomment-838065925
84 +# PIP_NO_BUILD_ISOLATION=0 disables build isolation
85 +PIP_NO_BUILD_ISOLATION=0 TORCH_CUDA_ARCH_LIST=9.0a+PTX pip install . --no-deps -v
86 +
vllm/config.py
+13
@@ -2008,6 +2008,19 @@ class ParallelConfig:
2008 aggregated_has_unfinished = bool(tensor.item())
2009 return aggregated_has_unfinished
2010
2011 + @staticmethod
2012 + def sync_kv_cache_memory_size(dp_group: "ProcessGroup",
2013 + kv_cache_memory: int) -> int:
2014 + if kv_cache_memory == -1:
2015 + kv_cache_memory = torch.iinfo(torch.int64).max
2016 + tensor = torch.tensor([kv_cache_memory],
2017 + dtype=torch.int64,
2018 + device="cpu")
2019 + # we cannot use broadcast for stateless dp group since it depends
2020 + # on global rank
2021 + torch.distributed.all_reduce(tensor, op=ReduceOp.MIN, group=dp_group)
2022 + return tensor.item()
2023 +
2024 def compute_hash(self):
2025 """
2026 Provide a hash that uniquely identifies all the configs
vllm/distributed/eplb/eplb_state.py
+219 -33
@@ -29,12 +29,15 @@ physical experts.
29 import time
30 from collections.abc import Sequence
31 from dataclasses import dataclass
32 +from typing import Optional, Union
33
34 import torch
34 -from torch.distributed import all_gather, all_reduce
35 +from torch.distributed import ProcessGroup, all_gather, all_reduce
36
37 from vllm.config import ParallelConfig
37 -from vllm.distributed.parallel_state import get_ep_group, get_node_count
38 +from vllm.distributed.parallel_state import (get_ep_group, get_node_count,
39 + in_the_same_node_as)
40 +from vllm.distributed.utils import StatelessProcessGroup
41 from vllm.logger import init_logger
42 from vllm.model_executor.models.interfaces import MixtureOfExperts
43
@@ -172,6 +175,9 @@ class EplbState:
175 model: MixtureOfExperts,
176 device: torch.device,
177 parallel_config: ParallelConfig,
178 + global_expert_load: Optional[torch.Tensor] = None,
179 + old_global_expert_indices: Optional[torch.Tensor] = None,
180 + rank_mapping: Optional[dict[int, int]] = None,
181 ) -> "EplbState":
182 """
183 Build the initial EPLB state.
@@ -185,8 +191,16 @@ class EplbState:
191 physical_to_logical_map_list,
192 device=device,
193 )
194 + # Assuming 8 GPUs per node, this supports up to
195 + # (1023 + 1) / 8 = 128 nodes for now.
196 + # TODO(rui): make this configurable
197 + MAX_EXPERT_REDUNDANCY = 1023
198 + assert model.num_redundant_experts <= MAX_EXPERT_REDUNDANCY, (
199 + f"num_redundant_experts {model.num_redundant_experts} "
200 + f"must be less than or equal to {MAX_EXPERT_REDUNDANCY}")
201 + max_slots_per_logical_expert = MAX_EXPERT_REDUNDANCY + 1
202 logical_to_physical_map = torch.full(
189 - (model.num_logical_experts, model.num_redundant_experts + 1),
203 + (model.num_logical_experts, max_slots_per_logical_expert),
204 -1,
205 device=device,
206 )
@@ -235,11 +249,63 @@ class EplbState:
249 expert_rearrangement_step = max(
250 0, eplb_step_interval - eplb_step_interval // 4)
251
252 + if global_expert_load is not None:
253 + ep_group = get_ep_group().device_group
254 + assert global_expert_load.shape == (model.num_moe_layers,
255 + model.num_logical_experts)
256 + assert global_expert_load.dtype == torch.int64
257 +
258 + num_replicas = model.num_physical_experts
259 + num_groups = model.num_expert_groups
260 + num_nodes = get_node_count()
261 + num_gpus = ep_group.size()
262 +
263 + if num_gpus % num_nodes != 0:
264 + num_nodes = 1
265 + logger.warning_once(
266 + f"num_gpus % num_nodes != 0, "
267 + "not using hierarchical rearrangement algorithm.\n"
268 + f"{num_gpus=}, {num_nodes=}")
269 +
270 + # Get new expert mappings
271 + (
272 + new_physical_to_logical_map,
273 + new_logical_to_physical_map,
274 + new_logical_replica_count,
275 + ) = (rebalance_experts(
276 + global_expert_load,
277 + num_replicas,
278 + num_groups,
279 + num_nodes,
280 + num_gpus,
281 + ))
282 +
283 + max_physical_slots = new_logical_to_physical_map.shape[-1]
284 + assert max_physical_slots <= logical_to_physical_map.shape[-1]
285 + new_logical_to_physical_map = torch.nn.functional.pad(
286 + new_logical_to_physical_map,
287 + (0, logical_to_physical_map.shape[-1] - max_physical_slots),
288 + value=-1,
289 + )
290 + physical_to_logical_map = new_physical_to_logical_map.to(device)
291 + logical_to_physical_map.copy_(new_logical_to_physical_map)
292 + logical_replica_count.copy_(new_logical_replica_count)
293 +
294 model.set_eplb_state(
295 expert_load_pass,
296 logical_to_physical_map,
297 logical_replica_count,
298 )
299 + if global_expert_load is not None:
300 + rearrange_expert_weights_inplace(
301 + old_global_expert_indices,
302 + new_physical_to_logical_map,
303 + model.expert_weights,
304 + ep_group,
305 + False,
306 + rank_mapping,
307 + )
308 + expert_rearrangement_step = 0
309
310 return cls(
311 physical_to_logical_map,
@@ -337,7 +403,10 @@ class EplbState:
403
404 def rearrange(self,
405 model: MixtureOfExperts,
340 - is_profile: bool = False) -> None:
406 + is_profile: bool = False,
407 + execute_shuffle: bool = True,
408 + global_expert_load: Optional[torch.Tensor] = None,
409 + rank_mapping: Optional[dict[int, int]] = None) -> None:
410 """
411 Rearrange the experts according to the current load.
412 """
@@ -353,42 +422,79 @@ class EplbState:
422 logger.info("Rearranging experts %s...",
423 "(profile)" if is_profile else "")
424
356 - # This mapping is only used here, so we do not store it in the state
357 - physical_expert_start = ep_rank * model.num_local_physical_experts
358 - physical_expert_end = (physical_expert_start +
359 - model.num_local_physical_experts)
360 - # (num_moe_layers, num_local_physical_experts)
361 - local_physical_to_logical_map = self.physical_to_logical_map[
362 - :,
363 - physical_expert_start:physical_expert_end,
364 - ]
425 + if global_expert_load is None:
426 + # This mapping is only used here, so we do not store it in the state
427 + physical_expert_start = ep_rank * model.num_local_physical_experts
428 + physical_expert_end = (physical_expert_start +
429 + model.num_local_physical_experts)
430 + # (num_moe_layers, num_local_physical_experts)
431 + local_physical_to_logical_map = self.physical_to_logical_map[
432 + :,
433 + physical_expert_start:physical_expert_end,
434 + ]
435
366 - # Map the local physical expert load to global logical experts
367 - logical_expert_load_window = torch.zeros(
368 - self.expert_load_window_size,
369 - model.num_moe_layers,
370 - model.num_logical_experts,
371 - dtype=self.expert_load_window.dtype,
372 - device=self.expert_load_window.device,
373 - )
374 - logical_expert_load_window.scatter_add_(
375 - dim=-1,
376 - index=local_physical_to_logical_map.unsqueeze(0).expand_as(
377 - self.expert_load_window).long(),
378 - src=self.expert_load_window,
379 - )
436 + # Map the local physical expert load to global logical experts
437 + logical_expert_load_window = torch.zeros(
438 + self.expert_load_window_size,
439 + model.num_moe_layers,
440 + model.num_logical_experts,
441 + dtype=self.expert_load_window.dtype,
442 + device=self.expert_load_window.device,
443 + )
444 + logical_expert_load_window.scatter_add_(
445 + dim=-1,
446 + index=local_physical_to_logical_map.unsqueeze(0).expand_as(
447 + self.expert_load_window).long(),
448 + src=self.expert_load_window,
449 + )
450
381 - # Perform all-reduce to get the expert load across all ranks
382 - global_expert_load_window = logical_expert_load_window.sum(dim=0)
383 - all_reduce(global_expert_load_window, group=ep_group)
451 + if not execute_shuffle:
452 + metadata = torch.tensor(
453 + [
454 + model.num_moe_layers, model.num_logical_experts,
455 + self.physical_to_logical_map.shape[1]
456 + ],
457 + dtype=torch.int32,
458 + device="cpu",
459 + )
460 + torch.distributed.broadcast(metadata,
461 + group=get_ep_group().cpu_group,
462 + group_src=0)
463 +
464 + # Perform all-reduce to get the expert load across all ranks
465 + global_expert_load_window = logical_expert_load_window.sum(dim=0)
466 + all_reduce(global_expert_load_window, group=ep_group)
467 +
468 + if not execute_shuffle:
469 + # (num_moe_layers, old_num_physical_experts)
470 + old_global_expert_indices = self.physical_to_logical_map
471 + torch.distributed.broadcast(old_global_expert_indices,
472 + group=ep_group,
473 + group_src=0)
474 + return global_expert_load_window
475 + else:
476 + assert execute_shuffle
477 + global_expert_load_window = global_expert_load
478
479 # TODO(bowen): Treat differently for prefill and decode nodes
480 num_replicas = model.num_physical_experts
481 num_groups = model.num_expert_groups
388 - num_nodes = get_node_count()
389 - num_gpus = ep_group.size()
482 + if rank_mapping is not None and len(rank_mapping) == ep_group.size():
483 + # NOTE(yongji): scale down, we need to rebalance the experts on
484 + # remaining GPUs, transfer the experts while we haven't shutdown
485 + # the GPUs to be released.
486 + cpu_group = get_ep_group().cpu_group
487 + num_nodes = _node_count_with_rank_mapping(cpu_group, rank_mapping)
488 + num_gpus = sum(new_rank != -1
489 + for new_rank in rank_mapping.values())
490 + num_replicas = num_replicas // ep_group.size(
491 + ) * num_gpus # handle num replicas change
492 + else:
493 + num_nodes = get_node_count()
494 + num_gpus = ep_group.size()
495
496 if num_gpus % num_nodes != 0:
497 + self.num_nodes = 1
498 logger.warning_once(
499 f"num_gpus % num_nodes != 0, "
500 "not using hierarchical rearrangement algorithm.\n"
@@ -414,10 +520,24 @@ class EplbState:
520 model.expert_weights,
521 ep_group,
522 is_profile,
523 + rank_mapping,
524 )
525
526 if not is_profile:
420 - self.physical_to_logical_map.copy_(new_physical_to_logical_map)
527 + if self.physical_to_logical_map.shape[
528 + 1] != new_physical_to_logical_map.shape[1]:
529 + self.physical_to_logical_map = new_physical_to_logical_map.to(
530 + self.physical_to_logical_map.device)
531 + else:
532 + self.physical_to_logical_map.copy_(new_physical_to_logical_map)
533 + max_physical_slots = new_logical_to_physical_map.shape[-1]
534 + assert max_physical_slots <= self.logical_to_physical_map.shape[-1]
535 + new_logical_to_physical_map = torch.nn.functional.pad(
536 + new_logical_to_physical_map,
537 + (0,
538 + self.logical_to_physical_map.shape[-1] - max_physical_slots),
539 + value=-1,
540 + )
541 self.logical_to_physical_map.copy_(new_logical_to_physical_map)
542 self.logical_replica_count.copy_(new_logical_replica_count)
543
@@ -430,3 +550,69 @@ class EplbState:
550 " (profile) " if is_profile else " ",
551 time_end - time_start,
552 )
553 +
554 + @staticmethod
555 + def recv_state() -> tuple[torch.Tensor, torch.Tensor]:
556 + """
557 + Receive the expert load and old placement from the master rank.
558 + """
559 + ep_group = get_ep_group()
560 + metadata = torch.empty(3, dtype=torch.int32, device="cpu")
561 + torch.distributed.broadcast(metadata,
562 + group=ep_group.cpu_group,
563 + group_src=0)
564 + num_moe_layers, num_logical_experts, num_old_physical_experts = (
565 + metadata.tolist())
566 + global_expert_load = torch.zeros(
567 + (num_moe_layers, num_logical_experts),
568 + dtype=torch.int64,
569 + device=ep_group.device,
570 + )
571 + all_reduce(global_expert_load, group=ep_group.device_group)
572 + old_global_expert_indices = torch.empty(
573 + (num_moe_layers, num_old_physical_experts),
574 + dtype=torch.int64,
575 + device=ep_group.device,
576 + )
577 + torch.distributed.broadcast(old_global_expert_indices,
578 + group=ep_group.device_group,
579 + group_src=0)
580 +
581 + return global_expert_load, old_global_expert_indices
582 +
583 +
584 +def _node_count_with_rank_mapping(
585 + pg: Union[ProcessGroup, StatelessProcessGroup],
586 + rank_mapping: dict[int, int],
587 +) -> int:
588 + if isinstance(pg, ProcessGroup):
589 + world_size = torch.distributed.get_world_size(group=pg)
590 + else:
591 + world_size = pg.world_size
592 +
593 + if world_size == 1:
594 + return 1
595 +
596 + # Build node assignment map
597 + node_assignment = [0] * world_size # rank -> node_id
598 + next_node_id = 0
599 +
600 + for current_rank in range(world_size):
601 + if node_assignment[current_rank] != 0:
602 + continue # Already assigned to a node
603 +
604 + assert current_rank in rank_mapping
605 + if rank_mapping[current_rank] == -1:
606 + continue # Pending shutdown
607 +
608 + # Assign current rank to a new node
609 + next_node_id += 1
610 + node_assignment[current_rank] = next_node_id
611 +
612 + # Find all ranks on the same node as current_rank
613 + same_node_flags = in_the_same_node_as(pg, current_rank)
614 + for other_rank, is_same_node in enumerate(same_node_flags):
615 + if is_same_node and node_assignment[other_rank] == 0:
616 + node_assignment[other_rank] = next_node_id
617 +
618 + return next_node_id
vllm/distributed/eplb/rebalance_execute.py
+117
@@ -8,6 +8,7 @@ This involves the exchange of expert weights between GPUs.
8
9 from collections.abc import Iterable, MutableSequence, Sequence
10 from functools import partial
11 +from typing import Optional
12
13 import torch
14 from torch.distributed import (P2POp, ProcessGroup, all_gather,
@@ -127,6 +128,8 @@ def shuffle_layer(
128 dst_global = local2global(dst)
129 if is_received_locally[dst]:
130 continue
131 + if old_indices[src_global] == -1 or new_indices[dst_global] == -1:
132 + continue
133 if old_indices[src_global] == new_indices[dst_global]:
134 is_received_locally[dst] = True
135 for weight, buffer in zip(expert_weights,
@@ -139,6 +142,8 @@ def shuffle_layer(
142 experts_send_loc: dict[int, int] = {}
143 for src in range(num_local_experts):
144 expert = old_indices[local2global(src)]
145 + if expert == -1:
146 + continue
147 if expert in experts_send_loc:
148 continue
149 experts_send_loc[expert] = src
@@ -181,6 +186,8 @@ def shuffle_layer(
186 if is_received_locally[dst]:
187 continue
188 expert = new_indices[local2global(dst)]
189 + if expert == -1:
190 + continue
191 if expert in experts_recv_loc:
192 continue
193 experts_recv_loc[expert] = dst
@@ -227,6 +234,8 @@ def shuffle_layer(
234 weight[dst].copy_(buffer[dst])
235 else:
236 expert = new_indices[local2global(dst)]
237 + if expert == -1:
238 + continue
239 src = experts_recv_loc[expert]
240 for weight, buffer in zip(expert_weights, expert_weights_buffer):
241 weight[dst].copy_(buffer[src])
@@ -238,6 +247,7 @@ def rearrange_expert_weights_inplace(
247 expert_weights: Sequence[Iterable[torch.Tensor]],
248 ep_group: ProcessGroup,
249 is_profile: bool = False,
250 + rank_mapping: Optional[dict[int, int]] = None,
251 ) -> None:
252 """
253 Rearranges the expert weights in place according to the new expert indices.
@@ -256,7 +266,28 @@ def rearrange_expert_weights_inplace(
266 is_profile (bool): If `True`, do not perform any actual weight copy.
267 This is used during profile run, where we only perform dummy
268 communications to reserve enough memory for the buffers.
269 + rank_mapping: A dictionary mapping old rank to new rank.
270 """
271 + if rank_mapping is not None:
272 + if len(rank_mapping) == ep_group.size():
273 + # scale down
274 + new_global_expert_indices = \
275 + _map_new_expert_indices_with_rank_mapping(
276 + new_global_expert_indices,
277 + rank_mapping,
278 + )
279 + else:
280 + # scale up
281 + old_global_expert_indices = \
282 + _map_old_expert_indices_with_rank_mapping(
283 + old_global_expert_indices,
284 + rank_mapping,
285 + ep_group.size(),
286 + )
287 +
288 + assert old_global_expert_indices.shape[
289 + 1] == new_global_expert_indices.shape[1]
290 +
291 num_moe_layers, num_physical_experts = old_global_expert_indices.shape
292 assert len(expert_weights) == num_moe_layers
293
@@ -304,4 +335,90 @@ def rearrange_expert_weights_inplace(
335 )
336
337
338 +def _map_old_expert_indices_with_rank_mapping(
339 + old_global_expert_indices: torch.Tensor,
340 + rank_mapping: dict[int, int],
341 + new_ep_size: int,
342 +) -> torch.Tensor:
343 + """
344 + Map the old global expert indices to the new global expert indices.
345 +
346 + Args:
347 + old_global_expert_indices:
348 + Shape (num_layers, old_ep_size * num_local_physical_experts).
349 + rank_mapping: Mapping from old rank to new rank.
350 + new_ep_size: New expert parallelism size.
351 +
352 + Returns:
353 + Mapped expert indices with shape
354 + (num_layers, new_ep_size * num_local_physical_experts).
355 + """
356 + num_layers, old_num_physical_experts = old_global_expert_indices.shape
357 + assert rank_mapping, "Rank mapping is required"
358 +
359 + # Get sizes from parameters and rank_mapping
360 + old_ep_size = len(rank_mapping)
361 + num_local_physical_experts = old_num_physical_experts // old_ep_size
362 + new_num_physical_experts = new_ep_size * num_local_physical_experts
363 +
364 + # Create mapped tensor with new shape, initialized to -1
365 + mapped_expert_indices = torch.full(
366 + (num_layers, new_num_physical_experts),
367 + fill_value=-1,
368 + dtype=old_global_expert_indices.dtype,
369 + device=old_global_expert_indices.device,
370 + )
371 +
372 + # Handle rank mapping (scale up/down with rank changes)
373 + for old_rank in range(old_ep_size):
374 + new_rank = rank_mapping.get(old_rank)
375 + if new_rank is not None and new_rank >= 0 and new_rank < new_ep_size:
376 + # This old rank exists in the new configuration
377 + old_start_idx = old_rank * num_local_physical_experts
378 + old_end_idx = (old_rank + 1) * num_local_physical_experts
379 + new_start_idx = new_rank * num_local_physical_experts
380 + new_end_idx = (new_rank + 1) * num_local_physical_experts
381 +
382 + mapped_expert_indices[:, new_start_idx:new_end_idx] = \
383 + old_global_expert_indices[:, old_start_idx:old_end_idx]
384 + # If new_rank is None or >= new_ep_size, the experts remain -1
385 + # (scale down case)
386 +
387 + return mapped_expert_indices
388 +
389 +
390 +def _map_new_expert_indices_with_rank_mapping(
391 + new_global_expert_indices: torch.Tensor,
392 + rank_mapping: dict[int, int],
393 +) -> torch.Tensor:
394 + num_layers, new_num_physical_experts = new_global_expert_indices.shape
395 + assert rank_mapping, "Rank mapping is required"
396 +
397 + # Get sizes from parameters and rank_mapping
398 + old_ep_size = len(rank_mapping)
399 + new_ep_size = sum(new_rank != -1 for new_rank in rank_mapping.values())
400 + num_local_physical_experts = new_num_physical_experts // new_ep_size
401 + old_num_physical_experts = old_ep_size * num_local_physical_experts
402 +
403 + mapped_expert_indices = torch.full(
404 + (num_layers, old_num_physical_experts),
405 + fill_value=-1,
406 + dtype=new_global_expert_indices.dtype,
407 + device=new_global_expert_indices.device,
408 + )
409 +
410 + for old_rank in range(old_ep_size):
411 + new_rank = rank_mapping[old_rank]
412 + if new_rank >= 0 and new_rank < new_ep_size:
413 + old_start_idx = old_rank * num_local_physical_experts
414 + old_end_idx = (old_rank + 1) * num_local_physical_experts
415 + new_start_idx = new_rank * num_local_physical_experts
416 + new_end_idx = (new_rank + 1) * num_local_physical_experts
417 +
418 + mapped_expert_indices[:, old_start_idx:old_end_idx] = \
419 + new_global_expert_indices[:, new_start_idx:new_end_idx]
420 +
421 + return mapped_expert_indices
422 +
423 +
424 __all__ = ["rearrange_expert_weights_inplace"]
vllm/engine/protocol.py
+6
@@ -324,3 +324,9 @@ class EngineClient(ABC):
324 async def add_lora(self, lora_request: LoRARequest) -> None:
325 """Load a new LoRA adapter into the engine for future requests."""
326 ...
327 +
328 + async def scale_elastic_ep(self,
329 + new_data_parallel_size: int,
330 + drain_timeout: int = 300) -> None:
331 + """Scale the engine"""
332 + raise NotImplementedError
vllm/entrypoints/openai/api_server.py
+105
@@ -1018,6 +1018,73 @@ if envs.VLLM_SERVER_DEV_MODE:
1018 return JSONResponse(content={"is_sleeping": is_sleeping})
1019
1020
1021 +@router.post("/scale_elastic_ep",
1022 + dependencies=[Depends(validate_json_request)],
1023 + responses={
1024 + HTTPStatus.OK.value: {
1025 + "model": dict
1026 + },
1027 + HTTPStatus.BAD_REQUEST.value: {
1028 + "model": ErrorResponse
1029 + },
1030 + HTTPStatus.REQUEST_TIMEOUT.value: {
1031 + "model": ErrorResponse
1032 + },
1033 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {
1034 + "model": ErrorResponse
1035 + },
1036 + })
1037 +async def scale_elastic_ep(raw_request: Request):
1038 + try:
1039 + body = await raw_request.json()
1040 + except json.JSONDecodeError as e:
1041 + raise HTTPException(status_code=400,
1042 + detail="Invalid JSON format") from e # noqa: B904
1043 +
1044 + new_data_parallel_size = body.get("new_data_parallel_size")
1045 + drain_timeout = body.get("drain_timeout", 120) # Default 2 minutes
1046 +
1047 + if new_data_parallel_size is None:
1048 + raise HTTPException(status_code=400,
1049 + detail="new_data_parallel_size is required")
1050 +
1051 + if not isinstance(new_data_parallel_size,
1052 + int) or new_data_parallel_size <= 0:
1053 + raise HTTPException(
1054 + status_code=400,
1055 + detail="new_data_parallel_size must be a positive integer")
1056 +
1057 + if not isinstance(drain_timeout, int) or drain_timeout <= 0:
1058 + raise HTTPException(status_code=400,
1059 + detail="drain_timeout must be a positive integer")
1060 +
1061 + # Set scaling flag to prevent new requests
1062 + global _scaling_elastic_ep
1063 + _scaling_elastic_ep = True
1064 + client = engine_client(raw_request)
1065 + try:
1066 + await client.scale_elastic_ep(new_data_parallel_size, drain_timeout)
1067 + return JSONResponse({
1068 + "message":
1069 + f"Scaled to {new_data_parallel_size} "
1070 + "data parallel engines",
1071 + })
1072 + except TimeoutError as e:
1073 + raise HTTPException(status_code=408,
1074 + detail="Scale failed due to request drain timeout "
1075 + f"after {drain_timeout} seconds") from e
1076 + except Exception as e:
1077 + logger.error("Scale failed: %s", e)
1078 + raise HTTPException(status_code=500, detail="Scale failed") from e
1079 + finally:
1080 + _scaling_elastic_ep = False
1081 +
1082 +
1083 +@router.post("/is_scaling_elastic_ep")
1084 +async def is_scaling_elastic_ep(raw_request: Request):
1085 + return JSONResponse({"is_scaling_elastic_ep": _scaling_elastic_ep})
1086 +
1087 +
1088 # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
1089 # (requires typing_extensions >= 4.13)
1090 RequestType = Any
@@ -1216,6 +1283,41 @@ class XRequestIdMiddleware:
1283 return self.app(scope, receive, send_with_request_id)
1284
1285
1286 +# Global variable to track scaling state
1287 +_scaling_elastic_ep = False
1288 +
1289 +
1290 +class ScalingMiddleware:
1291 + """
1292 + Middleware that checks if the model is currently scaling and
1293 + returns a 503 Service Unavailable response if it is.
1294 +
1295 + This middleware applies to all HTTP requests and prevents
1296 + processing when the model is in a scaling state.
1297 + """
1298 +
1299 + def __init__(self, app: ASGIApp) -> None:
1300 + self.app = app
1301 +
1302 + def __call__(self, scope: Scope, receive: Receive,
1303 + send: Send) -> Awaitable[None]:
1304 + if scope["type"] != "http":
1305 + return self.app(scope, receive, send)
1306 +
1307 + # Check global scaling state
1308 + global _scaling_elastic_ep
1309 + if _scaling_elastic_ep:
1310 + # Return 503 Service Unavailable response
1311 + response = JSONResponse(content={
1312 + "error":
1313 + "The model is currently scaling. Please try again later."
1314 + },
1315 + status_code=503)
1316 + return response(scope, receive, send)
1317 +
1318 + return self.app(scope, receive, send)
1319 +
1320 +
1321 def _extract_content_from_chunk(chunk_data: dict) -> str:
1322 """Extract content from a streaming response chunk."""
1323 try:
@@ -1404,6 +1506,9 @@ def build_app(args: Namespace) -> FastAPI:
1506 if args.enable_request_id_headers:
1507 app.add_middleware(XRequestIdMiddleware)
1508
1509 + # Add scaling middleware to check for scaling state
1510 + app.add_middleware(ScalingMiddleware)
1511 +
1512 if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE:
1513 logger.warning("CAUTION: Enabling log response in the API Server. "
1514 "This can include sensitive information and should be "
vllm/executor/uniproc_executor.py
+9
@@ -12,6 +12,7 @@ from vllm.executor.executor_base import ExecutorBase
12 from vllm.logger import init_logger
13 from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
14 run_method)
15 +from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
16 from vllm.worker.worker_base import WorkerWrapperBase
17
18 logger = init_logger(__name__)
@@ -62,6 +63,14 @@ class UniProcExecutor(ExecutorBase):
63 # it's running.
64 return
65
66 + def reinitialize_distributed(
67 + self, reconfig_request: ReconfigureDistributedRequest) -> None:
68 + self.driver_worker.reinitialize_distributed(reconfig_request)
69 + if reconfig_request.new_data_parallel_rank == \
70 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
71 + self.shutdown()
72 + return
73 +
74
75 UniProcExecutorAsync = UniProcExecutor
76
vllm/model_executor/layers/fused_moe/layer.py
+31 -8
@@ -265,9 +265,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
265 prepare_finalize: FusedMoEPrepareAndFinalize,
266 moe: FusedMoEConfig,
267 ) -> FusedMoEPermuteExpertsUnpermute:
268 -
269 - assert self.fused_experts == fused_experts
270 -
268 if (prepare_finalize.activation_format ==
269 FusedMoEActivationFormat.BatchedExperts):
270 logger.debug("BatchedTritonExperts %s", self.moe)
@@ -375,8 +372,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
372 logical_replica_count: Optional[torch.Tensor] = None,
373 ) -> torch.Tensor:
374 if enable_eplb:
378 - raise NotImplementedError(
379 - "EPLB not supported for `UnquantizedFusedMoEMethod` yet.")
375 + assert expert_load_view is not None
376 + assert logical_to_physical_map is not None
377 + assert logical_replica_count is not None
378 + assert isinstance(layer, FusedMoE)
379
380 return self.forward(
381 x=x,
@@ -393,7 +392,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
392 scoring_func=scoring_func,
393 e_score_correction_bias=e_score_correction_bias,
394 activation=activation,
396 - apply_router_weight_on_input=apply_router_weight_on_input)
395 + apply_router_weight_on_input=apply_router_weight_on_input,
396 + enable_eplb=enable_eplb,
397 + expert_load_view=expert_load_view,
398 + logical_to_physical_map=logical_to_physical_map,
399 + logical_replica_count=logical_replica_count,
400 + )
401
402 def forward_cuda(
403 self,
@@ -412,6 +416,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
416 e_score_correction_bias: Optional[torch.Tensor] = None,
417 apply_router_weight_on_input: bool = False,
418 activation: str = "silu",
419 + enable_eplb: bool = False,
420 + expert_load_view: Optional[torch.Tensor] = None,
421 + logical_to_physical_map: Optional[torch.Tensor] = None,
422 + logical_replica_count: Optional[torch.Tensor] = None,
423 ) -> torch.Tensor:
424
425 topk_weights, topk_ids = FusedMoE.select_experts(
@@ -425,7 +433,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
433 custom_routing_function=custom_routing_function,
434 scoring_func=scoring_func,
435 e_score_correction_bias=e_score_correction_bias,
428 - indices_type=self.topk_indices_dtype)
436 + indices_type=self.topk_indices_dtype,
437 + enable_eplb=enable_eplb,
438 + expert_map=expert_map,
439 + expert_load_view=expert_load_view,
440 + logical_to_physical_map=logical_to_physical_map,
441 + logical_replica_count=logical_replica_count)
442
443 if self.rocm_aiter_moe_enabled:
444 return self.rocm_aiter_fused_experts(
@@ -730,7 +743,8 @@ class FusedMoE(torch.nn.Module):
743 if self.enable_eplb:
744 from vllm.model_executor.layers.quantization.fp8 import (
745 Fp8MoEMethod)
733 - if not isinstance(quant_method, Fp8MoEMethod):
746 + if not isinstance(quant_method,
747 + (Fp8MoEMethod, UnquantizedFusedMoEMethod)):
748 # TODO: Add support for additional quantization methods.
749 # The implementation for other quantization methods does not
750 # contain essential differences, but the current quant API
@@ -821,6 +835,15 @@ class FusedMoE(torch.nn.Module):
835 def use_flashinfer_cutlass_kernels(self):
836 return self.moe_parallel_config.use_flashinfer_cutlass_kernels
837
838 + def update_expert_map(self):
839 + # ep_size and ep_rank should already be updated
840 + assert self.expert_map is not None
841 + with self.expert_map.device:
842 + self.local_num_experts, self.expert_map = determine_expert_map(
843 + ep_size=self.ep_size,
844 + ep_rank=self.ep_rank,
845 + global_num_experts=self.global_num_experts)
846 +
847 def _load_per_tensor_weight_scale(self, shard_id: str,
848 param: torch.nn.Parameter,
849 loaded_weight: torch.Tensor,
vllm/model_executor/models/deepseek_v2.py
+20 -3
@@ -776,6 +776,24 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
776 logical_replica_count=logical_replica_count,
777 )
778
779 + def update_physical_experts_metadata(
780 + self,
781 + num_physical_experts: int,
782 + num_local_physical_experts: int,
783 + ) -> None:
784 + assert self.num_local_physical_experts == num_local_physical_experts
785 + self.num_physical_experts = num_physical_experts
786 + self.num_local_physical_experts = num_local_physical_experts
787 + self.num_redundant_experts = (num_physical_experts -
788 + self.num_logical_experts)
789 + for layer in self.model.layers:
790 + if isinstance(layer.mlp, DeepseekV2MoE):
791 + moe = layer.mlp
792 + moe.n_local_physical_experts = num_local_physical_experts
793 + moe.n_physical_experts = num_physical_experts
794 + moe.n_redundant_experts = self.num_redundant_experts
795 + moe.experts.update_expert_map()
796 +
797 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
798 return self.model.get_input_embeddings(input_ids)
799
@@ -931,9 +949,8 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
949
950 def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,
951 weight_name: str) -> Optional[int]:
934 - if hasattr(config,
935 - "num_nextn_predict_layers") and (config.num_nextn_predict_layers
936 - > 0):
952 + if (hasattr(config, "num_nextn_predict_layers")
953 + and config.num_nextn_predict_layers > 0):
954 layer_idx = config.num_hidden_layers
955 for i in range(config.num_nextn_predict_layers):
956 if weight_name.startswith(f"model.layers.{layer_idx+i}."):
vllm/model_executor/models/interfaces.py
+7
@@ -543,6 +543,13 @@ class MixtureOfExperts(Protocol):
543 """
544 ...
545
546 + def update_physical_experts_metadata(
547 + self,
548 + num_physical_experts: int,
549 + num_local_physical_experts: int,
550 + ) -> None:
551 + ...
552 +
553
554 def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
555 return isinstance(model, MixtureOfExperts)
vllm/v1/engine/__init__.py
+16
@@ -177,3 +177,19 @@ class EngineCoreRequestType(enum.Enum):
177 UTILITY = b'\x03'
178 # Sentinel used within EngineCoreProc.
179 EXECUTOR_FAILED = b'\x04'
180 +
181 +
182 +class ReconfigureDistributedRequest(msgspec.Struct):
183 + new_data_parallel_size: int
184 + new_data_parallel_rank: int
185 + new_data_parallel_rank_local: int
186 + new_data_parallel_master_ip: str
187 + new_data_parallel_master_port: int
188 +
189 +
190 +class ReconfigureRankType(enum.IntEnum):
191 + """
192 + Rank type for reconfiguring distributed request.
193 + """
194 + KEEP_CURRENT_RANK = -1
195 + SHUTDOWN_CURRENT_RANK = -2
vllm/v1/engine/async_llm.py
+58
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 +import time
5 from collections.abc import AsyncGenerator, Mapping
6 from copy import copy
7 from typing import Any, Optional, Union
@@ -608,6 +609,63 @@ class AsyncLLM(EngineClient):
609 return await self.engine_core.collective_rpc_async(
610 method, timeout, args, kwargs)
611
612 + async def wait_for_requests_to_drain(self, drain_timeout: int = 300):
613 + """Wait for all requests to be drained."""
614 + start_time = time.time()
615 + while time.time() - start_time < drain_timeout:
616 + if not self.engine_core.dp_engines_running():
617 + logger.info("Engines are idle, requests have been drained")
618 + return
619 +
620 + logger.info(
621 + "Engines are still running, waiting for requests to drain...")
622 + await asyncio.sleep(1) # Wait 1 second before checking again
623 +
624 + raise TimeoutError(f"Timeout reached after {drain_timeout} seconds "
625 + "waiting for requests to drain.")
626 +
627 + async def scale_elastic_ep(self,
628 + new_data_parallel_size: int,
629 + drain_timeout: int = 300):
630 + """
631 + Scale up or down the data parallel size by adding or removing
632 + engine cores.
633 + Args:
634 + new_data_parallel_size: The new number of data parallel workers
635 + drain_timeout:
636 + Maximum time to wait for requests to drain (seconds)
637 + """
638 + old_data_parallel_size = \
639 + self.vllm_config.parallel_config.data_parallel_size
640 + if old_data_parallel_size == new_data_parallel_size:
641 + logger.info("Data parallel size is already %s, skipping scale",
642 + new_data_parallel_size)
643 + return
644 + logger.info(
645 + "Waiting for requests to drain before "
646 + "scaling up to %s engines...", new_data_parallel_size)
647 + await self.wait_for_requests_to_drain(drain_timeout)
648 + logger.info(
649 + "Requests have been drained, proceeding with scale "
650 + "to %s engines", new_data_parallel_size)
651 + await self.engine_core.scale_elastic_ep(new_data_parallel_size)
652 + self.vllm_config.parallel_config.data_parallel_size = \
653 + new_data_parallel_size
654 +
655 + # recreate stat loggers
656 + if new_data_parallel_size > old_data_parallel_size:
657 + stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
658 + vllm_config=self.vllm_config,
659 + log_stats=self.log_stats,
660 + engine_num=new_data_parallel_size,
661 + custom_stat_loggers=None,
662 + )
663 + num_new_engines = len(stat_loggers) - len(self.stat_loggers)
664 + self.stat_loggers.extend(stat_loggers[-num_new_engines:])
665 + else:
666 + for _ in range(old_data_parallel_size - new_data_parallel_size):
667 + self.stat_loggers.pop()
668 +
669 @property
670 def is_running(self) -> bool:
671 # Is None before the loop is started.
vllm/v1/engine/coordinator.py
+31 -1
@@ -200,11 +200,41 @@ class CoordinatorProc:
200 # Ignore subscription messages.
201 continue
202
203 + decoded = msgspec.msgpack.decode(buffer)
204 + if isinstance(decoded, (list, tuple)) and len(
205 + decoded) == 2 and decoded[0] == "SCALE_ELASTIC_EP":
206 + # Handle scale up notification
207 + new_engine_count = decoded[1]
208 + current_count = len(self.engines)
209 + if new_engine_count > current_count:
210 + for _ in range(new_engine_count - current_count):
211 + self.engines.append(EngineState())
212 + # NOTE(yongji): handle the case
213 + # where newly started engines have current_wave = 0
214 + # if existing engines just finished a wave
215 + # and engine_running isn't updated yet at
216 + # CoordinatorProc requests routed to newly started
217 + # engines may not wake up existing engines, as long
218 + # as 0 < request.wave < existing engines'
219 + # current_wave
220 + # we note that 0 is the wave number for the new
221 + # engine
222 + self.engines_running = False
223 + logger.info(
224 + "DPCoordinator scaled up from %s to %s "
225 + "engines", current_count, new_engine_count)
226 + else:
227 + self.engines = self.engines[:new_engine_count]
228 + logger.info(
229 + "DPCoordinator scaled down from %s to %s "
230 + "engines", current_count, new_engine_count)
231 + continue # Skip normal engine notification processing
232 +
233 # We received a message on the front-end XPUB socket,
234 # from an API server sending a new request while the
235 # engines are paused, so that we can wake the other
236 # engines.
207 - engine_to_exclude, wave = msgspec.msgpack.decode(buffer)
237 + engine_to_exclude, wave = decoded
238 if not self.engines_running:
239 if wave < self.current_wave:
240 # If the wave number is stale, ensure the message
vllm/v1/engine/core.py
+64 -5
@@ -32,7 +32,9 @@ from vllm.v1.core.sched.interface import SchedulerInterface
32 from vllm.v1.core.sched.output import SchedulerOutput
33 from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler
34 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
35 - EngineCoreRequestType, UtilityOutput)
35 + EngineCoreRequestType,
36 + ReconfigureDistributedRequest, ReconfigureRankType,
37 + UtilityOutput)
38 from vllm.v1.engine.mm_input_cache import MirroredProcessingCache
39 from vllm.v1.engine.utils import EngineHandshakeMetadata, EngineZmqAddresses
40 from vllm.v1.executor.abstract import Executor
@@ -77,6 +79,8 @@ class EngineCore:
79 self.model_executor.register_failure_callback(
80 executor_fail_callback)
81
82 + self.available_gpu_memory_for_kv_cache = -1
83 +
84 # Setup KV Caches and update CacheConfig after profiling.
85 num_gpu_blocks, num_cpu_blocks, kv_cache_config = \
86 self._initialize_kv_caches(vllm_config)
@@ -137,12 +141,23 @@ class EngineCore:
141 # Get all kv cache needed by the model
142 kv_cache_specs = self.model_executor.get_kv_cache_specs()
143
140 - # Profiles the peak memory usage of the model to determine how much
141 - # memory can be allocated for kv cache.
144 has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
145 if has_kv_cache:
144 - available_gpu_memory = \
145 - self.model_executor.determine_available_memory()
146 + if os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1":
147 + dp_group = getattr(self, "dp_group", None)
148 + assert dp_group is not None
149 + self.available_gpu_memory_for_kv_cache = \
150 + ParallelConfig.sync_kv_cache_memory_size(dp_group, -1)
151 + available_gpu_memory = [
152 + self.available_gpu_memory_for_kv_cache
153 + ] * len(kv_cache_specs)
154 + else:
155 + # Profiles the peak memory usage of the model to determine how
156 + # much memory can be allocated for kv cache.
157 + available_gpu_memory = (
158 + self.model_executor.determine_available_memory())
159 + self.available_gpu_memory_for_kv_cache = \
160 + available_gpu_memory[0]
161 else:
162 # Attention free models don't need memory for kv cache
163 available_gpu_memory = [0] * len(kv_cache_specs)
@@ -989,6 +1004,50 @@ class DPEngineCoreProc(EngineCoreProc):
1004 return ParallelConfig.has_unfinished_dp(self.dp_group,
1005 local_unfinished)
1006
1007 + def reinitialize_distributed(
1008 + self, reconfig_request: ReconfigureDistributedRequest) -> None:
1009 + stateless_destroy_torch_distributed_process_group(self.dp_group)
1010 + self.shutdown()
1011 +
1012 + parallel_config = self.vllm_config.parallel_config
1013 + old_dp_size = parallel_config.data_parallel_size
1014 + parallel_config.data_parallel_size = \
1015 + reconfig_request.new_data_parallel_size
1016 + if reconfig_request.new_data_parallel_rank != -1:
1017 + parallel_config.data_parallel_rank = \
1018 + reconfig_request.new_data_parallel_rank
1019 + # local rank specifies device visibility, it should not be changed
1020 + assert reconfig_request.new_data_parallel_rank_local == \
1021 + ReconfigureRankType.KEEP_CURRENT_RANK
1022 + parallel_config.data_parallel_master_ip = \
1023 + reconfig_request.new_data_parallel_master_ip
1024 + parallel_config.data_parallel_master_port = \
1025 + reconfig_request.new_data_parallel_master_port
1026 + if reconfig_request.new_data_parallel_rank != -2:
1027 + self.dp_rank = parallel_config.data_parallel_rank
1028 + self.dp_group = parallel_config.stateless_init_dp_group()
1029 + reconfig_request.new_data_parallel_master_port = \
1030 + parallel_config.data_parallel_master_port
1031 +
1032 + self.model_executor.reinitialize_distributed(reconfig_request)
1033 + if reconfig_request.new_data_parallel_size > old_dp_size:
1034 + assert self.available_gpu_memory_for_kv_cache > 0
1035 + # pass available_gpu_memory_for_kv_cache from existing
1036 + # engine-cores to new engine-cores so they can directly
1037 + # use it in _initialize_kv_caches() rather than profiling.
1038 + ParallelConfig.sync_kv_cache_memory_size(
1039 + self.dp_group, self.available_gpu_memory_for_kv_cache)
1040 + # NOTE(yongji): newly joined workers require dummy_run even
1041 + # CUDA graph is not used
1042 + self.model_executor.collective_rpc("compile_or_warm_up_model")
1043 + if reconfig_request.new_data_parallel_rank == \
1044 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
1045 + self.shutdown()
1046 + logger.info("DPEngineCoreProc %s shutdown", self.dp_rank)
1047 + else:
1048 + logger.info("Distributed environment reinitialized for DP rank %s",
1049 + self.dp_rank)
1050 +
1051
1052 class DPEngineCoreActor(DPEngineCoreProc):
1053 """
vllm/v1/engine/core_client.py
+182 -7
@@ -21,9 +21,11 @@ import zmq.asyncio
21 from vllm.config import VllmConfig
22 from vllm.logger import init_logger
23 from vllm.lora.request import LoRARequest
24 -from vllm.utils import get_open_zmq_inproc_path, make_zmq_socket
24 +from vllm.utils import get_open_port, get_open_zmq_inproc_path, make_zmq_socket
25 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
26 - EngineCoreRequestType, UtilityOutput)
26 + EngineCoreRequestType,
27 + ReconfigureDistributedRequest, ReconfigureRankType,
28 + UtilityOutput)
29 from vllm.v1.engine.coordinator import DPCoordinator
30 from vllm.v1.engine.core import EngineCore, EngineCoreProc
31 from vllm.v1.engine.exceptions import EngineDeadError
@@ -162,6 +164,9 @@ class EngineCoreClient(ABC):
164 running state."""
165 raise NotImplementedError
166
167 + async def scale_elastic_ep(self, new_data_parallel_size: int) -> None:
168 + raise NotImplementedError
169 +
170 async def get_output_async(self) -> EngineCoreOutputs:
171 raise NotImplementedError
172
@@ -910,14 +915,30 @@ class DPAsyncMPClient(AsyncMPClient):
915 events = await poller.poll()
916 if not self.engines_running and len(events) == 2 or (
917 events[0][0] == first_req_rcv_socket):
913 - # Send a message to notify the coordinator that
918 + # Check if this is a regular request notification or
919 + # scale up notification
920 + buf = first_req_rcv_socket.recv(
921 + flags=zmq.NOBLOCK).result()
922 +
923 + decoded = msgspec.msgpack.decode(buf)
924 + if isinstance(
925 + decoded,
926 + (list, tuple)) and len(decoded) == 2 and decoded[
927 + 0] == "SCALE_ELASTIC_EP":
928 + # Extract new engine count from the decoded message
929 + new_engine_count = decoded[1]
930 + # Send scale up notification to coordinator
931 + scale_msg = msgspec.msgpack.encode(
932 + ("SCALE_ELASTIC_EP", new_engine_count))
933 + await socket.send(scale_msg)
934 + continue
935 +
936 # we're sending a request while the engines are
937 # paused, so that it can wake the others up
938 # (to run dummy EP loop).
939 + assert decoded[0] == "FIRST_REQ"
940 + target_eng_index = decoded[1]
941 self.engines_running = True
918 - buf = first_req_rcv_socket.recv(
919 - flags=zmq.NOBLOCK).result()
920 - target_eng_index = int.from_bytes(buf, "little")
942 msg = msgspec.msgpack.encode(
943 (target_eng_index, self.current_wave))
944 await socket.send(msg)
@@ -953,7 +974,8 @@ class DPAsyncMPClient(AsyncMPClient):
974 chosen_engine)
975 if not self.engines_running:
976 # Notify coordinator that we're sending a request
956 - await self.first_req_send_socket.send(chosen_engine)
977 + req_msg = msgspec.msgpack.encode(("FIRST_REQ", chosen_engine))
978 + await self.first_req_send_socket.send(req_msg)
979
980 await to_await
981
@@ -1047,3 +1069,156 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1069 engine: EngineIdentity) -> None:
1070 await self._send_input(EngineCoreRequestType.ABORT, request_ids,
1071 engine)
1072 +
1073 + async def _send_reconfig_message(
1074 + self, reconfig_request: ReconfigureDistributedRequest,
1075 + engine: EngineIdentity) -> asyncio.Future:
1076 + """Send reconfiguration message and return the result future without
1077 + waiting for completion."""
1078 + call_id = uuid.uuid1().int >> 64
1079 + future = asyncio.get_running_loop().create_future()
1080 + self.utility_results[call_id] = future
1081 + message = (EngineCoreRequestType.UTILITY.value, *self.encoder.encode(
1082 + (self.client_index, call_id, "reinitialize_distributed",
1083 + (reconfig_request, ))))
1084 + await self._send_input_message(message, engine, reconfig_request)
1085 + self._ensure_output_queue_task()
1086 + return future
1087 +
1088 + async def scale_elastic_ep(self, new_data_parallel_size: int) -> None:
1089 + """Scale elastic EP data parallel size"""
1090 + cur_data_parallel_size = len(self.core_engines)
1091 +
1092 + assert new_data_parallel_size != cur_data_parallel_size, (
1093 + f"new_data_parallel_size {new_data_parallel_size} must be "
1094 + f"different from cur_data_parallel_size {cur_data_parallel_size}")
1095 +
1096 + assert self.vllm_config.parallel_config.data_parallel_backend == \
1097 + "ray", ("Only ray DP backend supports scaling elastic EP")
1098 +
1099 + scale_up = new_data_parallel_size > cur_data_parallel_size
1100 +
1101 + if scale_up:
1102 + await self._scale_up_elastic_ep(cur_data_parallel_size,
1103 + new_data_parallel_size)
1104 + else:
1105 + await self._scale_down_elastic_ep(cur_data_parallel_size,
1106 + new_data_parallel_size)
1107 +
1108 + async def _scale_up_elastic_ep(self, cur_data_parallel_size: int,
1109 + new_data_parallel_size: int) -> None:
1110 + """Scale up the data parallel size by creating new engine cores
1111 + and reconfiguring existing ones."""
1112 + cur_data_parallel_size = len(self.core_engines)
1113 +
1114 + # Phase 1: Send reconfigure messages to all existing engines and wait
1115 + # for them to be sent
1116 + reconfig_futures = []
1117 + self.vllm_config.parallel_config.data_parallel_master_port = \
1118 + get_open_port()
1119 + for engine in self.core_engines:
1120 + reconfig_request = ReconfigureDistributedRequest(
1121 + new_data_parallel_size=new_data_parallel_size,
1122 + new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1123 + new_data_parallel_rank_local=\
1124 + ReconfigureRankType.KEEP_CURRENT_RANK,
1125 + new_data_parallel_master_ip=self.vllm_config.parallel_config.
1126 + data_parallel_master_ip,
1127 + new_data_parallel_master_port=self.vllm_config.parallel_config.
1128 + data_parallel_master_port)
1129 + future = await self._send_reconfig_message(reconfig_request,
1130 + engine)
1131 + reconfig_futures.append(future)
1132 +
1133 + logger.info("All reconfigure messages sent, starting engine creation")
1134 +
1135 + # Phase 2: Create new engines now that reconfig messages have been sent
1136 + # self.resources.engine_manager is guaranteed to be
1137 + # CoreEngineActorManager for RayDPClient
1138 + assert isinstance(self.resources.engine_manager,
1139 + CoreEngineActorManager)
1140 + self.resources.engine_manager.scale_up_elastic_ep(
1141 + self.vllm_config, new_data_parallel_size)
1142 +
1143 + # Create new CoreEngine objects for the new engines
1144 + new_engine_identities = set()
1145 + for i in range(cur_data_parallel_size, new_data_parallel_size):
1146 + new_engine = i.to_bytes(2, "little")
1147 + self.core_engines.append(new_engine)
1148 + new_engine_identities.add(new_engine)
1149 +
1150 + # Wait for ready messages from new engines on the input socket
1151 + sync_input_socket = zmq.Socket.shadow(self.input_socket)
1152 + while new_engine_identities:
1153 + if not sync_input_socket.poll(timeout=600_000):
1154 + raise TimeoutError(
1155 + "Timed out waiting for new engines to send initial "
1156 + "message on input socket.")
1157 + identity, _ = sync_input_socket.recv_multipart()
1158 + new_engine_identities.discard(identity)
1159 +
1160 + # Phase 3: Wait for all existing engines to complete reconfiguration
1161 + logger.info("Waiting for existing engines to complete reconfiguration")
1162 + await asyncio.gather(*reconfig_futures)
1163 +
1164 + # Notify coordinator about scale up through existing
1165 + # stats_update_task connection
1166 + self._ensure_stats_update_task()
1167 + scale_up_marker = msgspec.msgpack.encode(
1168 + ("SCALE_ELASTIC_EP", new_data_parallel_size))
1169 + await self.first_req_send_socket.send(scale_up_marker)
1170 +
1171 + # Update the parallel config
1172 + self.vllm_config.parallel_config.data_parallel_size = \
1173 + new_data_parallel_size
1174 + logger.info(
1175 + "[Elastic EP] Scale up completed, new data parallel size: %s",
1176 + new_data_parallel_size)
1177 +
1178 + async def _scale_down_elastic_ep(self, cur_data_parallel_size: int,
1179 + new_data_parallel_size: int) -> None:
1180 + """Scale down the data parallel size by shutting down and
1181 + reconfiguring existing engine cores."""
1182 + cur_data_parallel_size = len(self.core_engines)
1183 +
1184 + self.vllm_config.parallel_config.data_parallel_master_port = \
1185 + get_open_port()
1186 +
1187 + reconfig_futures = []
1188 + for cur_dp_rank, engine in enumerate(self.core_engines):
1189 + reconfig_request = ReconfigureDistributedRequest(
1190 + new_data_parallel_size=new_data_parallel_size,
1191 + new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1192 + new_data_parallel_rank_local=\
1193 + ReconfigureRankType.KEEP_CURRENT_RANK,
1194 + new_data_parallel_master_ip=self.vllm_config.parallel_config.
1195 + data_parallel_master_ip,
1196 + new_data_parallel_master_port=self.vllm_config.parallel_config.
1197 + data_parallel_master_port)
1198 + if cur_dp_rank >= new_data_parallel_size:
1199 + reconfig_request.new_data_parallel_rank = \
1200 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK
1201 + future = await self._send_reconfig_message(reconfig_request,
1202 + engine)
1203 + reconfig_futures.append(future)
1204 +
1205 + for _ in range(new_data_parallel_size, cur_data_parallel_size):
1206 + self.core_engines.pop()
1207 +
1208 + await asyncio.gather(*reconfig_futures)
1209 +
1210 + assert isinstance(self.resources.engine_manager,
1211 + CoreEngineActorManager)
1212 + self.resources.engine_manager.scale_down_elastic_ep(
1213 + cur_data_parallel_size, new_data_parallel_size)
1214 +
1215 + self._ensure_stats_update_task()
1216 + scale_down_marker = msgspec.msgpack.encode(
1217 + ("SCALE_ELASTIC_EP", new_data_parallel_size))
1218 + await self.first_req_send_socket.send(scale_down_marker)
1219 +
1220 + self.vllm_config.parallel_config.data_parallel_size = \
1221 + new_data_parallel_size
1222 + logger.info(
1223 + "[Elastic EP] Scale down completed, new data parallel size: %s",
1224 + new_data_parallel_size)
vllm/v1/engine/utils.py
+217 -8
@@ -174,16 +174,21 @@ class CoreEngineActorManager:
174
175 self.local_engine_actors: list[ray.ActorHandle] = []
176 self.remote_engine_actors: list[ray.ActorHandle] = []
177 +
178 + env_vars_list = get_env_vars_to_copy(destination="DPEngineCoreActor")
179 + self.env_vars_dict = {
180 + name: os.environ[name]
181 + for name in env_vars_list if name in os.environ
182 + }
183 + runtime_env = RuntimeEnv(env_vars=self.env_vars_dict)
184 +
185 + self.addresses = addresses
186 + self.executor_class = executor_class
187 + self.log_stats = log_stats
188 dp_size = vllm_config.parallel_config.data_parallel_size
189 local_engine_count = \
190 vllm_config.parallel_config.data_parallel_size_local
191 world_size = vllm_config.parallel_config.world_size
181 - env_vars_set = get_env_vars_to_copy(destination="DPEngineCoreActor")
182 - env_vars_dict = {
183 - name: os.environ[name]
184 - for name in env_vars_set if name in os.environ
185 - }
186 - runtime_env = RuntimeEnv(env_vars=env_vars_dict)
192
193 if ray.is_initialized():
194 logger.info(
@@ -208,6 +213,7 @@ class CoreEngineActorManager:
213 assert len(placement_groups) == dp_size, (
214 "Number of placement groups must match data parallel size")
215
216 + self.placement_group_is_local = []
217 refs = []
218 for index in range(dp_size):
219 local_index = local_dp_ranks[index]
@@ -231,6 +237,7 @@ class CoreEngineActorManager:
237 self.local_engine_actors.append(actor)
238 else:
239 self.remote_engine_actors.append(actor)
240 + self.placement_group_is_local.append(local_client)
241 refs.append(actor.wait_for_init.remote())
242
243 ray.get(refs)
@@ -242,6 +249,9 @@ class CoreEngineActorManager:
249 def create_dp_placement_groups(
250 vllm_config: VllmConfig
251 ) -> tuple[list["PlacementGroup"], list[int]]:
252 + """
253 + Create placement groups for data parallel.
254 + """
255
256 import ray
257 from ray._private.state import available_resources_per_node
@@ -250,10 +260,11 @@ class CoreEngineActorManager:
260 logger.info("Creating placement groups for data parallel")
261 dp_master_ip = \
262 vllm_config.parallel_config.data_parallel_master_ip
253 - dp_size = vllm_config.parallel_config.data_parallel_size
263 + num_pg_to_create = vllm_config.parallel_config.data_parallel_size
264 local_engine_count = \
265 vllm_config.parallel_config.data_parallel_size_local
266
267 + nodes = list_nodes()
268 nodes = sorted(list_nodes(),
269 key=lambda node: node.node_ip != dp_master_ip)
270 assert nodes[0].node_ip == dp_master_ip, (
@@ -293,7 +304,7 @@ class CoreEngineActorManager:
304 local_dp_ranks.append(i)
305 else:
306 for i in range(available_engine_count):
296 - if len(placement_groups) == dp_size:
307 + if len(placement_groups) == num_pg_to_create:
308 break
309 bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
310 pg = ray.util.placement_group(
@@ -305,6 +316,204 @@ class CoreEngineActorManager:
316 local_dp_ranks.append(i)
317 return placement_groups, local_dp_ranks
318
319 + @staticmethod
320 + def add_dp_placement_groups(
321 + old_vllm_config: VllmConfig, new_data_parallel_size: int
322 + ) -> tuple[list["PlacementGroup"], list[int]]:
323 + """
324 + Add placement groups for new data parallel size.
325 + """
326 + import ray
327 + from ray._private.state import (available_resources_per_node,
328 + total_resources_per_node)
329 + from ray.util.state import list_nodes
330 +
331 + old_dp_size = old_vllm_config.parallel_config.data_parallel_size
332 + num_pg_to_create = new_data_parallel_size - old_dp_size
333 +
334 + if num_pg_to_create <= 0:
335 + return [], []
336 +
337 + dp_master_ip = old_vllm_config.parallel_config.data_parallel_master_ip
338 + world_size = old_vllm_config.parallel_config.world_size
339 +
340 + nodes = list_nodes()
341 + nodes = sorted(nodes, key=lambda node: node.node_ip != dp_master_ip)
342 + assert nodes[0].node_ip == dp_master_ip, (
343 + "The first node must be the head node")
344 + assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, (
345 + "There can only be one head node")
346 +
347 + available_resources = available_resources_per_node()
348 + total_resources = total_resources_per_node()
349 +
350 + placement_groups = []
351 + local_dp_ranks = []
352 + num_pg_created = 0
353 +
354 + for node in nodes:
355 + if num_pg_created >= num_pg_to_create:
356 + break
357 +
358 + node_ip = node.node_ip
359 + node_id = node.node_id
360 + available_gpus = int(available_resources[node_id]["GPU"])
361 +
362 + # Get total GPUs on this node from the node's resources
363 + # Ray stores node resources with node ID as key
364 + total_gpus = int(total_resources[node_id]["GPU"])
365 +
366 + # Calculate used GPUs and used engines on this node
367 + used_gpus = max(0, total_gpus - available_gpus)
368 + used_engines_on_node = used_gpus // world_size
369 +
370 + # Calculate how many new engines this node can accommodate
371 + available_engine_count = available_gpus // world_size
372 +
373 + # Create placement groups for new engines on this node
374 + for i in range(available_engine_count):
375 + if num_pg_created >= num_pg_to_create:
376 + break
377 +
378 + rank = old_dp_size + num_pg_created
379 +
380 + # Create bundles with node constraint for master node
381 + if node_ip == dp_master_ip:
382 + bundles = [{
383 + "GPU": 1.0,
384 + "node:" + dp_master_ip: 0.001
385 + }] * world_size + [{
386 + "CPU": 1.0
387 + }]
388 + else:
389 + bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
390 +
391 + pg = ray.util.placement_group(
392 + name=f"dp_rank_{rank}",
393 + strategy="STRICT_PACK",
394 + bundles=bundles,
395 + )
396 + placement_groups.append(pg)
397 +
398 + # Local rank starts from the number of engines already used
399 + # on this node
400 + local_rank = used_engines_on_node + i
401 + local_dp_ranks.append(local_rank)
402 + num_pg_created += 1
403 +
404 + return placement_groups, local_dp_ranks
405 +
406 + def scale_up_elastic_ep(self, cur_vllm_config: VllmConfig,
407 + new_data_parallel_size: int) -> None:
408 + import copy
409 +
410 + import ray
411 + from ray.runtime_env import RuntimeEnv
412 + from ray.util.scheduling_strategies import (
413 + PlacementGroupSchedulingStrategy)
414 +
415 + from vllm.v1.engine.core import DPEngineCoreActor
416 +
417 + cur_data_parallel_size = len(self.local_engine_actors) + \
418 + len(self.remote_engine_actors)
419 +
420 + assert new_data_parallel_size > cur_data_parallel_size, (
421 + f"New data parallel size {new_data_parallel_size} must be greater "
422 + f"than current data parallel size {cur_data_parallel_size} "
423 + "for scale up")
424 +
425 + placement_groups, local_dp_ranks = \
426 + self.add_dp_placement_groups(
427 + cur_vllm_config, new_data_parallel_size)
428 +
429 + world_size = cur_vllm_config.parallel_config.world_size
430 + dp_master_ip = cur_vllm_config.parallel_config.data_parallel_master_ip
431 + new_local_engines = 0
432 +
433 + runtime_env = RuntimeEnv(env_vars=self.env_vars_dict
434 + | {"VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": "1"})
435 + for i, (pg,
436 + local_rank) in enumerate(zip(placement_groups,
437 + local_dp_ranks)):
438 + rank = cur_data_parallel_size + i
439 + dp_vllm_config = copy.deepcopy(cur_vllm_config)
440 + dp_vllm_config.parallel_config.data_parallel_size = \
441 + new_data_parallel_size
442 + dp_vllm_config.parallel_config.placement_group = pg
443 +
444 + # Check if this placement group is on the head node
445 + local_client = any(
446 + bundle.get("node:" + dp_master_ip, 0) > 0
447 + for bundle in pg.bundle_specs)
448 +
449 + if local_client:
450 + new_local_engines += 1
451 + # Update data_parallel_size_local
452 + dp_vllm_config.parallel_config.data_parallel_size_local = (
453 + cur_vllm_config.parallel_config.data_parallel_size_local +
454 + new_local_engines)
455 +
456 + actor = ray.remote(DPEngineCoreActor).options(
457 + scheduling_strategy=PlacementGroupSchedulingStrategy(
458 + placement_group=pg,
459 + placement_group_bundle_index=world_size,
460 + ),
461 + runtime_env=runtime_env).remote(
462 + vllm_config=dp_vllm_config,
463 + executor_class=self.executor_class,
464 + log_stats=self.log_stats,
465 + local_client=local_client,
466 + addresses=self.addresses,
467 + dp_rank=rank,
468 + local_dp_rank=local_rank)
469 +
470 + if local_client:
471 + self.local_engine_actors.append(actor)
472 + else:
473 + self.remote_engine_actors.append(actor)
474 + self.created_placement_groups.append(pg)
475 + self.placement_group_is_local.append(local_client)
476 +
477 + ray.get([
478 + actor.wait_for_init.remote()
479 + for actor in (self.local_engine_actors[-new_local_engines:]
480 + if new_local_engines > 0 else []) +
481 + self.remote_engine_actors[-(len(placement_groups) -
482 + new_local_engines):]
483 + ])
484 +
485 + actors = (self.local_engine_actors[-new_local_engines:]
486 + if new_local_engines > 0 else []) + \
487 + self.remote_engine_actors[-(len(placement_groups) -
488 + new_local_engines):]
489 +
490 + for actor in actors:
491 + self.run_refs.append(actor.run.remote())
492 +
493 + cur_vllm_config.parallel_config.data_parallel_size = \
494 + new_data_parallel_size
495 + # Update old_vllm_config with new data_parallel_size_local if any new
496 + # local engines were added
497 + if new_local_engines > 0:
498 + cur_vllm_config.parallel_config.data_parallel_size_local += \
499 + new_local_engines
500 +
501 + def scale_down_elastic_ep(self, cur_data_parallel_size: int,
502 + new_data_parallel_size: int) -> None:
503 + import ray
504 + assert cur_data_parallel_size > new_data_parallel_size, (
505 + f"cur_data_parallel_size {cur_data_parallel_size} must be greater "
506 + f"than new_data_parallel_size {new_data_parallel_size} "
507 + "for scale down")
508 + for _ in range(cur_data_parallel_size - new_data_parallel_size):
509 + pg = self.created_placement_groups.pop()
510 + is_local = self.placement_group_is_local.pop()
511 + if is_local:
512 + self.local_engine_actors.pop()
513 + else:
514 + self.remote_engine_actors.pop()
515 + ray.util.remove_placement_group(pg)
516 +
517 def get_run_refs(self):
518 return self.run_refs
519
vllm/v1/executor/ray_distributed_executor.py
+9
@@ -6,6 +6,7 @@ from typing import Union
6
7 from vllm.executor.ray_distributed_executor import ( # noqa
8 RayDistributedExecutor as RayDistributedExecutorV0)
9 +from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
10 from vllm.v1.executor.abstract import Executor
11 from vllm.v1.outputs import ModelRunnerOutput
12
@@ -62,3 +63,11 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
63 # When PP is used, we return a FutureWrapper immediately so that
64 # the scheduler can yield to the next batch.
65 return FutureWrapper(refs[0])
66 +
67 + def reinitialize_distributed(
68 + self, reconfig_request: ReconfigureDistributedRequest) -> None:
69 + self._run_workers("reinitialize_distributed", reconfig_request)
70 + if reconfig_request.new_data_parallel_rank == \
71 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
72 + self.shutdown()
73 + return
vllm/v1/worker/cpu_model_runner.py
+1 -1
@@ -49,7 +49,7 @@ class CPUModelRunner(GPUModelRunner):
49 if k.endswith("_cpu") and isinstance(v, torch.Tensor):
50 replace_tensor(self.input_batch.block_table, k, k[:-4])
51
52 - def load_model(self) -> None:
52 + def load_model(self, eep_scale_up: bool = False) -> None:
53 logger.info("Starting to load model %s...", self.model_config.model)
54 self.model = get_model(vllm_config=self.vllm_config)
55
vllm/v1/worker/gpu_model_runner.py
+36 -1
@@ -1745,8 +1745,40 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1745 new_config = update_config(config, config_overrides)
1746 setattr(self, config_name, new_config)
1747
1748 - def load_model(self) -> None:
1748 + def load_model(self, eep_scale_up: bool = False) -> None:
1749 + """
1750 + Args:
1751 + eep_scale_up: the model loading is for elastic EP scale up.
1752 + """
1753 logger.info("Starting to load model %s...", self.model_config.model)
1754 + if eep_scale_up:
1755 + from vllm.distributed.parallel_state import get_ep_group
1756 + num_local_physical_experts = torch.empty(1,
1757 + dtype=torch.int32,
1758 + device="cpu")
1759 + torch.distributed.broadcast(num_local_physical_experts,
1760 + group=get_ep_group().cpu_group,
1761 + group_src=0)
1762 + num_local_physical_experts = int(num_local_physical_experts.item())
1763 + new_ep_size = get_ep_group().world_size
1764 + global_expert_load, old_global_expert_indices = (
1765 + EplbState.recv_state())
1766 + num_logical_experts = global_expert_load.shape[1]
1767 + self.parallel_config.num_redundant_experts = (
1768 + num_local_physical_experts * new_ep_size - num_logical_experts)
1769 + assert old_global_expert_indices.shape[
1770 + 1] % num_local_physical_experts == 0
1771 + old_ep_size = old_global_expert_indices.shape[
1772 + 1] // num_local_physical_experts
1773 + rank_mapping = {
1774 + old_ep_rank: old_ep_rank
1775 + for old_ep_rank in range(old_ep_size)
1776 + }
1777 + else:
1778 + global_expert_load = None
1779 + old_global_expert_indices = None
1780 + rank_mapping = None
1781 +
1782 with DeviceMemoryProfiler() as m: # noqa: SIM117
1783 time_before_load = time.perf_counter()
1784 model_loader = get_model_loader(self.load_config)
@@ -1788,6 +1820,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1820 self.model,
1821 self.device,
1822 self.parallel_config,
1823 + global_expert_load,
1824 + old_global_expert_indices,
1825 + rank_mapping,
1826 )
1827
1828 def save_tensorized_model(
vllm/v1/worker/gpu_worker.py
+158 -1
@@ -26,6 +26,7 @@ from vllm.platforms import current_platform
26 from vllm.pooling_params import PoolingTask
27 from vllm.sequence import IntermediateTensors
28 from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling
29 +from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
30 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
31 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput
32 from vllm.v1.utils import report_usage_stats
@@ -191,8 +192,9 @@ class Worker(WorkerBase):
192 else:
193 from contextlib import nullcontext
194 context = nullcontext()
195 + eep_scale_up = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"
196 with context:
195 - self.model_runner.load_model()
197 + self.model_runner.load_model(eep_scale_up=eep_scale_up)
198
199 def update_config(self, overrides: dict[str, Any]) -> None:
200 self.model_runner.update_config(overrides)
@@ -384,6 +386,161 @@ class Worker(WorkerBase):
386 # worker will always be healthy as long as it's running.
387 return
388
389 + def _eplb_before_scale_down(self, old_ep_size: int,
390 + new_ep_size: int) -> None:
391 + from vllm.distributed.parallel_state import get_ep_group
392 + if get_ep_group().rank == 0:
393 + logger.info("[Elastic EP] Starting expert resharding "
394 + "before scaling down...")
395 + rank_mapping = {
396 + old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
397 + for old_ep_rank in range(old_ep_size)
398 + }
399 + assert self.model_runner.eplb_state is not None
400 + self.model_runner.eplb_state.rearrange(self.model_runner.model,
401 + execute_shuffle=True,
402 + global_expert_load=None,
403 + rank_mapping=rank_mapping)
404 + torch.cuda.synchronize()
405 + if get_ep_group().rank == 0:
406 + logger.info("[Elastic EP] Expert resharding completed!")
407 +
408 + def _eplb_after_scale_up(
409 + self, old_ep_size: int, new_ep_size: int,
410 + global_expert_load: Optional[torch.Tensor]) -> None:
411 + from vllm.distributed.parallel_state import get_ep_group
412 + if get_ep_group().rank == 0:
413 + logger.info("[Elastic EP] Starting expert resharding "
414 + "after scaling up...")
415 + rank_mapping = {
416 + old_ep_rank: old_ep_rank
417 + for old_ep_rank in range(old_ep_size)
418 + }
419 + assert self.model_runner.eplb_state is not None
420 + self.model_runner.eplb_state.rearrange(
421 + self.model_runner.model,
422 + execute_shuffle=True,
423 + global_expert_load=global_expert_load,
424 + rank_mapping=rank_mapping)
425 + if get_ep_group().rank == 0:
426 + logger.info("[Elastic EP] Expert resharding completed!")
427 +
428 + def _reconfigure_parallel_config(
429 + self, reconfig_request: ReconfigureDistributedRequest) -> None:
430 + """
431 + Update parallel config with provided reconfig_request
432 + """
433 + parallel_config = self.vllm_config.parallel_config
434 + parallel_config.data_parallel_size = \
435 + reconfig_request.new_data_parallel_size
436 + if reconfig_request.new_data_parallel_rank != \
437 + ReconfigureRankType.KEEP_CURRENT_RANK:
438 + parallel_config.data_parallel_rank = \
439 + reconfig_request.new_data_parallel_rank
440 + if reconfig_request.new_data_parallel_rank_local != \
441 + ReconfigureRankType.KEEP_CURRENT_RANK:
442 + parallel_config.data_parallel_rank_local = \
443 + reconfig_request.new_data_parallel_rank_local
444 + parallel_config.data_parallel_master_ip = \
445 + reconfig_request.new_data_parallel_master_ip
446 + parallel_config.data_parallel_master_port = \
447 + reconfig_request.new_data_parallel_master_port
448 +
449 + def _reconfigure_moe(self, old_ep_size: int,
450 + new_ep_size: int) -> Optional[torch.Tensor]:
451 + """
452 + Reconfigure MoE modules with provided reconfig_request
453 +
454 + Return the global expert load if new_ep_size > old_ep_size,
455 + otherwise None
456 + """
457 + from vllm.distributed.parallel_state import (
458 + get_dp_group, get_ep_group, prepare_communication_buffer_for_model)
459 + from vllm.model_executor.layers.fused_moe.layer import (
460 + FusedMoEParallelConfig)
461 +
462 + parallel_config = self.vllm_config.parallel_config
463 + moe_modules = [
464 + module for module in self.model_runner.model.modules()
465 + if module.__class__.__name__ == "FusedMoE"
466 + ]
467 + num_local_experts = moe_modules[0].moe_config.num_local_experts
468 + assert all(module.moe_config.num_local_experts == num_local_experts
469 + for module in moe_modules), (
470 + "All MoE modules must have the same number of experts")
471 + for module in moe_modules:
472 + module.moe_config.num_experts = num_local_experts * new_ep_size
473 + module.global_num_experts = module.moe_config.num_experts
474 + module.moe_parallel_config = FusedMoEParallelConfig.make(
475 + tp_size_=get_tp_group().world_size,
476 + dp_size_=get_dp_group().world_size,
477 + vllm_parallel_config=parallel_config,
478 + )
479 + module.moe_config.moe_parallel_config = module.moe_parallel_config
480 + if new_ep_size < old_ep_size:
481 + num_local_physical_experts = num_local_experts
482 + assert self.model_runner.eplb_state is not None
483 + new_physical_experts = \
484 + self.model_runner.eplb_state.physical_to_logical_map.shape[1]
485 + parallel_config.num_redundant_experts = (
486 + new_physical_experts -
487 + self.model_runner.eplb_state.logical_replica_count.shape[1])
488 + global_expert_load = None
489 + else:
490 + num_local_physical_experts = torch.tensor([num_local_experts],
491 + dtype=torch.int32,
492 + device="cpu")
493 + torch.distributed.broadcast(num_local_physical_experts,
494 + group=get_ep_group().cpu_group,
495 + group_src=0)
496 + num_local_physical_experts = num_local_physical_experts.item()
497 + new_physical_experts = num_local_physical_experts * new_ep_size
498 + assert self.model_runner.eplb_state is not None
499 + global_expert_load = self.model_runner.eplb_state.rearrange(
500 + self.model_runner.model, execute_shuffle=False)
501 + parallel_config.num_redundant_experts = (
502 + new_physical_experts - global_expert_load.shape[1])
503 + prepare_communication_buffer_for_model(self.model_runner.model)
504 + self.model_runner.model.update_physical_experts_metadata(
505 + num_physical_experts=new_physical_experts,
506 + num_local_physical_experts=num_local_physical_experts)
507 + return global_expert_load
508 +
509 + def reinitialize_distributed(
510 + self, reconfig_request: ReconfigureDistributedRequest) -> None:
511 + from vllm.config import set_current_vllm_config
512 + from vllm.distributed.parallel_state import (
513 + cleanup_dist_env_and_memory, get_ep_group)
514 +
515 + old_ep_size = get_ep_group().world_size
516 + old_ep_rank = get_ep_group().rank
517 + new_ep_size = reconfig_request.new_data_parallel_size * get_tp_group(
518 + ).world_size * get_pp_group().world_size
519 + if new_ep_size < old_ep_size:
520 + self._eplb_before_scale_down(old_ep_size, new_ep_size)
521 +
522 + cleanup_dist_env_and_memory()
523 +
524 + if reconfig_request.new_data_parallel_rank == \
525 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
526 + assert old_ep_rank >= new_ep_size
527 + # shutdown
528 + return
529 +
530 + self._reconfigure_parallel_config(reconfig_request)
531 +
532 + with set_current_vllm_config(self.vllm_config):
533 + init_worker_distributed_environment(self.vllm_config, self.rank,
534 + self.distributed_init_method,
535 + self.local_rank)
536 +
537 + global_expert_load = self._reconfigure_moe(old_ep_size, new_ep_size)
538 +
539 + if new_ep_size > old_ep_size:
540 + assert global_expert_load is not None
541 + self._eplb_after_scale_up(old_ep_size, new_ep_size,
542 + global_expert_load)
543 +
544 def save_sharded_state(
545 self,
546 path: str,