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 14, 05:30 UTC
[Misc] Further refine type annotations in parallel state (#22499)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Aug 10, 2025 at 20:49 UTC
d411df029648ff8107bddf89594b101879960491
2 files changed
+19
-20
vllm/distributed/eplb/eplb_state.py
-3
@@ -259,7 +259,6 @@ class EplbState:
259
260
if global_expert_load is not None:
261
ep_group = get_ep_group().device_group
262
- assert ep_group is not None
262
assert global_expert_load.shape == (model.num_moe_layers,
263
model.num_logical_experts)
264
assert global_expert_load.dtype == torch.int64
@@ -366,7 +365,6 @@ class EplbState:
365
366
# Collect load metrics from all ranks
367
ep_group = get_ep_group().device_group
369
- assert ep_group is not None
368
all_reduce(total_expert_load_pass, group=ep_group)
369
370
# num_tokens_per_rank: (num_moe_layers, num_ranks)
@@ -422,7 +420,6 @@ class EplbState:
420
"""
421
422
ep_group = get_ep_group().device_group
425
- assert ep_group is not None
423
ep_rank = ep_group.rank()
424
425
time_start = None
vllm/distributed/parallel_state.py
+19
-17
@@ -197,11 +197,10 @@ class GroupCoordinator:
197
# 3 | 1 | 3 | 1 | 3
198
local_rank: int # local rank used to assign devices
199
rank_in_group: int # rank inside the group
200
- cpu_group: Optional[ProcessGroup] # group for CPU communication
201
- device_group: Optional[ProcessGroup] # group for device communication
202
- use_device_communicator: bool # whether to use device communicator
203
- device_communicator: Optional[
204
- DeviceCommunicatorBase] # device communicator
200
+ cpu_group: ProcessGroup # group for CPU communication
201
+ device_group: ProcessGroup # group for device communication
202
+ # device communicator (if use_device_communicator=True)
203
+ device_communicator: Optional[DeviceCommunicatorBase]
204
mq_broadcaster: Optional[Any] # shared memory broadcaster
205
206
def __init__(
@@ -209,7 +208,7 @@ class GroupCoordinator:
208
group_ranks: list[list[int]],
209
local_rank: int,
210
torch_distributed_backend: Union[str, Backend],
212
- use_device_communicator: bool,
211
+ use_device_communicator: bool, # whether to use device communicator
212
use_message_queue_broadcaster: bool = False,
213
group_name: Optional[str] = None,
214
):
@@ -219,8 +218,9 @@ class GroupCoordinator:
218
219
self.rank = torch.distributed.get_rank()
220
self.local_rank = local_rank
222
- self.device_group = None
223
- self.cpu_group = None
221
+
222
+ self_device_group = None
223
+ self_cpu_group = None
224
225
for ranks in group_ranks:
226
device_group = torch.distributed.new_group(
@@ -232,11 +232,14 @@ class GroupCoordinator:
232
self.ranks = ranks
233
self.world_size = len(ranks)
234
self.rank_in_group = ranks.index(self.rank)
235
- self.device_group = device_group
236
- self.cpu_group = cpu_group
235
+ self_device_group = device_group
236
+ self_cpu_group = cpu_group
237
+
238
+ assert self_cpu_group is not None
239
+ assert self_device_group is not None
240
238
- assert self.cpu_group is not None
239
- assert self.device_group is not None
241
+ self.cpu_group = self_cpu_group
242
+ self.device_group = self_device_group
243
244
from vllm.platforms import current_platform
245
@@ -251,7 +254,6 @@ class GroupCoordinator:
254
self.device = torch.device("cpu")
255
256
self.use_device_communicator = use_device_communicator
254
-
257
self.device_communicator = None
258
if use_device_communicator and self.world_size > 1:
259
device_comm_cls = resolve_obj_by_qualname(
@@ -817,12 +819,12 @@ class GroupCoordinator:
819
return self.device_communicator.recv(size, dtype, src)
820
821
def destroy(self):
820
- if self.device_group is not None:
822
+ if hasattr(self, "device_group"):
823
torch.distributed.destroy_process_group(self.device_group)
822
- self.device_group = None
823
- if self.cpu_group is not None:
824
+ del self.device_group
825
+ if hasattr(self, "cpu_group"):
826
torch.distributed.destroy_process_group(self.cpu_group)
825
- self.cpu_group = None
827
+ del self.cpu_group
828
if self.device_communicator is not None:
829
self.device_communicator.destroy()
830
if self.mq_broadcaster is not None: