@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 74333ae2f6

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 14, 05:30 UTC

[Misc] correct static type check for GroupCoordinator (#21946)

Signed-off-by: Andy Xie <andy.xning@gmail.com>

Ning Xie committed Aug 5, 2025 at 18:17 UTC 74333ae2f6c3c4aa4b55301e5ed7aba03a5b09f8

3 files changed +29 -4

vllm/distributed/device_communicators/ray_communicator.py

@@ -70,6 +70,7 @@ class RayPPCommunicator(Communicator):

70 assert ray.get_gpu_ids(), "RayPPCommunicator has no GPUs assigned"

72 self._comm = get_pp_group().device_communicator

73 + assert self._comm is not None

75 # Since we wrap around the vLLM _PP communicator, we use

76 # the rank from the vLLM communicator, and ignore the rank

vllm/distributed/eplb/eplb_state.py

@@ -251,6 +251,7 @@ class EplbState:

251

252 if global_expert_load is not None:

253 ep_group = get_ep_group().device_group

254 + assert ep_group is not None

255 assert global_expert_load.shape == (model.num_moe_layers,

256 model.num_logical_experts)

257 assert global_expert_load.dtype == torch.int64

@@ -357,6 +358,7 @@ class EplbState:

358

359 # Collect load metrics from all ranks

360 ep_group = get_ep_group().device_group

361 + assert ep_group is not None

362 num_tokens_list = [

363 torch.empty_like(num_tokens) for _ in range(ep_group.size())

364 ]

@@ -412,6 +414,7 @@ class EplbState:

414 """

415

416 ep_group = get_ep_group().device_group

417 + assert ep_group is not None

418 ep_rank = ep_group.rank()

419

420 time_start = None

vllm/distributed/parallel_state.py

+25 -4

@@ -196,10 +196,11 @@ class GroupCoordinator:

196 # 3 | 1 | 3 | 1 | 3

197 local_rank: int # local rank used to assign devices

198 rank_in_group: int # rank inside the group

199 - cpu_group: ProcessGroup # group for CPU communication

200 - device_group: ProcessGroup # group for device communication

199 + cpu_group: Optional[ProcessGroup] # group for CPU communication

200 + device_group: Optional[ProcessGroup] # group for device communication

201 use_device_communicator: bool # whether to use device communicator

202 - device_communicator: DeviceCommunicatorBase # device communicator

202 + device_communicator: Optional[

203 + DeviceCommunicatorBase] # device communicator

204 mq_broadcaster: Optional[Any] # shared memory broadcaster

205

206 def __init__(

@@ -250,7 +251,7 @@ class GroupCoordinator:

251

252 self.use_device_communicator = use_device_communicator

253

253 - self.device_communicator: DeviceCommunicatorBase = None # type: ignore

254 + self.device_communicator = None

255 if use_device_communicator and self.world_size > 1:

256 device_comm_cls = resolve_obj_by_qualname(

257 current_platform.get_device_communicator_cls())

@@ -364,6 +365,8 @@ class GroupCoordinator:

365 return self._all_reduce_out_place(input_)

366

367 def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:

368 + if self.device_communicator is None:

369 + raise ValueError("No device communicator found")

370 return self.device_communicator.all_reduce(input_)

371

372 def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

@@ -384,12 +387,16 @@ class GroupCoordinator:

387

388 def _all_gather_out_place(self, input_: torch.Tensor,

389 dim: int) -> torch.Tensor:

390 + if self.device_communicator is None:

391 + raise ValueError("No device communicator found")

392 return self.device_communicator.all_gather(input_, dim)

393

394 def all_gatherv(self,

395 input_: Union[torch.Tensor, list[torch.Tensor]],

396 dim: int = 0,

397 sizes: Optional[list[int]] = None):

398 + if self.device_communicator is None:

399 + raise ValueError("No device communicator found")

400 return self.device_communicator.all_gatherv(input_, dim, sizes)

401

402 def reduce_scatter(self,

@@ -414,10 +421,14 @@ class GroupCoordinator:

421 input_: torch.Tensor,

422 dim: int = -1,

423 sizes: Optional[list[int]] = None) -> torch.Tensor:

424 + if self.device_communicator is None:

425 + raise ValueError("No device communicator found")

426 return self.device_communicator.reduce_scatterv(input_, dim, sizes)

427

428 def _reduce_scatter_out_place(self, input_: torch.Tensor,

429 dim: int) -> torch.Tensor:

430 + if self.device_communicator is None:

431 + raise ValueError("No device communicator found")

432 return self.device_communicator.reduce_scatter(input_, dim)

433

434 def gather(self,

@@ -433,6 +444,8 @@ class GroupCoordinator:

444 # Bypass the function if we are using only 1 GPU.

445 if world_size == 1:

446 return input_

447 + if self.device_communicator is None:

448 + raise ValueError("No device communicator found")

449 return self.device_communicator.gather(input_, dst, dim)

450

451 def broadcast(self, input_: torch.Tensor, src: int = 0):

@@ -667,6 +680,8 @@ class GroupCoordinator:

680 assert dst < self.world_size, f"Invalid dst rank ({dst})"

681

682 if self.use_cpu_custom_send_recv:

683 + if self.device_communicator is None:

684 + raise ValueError("No device communicator found")

685 self.device_communicator.send_tensor_dict( # type: ignore

686 tensor_dict, dst)

687 return None

@@ -727,6 +742,8 @@ class GroupCoordinator:

742 assert src < self.world_size, f"Invalid src rank ({src})"

743

744 if self.use_cpu_custom_send_recv:

745 + if self.device_communicator is None:

746 + raise ValueError("No device communicator found")

747 return self.device_communicator.recv_tensor_dict( # type: ignore

748 src)

749

@@ -784,6 +801,8 @@ class GroupCoordinator:

801 def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:

802 """Sends a tensor to the destination rank in a blocking way"""

803 """NOTE: `dst` is the local rank of the destination rank."""

804 + if self.device_communicator is None:

805 + raise ValueError("No device communicator found")

806 self.device_communicator.send(tensor, dst)

807

808 def recv(self,

@@ -792,6 +811,8 @@ class GroupCoordinator:

811 src: Optional[int] = None) -> torch.Tensor:

812 """Receives a tensor from the source rank."""

813 """NOTE: `src` is the local rank of the source rank."""

814 + if self.device_communicator is None:

815 + raise ValueError("No device communicator found")

816 return self.device_communicator.recv(size, dtype, src)

817

818 def destroy(self):