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

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

tests/diffusion/attention/test_attention_sp.py

@@ -6,7 +6,6 @@ import tempfile
 
 import pytest
 import torch
-from vllm.platforms import current_platform
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
@@ -19,15 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-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]):
@@ -171,7 +162,7 @@ def test_sequence_parallel(
     sequence_parallel_size = ulysses_degree * ring_degree
 
     # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
     if available_gpus < sequence_parallel_size:
         pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
 
@@ -338,13 +329,13 @@ def ulysses_attention_on_test_model(
     """Run Ulysses attention test on a test model and save results for comparison."""
     # Use fixed seed for reproducibility across baseline and SP runs
     RANDOM_SEED = 42
-    current_platform.seed_everything(RANDOM_SEED)
+    current_omni_platform.seed_everything(RANDOM_SEED)
 
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
 
-    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)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
@@ -423,7 +414,7 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            torch_device.manual_seed_all(42)
+            current_omni_platform.seed_everything(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

tests/diffusion/distributed/test_comm.py

@@ -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

tests/diffusion/test_diffusion_worker.py

@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Unit tests for GPUDiffusionWorker class.
+Unit tests for DiffusionWorker class.
 
-This module tests the GPUDiffusionWorker implementation:
+This module tests the DiffusionWorker implementation:
 - load_weights: Loading model weights
 - sleep: Putting worker into sleep mode (levels 1 and 2)
 - wake_up: Waking worker from sleep mode
@@ -15,7 +15,7 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
 
 @pytest.fixture
@@ -33,18 +33,19 @@ def mock_od_config():
 
 @pytest.fixture
 def mock_gpu_worker(mock_od_config):
-    """Create a GPUDiffusionWorker with mocked initialization."""
-    with patch.object(GPUDiffusionWorker, "init_device"):
-        worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
+    """Create a DiffusionWorker with mocked initialization."""
+    with patch.object(DiffusionWorker, "init_device"):
+        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
         # Mock the model_runner with pipeline
         worker.model_runner = Mock()
         worker.model_runner.pipeline = Mock()
+        worker.device = torch.device("cuda", 0)
         worker._sleep_saved_buffers = {}
         return worker
 
 
-class TestGPUDiffusionWorkerLoadWeights:
-    """Test GPUDiffusionWorker.load_weights method."""
+class TestDiffusionWorkerLoadWeights:
+    """Test DiffusionWorker.load_weights method."""
 
     def test_load_weights_calls_pipeline(self, mock_gpu_worker):
         """Test that load_weights delegates to model_runner.load_weights."""
@@ -75,20 +76,21 @@ class TestGPUDiffusionWorkerLoadWeights:
         assert result == set()
 
 
-class TestGPUDiffusionWorkerSleep:
-    """Test GPUDiffusionWorker.sleep method."""
+class TestDiffusionWorkerSleep:
+    """Test DiffusionWorker.sleep method."""
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
         # Setup memory info mocks
-        # Before sleep: 1GB free, 8GB total
-        # After sleep: 3GB free, 8GB total (freed 2GB)
-        mock_mem_info.side_effect = [
-            (1 * 1024**3, 8 * 1024**3),  # Before sleep
-            (3 * 1024**3, 8 * 1024**3),  # After sleep
+        # Before sleep: 1GB free
+        # After sleep: 3GB free (freed 2GB)
+        mock_platform.get_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            3 * 1024**3,  # After sleep
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -104,15 +106,16 @@ class TestGPUDiffusionWorkerSleep:
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
         # Setup memory info mocks
-        mock_mem_info.side_effect = [
-            (1 * 1024**3, 8 * 1024**3),  # Before sleep
-            (5 * 1024**3, 8 * 1024**3),  # After sleep (freed 4GB)
+        mock_platform.get_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            5 * 1024**3,  # After sleep (freed 4GB)
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -141,15 +144,16 @@ class TestGPUDiffusionWorkerSleep:
         assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
         assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         # Simulate memory increase (should trigger assertion error)
-        mock_mem_info.side_effect = [
-            (3 * 1024**3, 8 * 1024**3),  # Before sleep: 3GB free
-            (1 * 1024**3, 8 * 1024**3),  # After sleep: 1GB free (negative freed!)
+        mock_platform.get_free_memory.side_effect = [
+            3 * 1024**3,  # Before sleep: 3GB free
+            1 * 1024**3,  # After sleep: 1GB free (negative freed!)
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         mock_allocator = Mock()
         mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
@@ -160,8 +164,8 @@ class TestGPUDiffusionWorkerSleep:
             mock_gpu_worker.sleep(level=1)
 
 
-class TestGPUDiffusionWorkerWakeUp:
-    """Test GPUDiffusionWorker.wake_up method."""
+class TestDiffusionWorkerWakeUp:
+    """Test DiffusionWorker.wake_up method."""
 
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
     def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -10,7 +10,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -40,7 +40,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -70,7 +70,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml

@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml

@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -7,7 +7,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.utils import GPUMemoryMonitor
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -20,7 +20,7 @@ models = ["riverclouds/qwen_image_random"]
 
 
 def inference(model_name: str, offload: bool = True):
-    torch.cuda.empty_cache()
+    current_omni_platform.empty_cache()
     device_index = torch.cuda.current_device()
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
@@ -45,7 +45,7 @@ def inference(model_name: str, offload: bool = True):
     return peak
 
 
-@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
+@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -9,7 +9,6 @@ from safetensors.torch import save_file
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -25,10 +24,6 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # model to reduce runtime and avoid extra downloads.
 models = ["Tongyi-MAI/Z-Image-Turbo"]
 
-# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly.
-if is_npu():
-    pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.", allow_module_level=True)
-
 
 @pytest.mark.parametrize("model_name", models)
 def test_diffusion_model(model_name: str, tmp_path: Path):

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -13,7 +13,7 @@ from vllm.assets.video import VideoAsset
 from vllm.envs import VLLM_USE_MODELSCOPE
 from vllm.multimodal.image import convert_image_mode
 
-from vllm_omni.utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 from .conftest import OmniRunner
 from .utils import create_new_process_for_each_test
@@ -21,9 +21,9 @@ from .utils import create_new_process_for_each_test
 models = ["Qwen/Qwen2.5-Omni-3B"]
 
 # CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
-if is_npu():
+if current_omni_platform.is_npu():
     stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif is_rocm():
+elif current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
 else:

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -14,14 +14,14 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
-from vllm_omni.utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 from .conftest import OmniRunner
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if is_rocm():
+if current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -30,8 +30,7 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.parallel_state import device_count
-from vllm_omni.diffusion.envs import get_device_name
+from vllm_omni.platforms import current_omni_platform
 
 models = ["riverclouds/qwen_image_random"]
 
@@ -63,9 +62,9 @@ def _cleanup_distributed():
 
     gc.collect()
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
 
     time.sleep(5)
 
@@ -88,7 +87,7 @@ def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height
                 width=width,
                 num_inference_steps=4,
                 guidance_scale=0.0,
-                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -125,7 +124,7 @@ def _run_sp(
                 width=width,
                 num_inference_steps=4,
                 guidance_scale=0.0,
-                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -157,8 +156,8 @@ def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
 def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
     """Test SP inference only (ulysses=2)."""
-    if device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+    if current_omni_platform.get_device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
 
     height = 256
     width = 256
@@ -177,8 +176,8 @@ def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
 def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
     """Test SP inference only (ring=2)."""
-    if device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+    if current_omni_platform.get_device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
 
     height = 256
     width = 256
@@ -212,8 +211,8 @@ def test_sequence_parallel(
         )
 
     sp_size = ulysses_degree * ring_degree
-    if device_count() < sp_size:
-        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
+    if current_omni_platform.get_device_count() < sp_size:
+        pytest.skip(f"Test requires {sp_size} GPUs but only {current_omni_platform.get_device_count()} available")
 
     height = 256
     width = 256
@@ -243,8 +242,10 @@ def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_ba
     ulysses_degree = 4
     ring_degree = 1
 
-    if device_count() < ulysses_degree * ring_degree:
-        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+    if current_omni_platform.get_device_count() < ulysses_degree * ring_degree:
+        pytest.skip(
+            f"Test requires {ulysses_degree * ring_degree} GPUs but only {current_omni_platform.get_device_count()} available"
+        )
 
     height = 272
     width = 272

tests/e2e/offline_inference/test_t2i_model.py

@@ -7,7 +7,7 @@ import torch
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -21,12 +21,11 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
 
-# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly
 # Modelscope can't find riverclouds/qwen_image_random
 # TODO: When NPU support is ready, remove this branch.
-if is_npu():
-    models = ["Qwen/Qwen-Image"]
-elif is_rocm():
+if current_omni_platform.is_npu():
+    models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
+elif current_omni_platform.is_rocm():
     # TODO: When ROCm support is ready, remove this branch.
     # vLLM V0.11.0 has issues running riverclouds/qwen_image_random
     # on ROCm

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -23,7 +23,7 @@ from tests.utils import GPUMemoryMonitor
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -129,7 +129,7 @@ def _run_zimage_generate(
 
 @pytest.mark.integration
 def test_zimage_tensor_parallel_tp2(tmp_path: Path):
-    if is_npu() or is_rocm():
+    if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
         pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
     if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
         pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -6,13 +6,14 @@
 # The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "0,1"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -20,6 +21,7 @@ stage_args:
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
+      max_num_batched_tokens: 32768
       hf_config_name: thinker_config
       tensor_parallel_size: 2
       load_format: dummy
@@ -36,13 +38,14 @@ stage_args:
       repetition_penalty: 1.05
 
   - stage_id: 1
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
        devices: "1"
        max_batch_size: 5
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: false
@@ -67,13 +70,14 @@ stage_args:
       stop_token_ids: [2150]
 
   - stage_id: 2
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml

@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true

tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/online_serving/test_images_generations_lora.py

@@ -27,8 +27,6 @@ from PIL import Image
 from safetensors.torch import save_file
 from vllm.utils.network_utils import get_open_port
 
-from vllm_omni.utils.platform_utils import is_npu
-
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 MODEL = "Tongyi-MAI/Z-Image-Turbo"
@@ -122,8 +120,6 @@ class OmniServer:
 
 @pytest.fixture(scope="module")
 def omni_server():
-    if is_npu():
-        pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.")
     with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
         yield server
 

tests/e2e/online_serving/test_qwen3_omni.py

@@ -27,7 +27,7 @@ from tests.conftest import (
     generate_synthetic_video,
     modify_stage_config,
 )
-from vllm_omni.utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
@@ -57,7 +57,7 @@ def get_chunk_config():
 
 CHUNK_CONFIG_PATH = get_chunk_config()
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if is_rocm():
+if current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
@@ -144,7 +144,9 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
+)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
     """
@@ -224,7 +226,9 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     assert similarity > 0.9, "The audio content is not same as the text"
 
 
-@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
+)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
     """

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
   engine_args:
     model_stage: thinker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
     gpu_memory_utilization: 0.6
     enforce_eager: false
@@ -43,7 +43,7 @@ stage_args:
   engine_args:
     model_stage: talker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
     gpu_memory_utilization: 0.3
     enforce_eager: false
@@ -71,7 +71,7 @@ stage_args:
   engine_args:
     model_stage: code2wav
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+    worker_type: generation
     scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
     enforce_eager: true
     trust_remote_code: true

tests/entrypoints/test_stage_utils.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -42,6 +43,14 @@ def _make_dummy_torch(call_log):
     return _Torch
 
 
+def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
+    """Create a mock platform for testing."""
+    mock_platform = MagicMock()
+    mock_platform.device_type = device_type
+    mock_platform.device_control_env_var = env_var
+    return mock_platform
+
+
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_respects_logical_ids(monkeypatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -49,9 +58,51 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
     call_log: list[int] = []
     dummy_torch = _make_dummy_torch(call_log)
     monkeypatch.setitem(sys.modules, "torch", dummy_torch)
-    monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
-    monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
+
+    # Mock the platform at the source module where it's defined
+    mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform",
+        mock_platform,
+    )
 
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
+
+
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+def test_set_stage_devices_npu_platform(monkeypatch):
+    """Test that set_stage_devices works correctly for NPU platform."""
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
+    call_log: list[int] = []
+
+    # Create NPU mock torch
+    class _Npu:
+        @staticmethod
+        def is_available():
+            return True
+
+        @staticmethod
+        def set_device(idx):
+            call_log.append(idx)
+
+        @staticmethod
+        def device_count():
+            return 2
+
+    class _NpuTorch:
+        npu = _Npu
+
+    monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
+
+    # Mock NPU platform at the source module where it's defined
+    mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform",
+        mock_platform,
+    )
+
+    set_stage_devices(stage_id=0, devices="0,1")
+
+    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"