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
[Platform] allow platform to init dp group (#22243)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
wangxiyuan committed
Oct 15, 2025 at 17:32 UTC
db1764e4e05b06c93073b9f26df7b1f3b684e638
5 files changed
+15
-83
vllm/config/parallel.py
+1
-1
@@ -334,7 +334,7 @@ class ParallelConfig:
334
self.get_next_dp_init_port(),
335
self.data_parallel_rank,
336
self.data_parallel_size,
337
- backend="gloo",
337
+ backend=current_platform.dist_backend,
338
)
339
except DistNetworkError as e:
340
# We only want to retry when the root cause is EADDRINUSE.
vllm/distributed/utils.py
+13
-13
@@ -415,7 +415,6 @@ class StatelessProcessGroup:
415
416
417
def init_gloo_process_group(
418
- backend: Backend,
418
prefix_store: PrefixStore,
419
group_rank: int,
420
group_size: int,
@@ -432,7 +431,7 @@ def init_gloo_process_group(
431
group_size,
432
)
433
else:
435
- options = ProcessGroup.Options(backend=backend)
434
+ options = ProcessGroup.Options(backend="gloo")
435
pg = ProcessGroup(
436
prefix_store,
437
group_rank,
@@ -504,24 +503,25 @@ def stateless_init_torch_distributed_process_group(
503
# Use a PrefixStore to avoid accidental overrides of keys used by
504
# different systems (e.g. RPC) in case the store is multi-tenant.
505
prefix_store = PrefixStore(init_method, store)
506
+ try:
507
+ from vllm.platforms import current_platform
508
508
- if backend == "gloo":
509
- return init_gloo_process_group(
509
+ return current_platform.stateless_init_device_torch_dist_pg(
510
backend=backend,
511
prefix_store=prefix_store,
512
group_rank=group_rank,
513
group_size=group_size,
514
timeout=timeout,
515
)
516
- from vllm.platforms import current_platform
517
-
518
- return current_platform.stateless_init_device_torch_dist_pg(
519
- backend=backend,
520
- prefix_store=prefix_store,
521
- group_rank=group_rank,
522
- group_size=group_size,
523
- timeout=timeout,
524
- )
516
+ except NotImplementedError:
517
+ # If platform doesn't implement stateless_init_device_torch_dist_pg, it
518
+ # will raise a NotImplementedError. In this case, we fall back to gloo.
519
+ return init_gloo_process_group(
520
+ prefix_store=prefix_store,
521
+ group_rank=group_rank,
522
+ group_size=group_size,
523
+ timeout=timeout,
524
+ )
525
526
527
def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None:
vllm/platforms/cuda.py
-34
@@ -6,13 +6,10 @@ pynvml. However, it should not initialize cuda context.
6
7
import os
8
from collections.abc import Callable
9
-from datetime import timedelta
9
from functools import cache, wraps
10
from typing import TYPE_CHECKING, TypeVar
11
12
import torch
14
-from torch.distributed import PrefixStore, ProcessGroup
15
-from torch.distributed.distributed_c10d import is_nccl_available
13
from typing_extensions import ParamSpec
14
15
# import custom ops, trigger op registration
@@ -455,37 +452,6 @@ class CudaPlatformBase(Platform):
452
def get_static_graph_wrapper_cls(cls) -> str:
453
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
454
458
- @classmethod
459
- def stateless_init_device_torch_dist_pg(
460
- cls,
461
- backend: str,
462
- prefix_store: PrefixStore,
463
- group_rank: int,
464
- group_size: int,
465
- timeout: timedelta,
466
- ) -> ProcessGroup:
467
- assert is_nccl_available()
468
- pg: ProcessGroup = ProcessGroup(
469
- prefix_store,
470
- group_rank,
471
- group_size,
472
- )
473
- from torch.distributed.distributed_c10d import ProcessGroupNCCL
474
-
475
- backend_options = ProcessGroupNCCL.Options()
476
- backend_options._timeout = timeout
477
-
478
- backend_class = ProcessGroupNCCL(
479
- prefix_store, group_rank, group_size, backend_options
480
- )
481
- backend_type = ProcessGroup.BackendType.NCCL
482
- device = torch.device("cuda")
483
- pg._set_default_backend(backend_type)
484
- backend_class._set_sequence_number_for_group()
485
-
486
- pg._register_backend(device, backend_type, backend_class)
487
- return pg
488
-
455
@classmethod
456
def device_count(cls) -> int:
457
return cuda_device_count_stateless()
vllm/platforms/interface.py
+1
-1
@@ -551,7 +551,7 @@ class Platform:
551
"""
552
Init platform-specific torch distributed process group.
553
"""
554
- raise RuntimeError(f"Unsupported torch distributed backend: {backend}")
554
+ raise NotImplementedError
555
556
@classmethod
557
def is_kv_cache_dtype_supported(
vllm/platforms/rocm.py
-34
@@ -2,13 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
-from datetime import timedelta
5
from functools import cache, lru_cache, wraps
6
from typing import TYPE_CHECKING
7
8
import torch
10
-from torch.distributed import PrefixStore, ProcessGroup
11
-from torch.distributed.distributed_c10d import is_nccl_available
9
10
import vllm.envs as envs
11
from vllm.logger import init_logger
@@ -476,37 +473,6 @@ class RocmPlatform(Platform):
473
def get_static_graph_wrapper_cls(cls) -> str:
474
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
475
479
- @classmethod
480
- def stateless_init_device_torch_dist_pg(
481
- cls,
482
- backend: str,
483
- prefix_store: PrefixStore,
484
- group_rank: int,
485
- group_size: int,
486
- timeout: timedelta,
487
- ) -> ProcessGroup:
488
- assert is_nccl_available()
489
- pg: ProcessGroup = ProcessGroup(
490
- prefix_store,
491
- group_rank,
492
- group_size,
493
- )
494
- from torch.distributed.distributed_c10d import ProcessGroupNCCL
495
-
496
- backend_options = ProcessGroupNCCL.Options()
497
- backend_options._timeout = timeout
498
-
499
- backend_class = ProcessGroupNCCL(
500
- prefix_store, group_rank, group_size, backend_options
501
- )
502
- backend_type = ProcessGroup.BackendType.NCCL
503
- device = torch.device("cuda")
504
- pg._set_default_backend(backend_type)
505
- backend_class._set_sequence_number_for_group()
506
-
507
- pg._register_backend(device, backend_type, backend_class)
508
- return pg
509
-
476
@classmethod
477
def device_count(cls) -> int:
478
return cuda_device_count_stateless()