[Hardware] Support platforms and plugin system (#774)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/diffusion/distributed/comm.py

@@ -8,6 +8,8 @@ import torch
 import torch.distributed as dist
 from torch import Tensor
 
+from vllm_omni.platforms import current_omni_platform
+
 __all__ = ["all_to_all_4D", "all_to_all_5D", "SeqAllToAll4D", "SeqAllToAll5D", "RingComm"]
 
 
@@ -48,7 +50,7 @@ def all_to_all_4D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
         # if scattering the seq-dim, transpose the heads back to the original dimension
@@ -83,7 +85,7 @@ def all_to_all_4D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 
@@ -155,7 +157,7 @@ def all_to_all_5D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 
@@ -190,7 +192,7 @@ def all_to_all_5D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -9,18 +9,14 @@ from typing import Any
 
 import torch
 import torch.distributed
-from torch.cuda import synchronize
 from torch.distributed import Backend, ProcessGroup
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-if envs._is_npu():
-    logger.info("torch.npu synchronize")
-    from torch.npu import synchronize
-
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
@@ -124,7 +120,7 @@ class GroupCoordinator:
         assert self.cpu_group is not None
         assert self.device_group is not None
 
-        self.device = envs.get_device(local_rank)
+        self.device = current_omni_platform.get_torch_device(local_rank)
 
     @property
     def first_rank(self):
@@ -618,7 +614,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
         assert self.cpu_group is not None
         assert self.device_group is not None
 
-        self.device = envs.get_device(local_rank)
+        self.device = current_omni_platform.get_torch_device(local_rank)
 
         self.recv_buffer_set: bool = False
         self.recv_tasks_queue: list[tuple[str, int]] = []
@@ -768,7 +764,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
 
         # To protect against race condition when using batch_isend_irecv().
         # should take this out once the bug with batch_isend_irecv is resolved.
-        synchronize()
+        current_omni_platform.synchronize()
 
         ops = []
         recv_prev_shape_tensor = None
@@ -799,7 +795,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
             for req in reqs:
                 req.wait()
 
-        synchronize()
+        current_omni_platform.synchronize()
 
         recv_prev_shape = [0, 0, 0]
         if recv_prev_shape_tensor is not None:

vllm_omni/diffusion/distributed/parallel_state.py

@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.platforms import current_omni_platform
 
 from .group_coordinator import (
     GroupCoordinator,
@@ -42,14 +43,6 @@ from .group_coordinator import (
     SequenceParallelGroupCoordinator,
 )
 
-if envs._is_npu():
-    from torch.npu import device_count, set_device
-elif envs._is_musa():
-    from torch_musa.core.device import device_count, set_device
-else:
-    from torch.cuda import device_count, set_device
-
-
 env_info = envs.PACKAGES_CHECKER.get_packages_info()
 
 HAS_FLASH_ATTN = env_info["has_flash_attn"]
@@ -396,7 +389,7 @@ def init_distributed_environment(
     backend: str | None = None,
 ):
     if backend is None:
-        backend = envs.get_torch_distributed_backend()
+        backend = current_omni_platform.dist_backend
     logger.debug(
         "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
         world_size,
@@ -416,7 +409,8 @@ def init_distributed_environment(
             world_size=world_size,
             rank=rank,
         )
-        set_device(torch.distributed.get_rank() % device_count())
+        device_id = torch.distributed.get_rank() % current_omni_platform.get_device_count()
+        current_omni_platform.set_device(current_omni_platform.get_torch_device(device_id))
     # set the local rank
     # local_rank is not available in torch ProcessGroup,
     # see https://github.com/pytorch/pytorch/issues/122816
@@ -572,7 +566,7 @@ def initialize_model_parallel(
     backend: str | None = None,
 ) -> None:
     if backend is None:
-        backend = envs.get_torch_distributed_backend()
+        backend = current_omni_platform.dist_backend
     """
     Initialize model parallel groups.
 
@@ -632,7 +626,7 @@ def initialize_model_parallel(
 
     # FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch,
     # the pipefusion is not ready for npu yet
-    if envs._is_npu():
+    if current_omni_platform.is_npu():
         assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU"
 
     dit_parallel_size = (

vllm_omni/diffusion/distributed/utils.py

@@ -5,11 +5,10 @@ import os
 
 import torch
 
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.platforms import current_omni_platform
 
 
 def get_local_device() -> torch.device:
     """Return the torch device for the current rank based on detected device type."""
-    device_type = detect_device_type()
-    local_rank = os.environ.get("LOCAL_RANK", 0)
-    return torch.device(f"{device_type}:{local_rank}")
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    return current_omni_platform.get_torch_device(local_rank)