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

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo authored and GitHub committed 2026-01-28 22:12:52 +08:00
1 parent b11d436e19
commit c9392dff01
103 files changed
+1770 -1218

No files matched your search

+10 -18
View File
@@ -14,15 +14,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.utils.platform_utils import detect_device_type
device_type = detect_device_type()
if device_type == "cuda":
torch_device = torch.cuda
elif device_type == "npu":
torch_device = torch.npu
else:
raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
from vllm_omni.platforms import current_omni_platform
def update_environment_variables(envs_dict: dict[str, str]):
@@ -49,7 +41,7 @@ def test_4d_identity(
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -85,8 +77,8 @@ def _test_4d_identity_worker(
):
"""Worker function for test_4d_identity."""
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
@@ -183,7 +175,7 @@ def test_5d_identity(
):
"""Test that two consecutive all-to-all operations return the original input."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -219,8 +211,8 @@ def _test_5d_identity_worker(
):
"""Worker function for test_5d_identity."""
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set environment variables for distributed training
update_environment_variables(
@@ -316,7 +308,7 @@ def test_ring_p2p(
):
"""Test Ring P2P communication (send_recv)."""
# Skip if not enough GPUs available
available_gpus = torch_device.device_count()
available_gpus = current_omni_platform.get_device_count()
if available_gpus < world_size:
pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
@@ -339,8 +331,8 @@ def _test_ring_p2p_worker(
import sys
# Set device
device = torch.device(f"{device_type}:{local_rank}")
torch_device.set_device(device)
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
# Set env vars
# Use a different port to avoid conflict with other tests if run in parallel