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

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

.buildkite/pipeline.yml

@@ -18,11 +18,23 @@ steps:
   #     queue: "ascend"
 
   - label: "Simple Unit Test"
-    depends_on: ~
+    depends_on: image-build
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - pytest -v -s tests/entrypoints/
+      - pytest -v -s tests/diffusion/cache/
+      - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
+      - pytest -v -s tests/worker/
     agents:
-      queue: "cpu_queue_premerge"
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
     timeout_in_minutes: 20
@@ -149,7 +161,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
-      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:

.buildkite/test-amd.yaml

@@ -54,7 +54,7 @@ steps:
   commands:
     - export MIOPEN_DEBUG_CONV_DIRECT=0
     - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15

docs/configuration/stage_configs.md

@@ -40,7 +40,7 @@ stage_args:
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+      worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
       gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
       enforce_eager: true  # Now we only support eager mode
@@ -66,7 +66,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
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -92,7 +92,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

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -8,7 +8,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
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -34,7 +34,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
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -60,7 +60,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

examples/offline_inference/image_to_image/image_edit.py

@@ -81,7 +81,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -280,6 +280,16 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -306,12 +316,8 @@ def main():
     else:
         input_image = input_images
 
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
@@ -344,8 +350,8 @@ def main():
     # Initialize Omni with appropriate pipeline
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,

examples/offline_inference/image_to_image/image_to_image.md

@@ -47,4 +47,8 @@ Key arguments:
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

examples/offline_inference/image_to_video/README.md

@@ -54,4 +54,8 @@ Key arguments:
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -29,7 +29,7 @@ import torch
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -59,6 +59,16 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -80,8 +90,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
@@ -98,17 +107,13 @@ def main():
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         enable_cpu_offload=args.enable_cpu_offload,

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -22,7 +22,7 @@ import torch
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -118,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     print(f"\n{'=' * 60}")
     print("Stable Audio Open - Text-to-Audio Generation")

examples/offline_inference/text_to_image/README.md

@@ -96,8 +96,12 @@ Key arguments:
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 
 ## Web UI Demo

examples/offline_inference/text_to_image/gradio_demo.py

@@ -7,7 +7,7 @@ import torch
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -62,8 +62,8 @@ def parse_args() -> argparse.Namespace:
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
     # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    vae_use_slicing = current_omni_platform.is_npu()
+    vae_use_tiling = current_omni_platform.is_npu()
     return Omni(
         model=model_name,
         vae_use_slicing=vae_use_slicing,
@@ -72,7 +72,6 @@ def get_omni(model_name: str) -> Omni:
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
-    device = detect_device_type()
     omni = get_omni(args.model)
 
     def run_inference(
@@ -99,7 +98,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
             raise gr.Error("Inference steps must be a positive integer.")
         if num_images not in {1, 2, 3, 4}:
             raise gr.Error("Number of images must be 1, 2, 3, or 4.")
-        generator = torch.Generator(device=device).manual_seed(seed)
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
             OmniDiffusionSamplingParams(

examples/offline_inference/text_to_image/text_to_image.py

@@ -12,7 +12,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -113,17 +113,22 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Configure cache based on backend type
     cache_config = None
@@ -167,8 +172,8 @@ def main():
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,

examples/offline_inference/text_to_video/text_to_video.md

@@ -29,4 +29,8 @@ Key arguments:
 - `--boundary_ratio`: Boundary split ratio for low/high DiT.
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -13,7 +13,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -54,6 +54,16 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enforce_eager",
         action="store_true",
@@ -81,12 +91,7 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
     cache_config = None
@@ -118,8 +123,8 @@ def main():
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         cache_backend=args.cache_backend,

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -13,7 +13,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.9
       enforce_eager: true

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"

vllm_omni/diffusion/attention/backends/abstract.py

@@ -7,12 +7,18 @@ from typing import Generic, TypeVar
 
 import torch
 
+from vllm_omni.platforms import current_omni_platform
+
 
 class AttentionBackend(ABC):
     """Abstract class for diffusion attention backends."""
 
     accept_output_buffer: bool = False
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return False
+
     @staticmethod
     @abstractmethod
     def get_name() -> str:
@@ -77,12 +83,58 @@ class AttentionImpl(ABC, Generic[T]):
     ) -> None:
         raise NotImplementedError
 
-    @abstractmethod
     def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: T,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        """Dispatch to platform-specific forward implementation."""
+        if current_omni_platform.is_rocm():
+            return self.forward_hip(query, key, value, attn_metadata)
+        elif current_omni_platform.is_cuda():
+            return self.forward_cuda(query, key, value, attn_metadata)
+        elif current_omni_platform.is_npu():
+            return self.forward_npu(query, key, value, attn_metadata)
+        elif current_omni_platform.is_xpu():
+            return self.forward_xpu(query, key, value, attn_metadata)
+        else:
+            raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
     ) -> torch.Tensor:
         raise NotImplementedError
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_hip(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        # By default, HIP ops are compatible with CUDA ops.
+        return self.forward_cuda(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/ascend_attn.py

@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from importlib.util import find_spec
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-
-logger = init_logger(__name__)
-
-
-class AscendAttentionBackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [x for x in range(1024)]  # todo
-
-    @staticmethod
-    def get_name() -> str:
-        return "ASCEND"
-
-    @staticmethod
-    def get_impl_cls() -> type["AscendAttentionBackendImpl"]:
-        return AscendAttentionBackendImpl
-
-
-class AscendAttentionBackendImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-
-    def forward_native(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
-        attention_mask = attn_metadata.attn_mask if attn_metadata else None
-
-        output = torch.nn.functional.scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            attn_mask=attention_mask,
-            dropout_p=0.0,
-            is_causal=self.causal,
-            scale=self.softmax_scale,
-        )
-        out = output.permute(0, 2, 1, 3)
-        return out
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        if find_spec("mindiesd"):
-            from mindiesd import attention_forward
-
-            attention_mask = attn_metadata.attn_mask if attn_metadata else None
-            output = attention_forward(
-                query,
-                key,
-                value,
-                attn_mask=attention_mask,
-                opt_mode="manual",
-                op_type="fused_attn_score",
-                layout="BNSD",
-            )
-            return output
-        else:
-            return self.forward_native(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -9,13 +9,6 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
-from vllm_omni.diffusion.attention.backends.utils.fa import (
-    _pad_input,
-    _unpad_input,
-    _upad_input,
-    flash_attn_func,
-    flash_attn_varlen_func,
-)
 
 logger = init_logger(__name__)
 
@@ -23,6 +16,10 @@ logger = init_logger(__name__)
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [64, 96, 128, 192, 256]
@@ -51,25 +48,22 @@ class FlashAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward(
+    def forward_cuda(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        """
-        Flash attention implementation.
+        """CUDA/ROCm flash attention implementation."""
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+            flash_attn_func,
+            flash_attn_varlen_func,
+        )
 
-        Args:
-            query: (batch_size, seq_len, num_heads, head_dim)
-            key: (batch_size, seq_len, num_heads, head_dim)
-            value: (batch_size, seq_len, num_heads, head_dim)
-            attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask
-
-        Returns:
-            (batch_size, seq_len, num_heads, head_dim)
-        """
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -106,3 +100,25 @@ class FlashAttentionImpl(AttentionImpl):
                 softmax_scale=self.softmax_scale,
             )
         return out
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        """NPU attention implementation using mindiesd."""
+        from mindiesd import attention_forward
+
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        output = attention_forward(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            opt_mode="manual",
+            op_type="fused_attn_score",
+            layout="BNSD",
+        )
+        return output

vllm_omni/diffusion/attention/backends/registry.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion attention backend registry.
+
+This module provides an enum-based registry for diffusion attention backends,
+similar to vLLM's AttentionBackendEnum. Each backend registers its class path,
+and platforms can override or extend backends using register_backend().
+"""
+
+from collections.abc import Callable
+from enum import Enum, EnumMeta
+from typing import TYPE_CHECKING
+
+from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend
+
+logger = init_logger(__name__)
+
+
+class _DiffusionBackendEnumMeta(EnumMeta):
+    """Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
+
+    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
+        """Get backend by name with helpful error messages."""
+        try:
+            return super().__getitem__(name)  # type: ignore[return-value]
+        except KeyError:
+            members = list(cls.__members__.keys())
+            valid_backends = ", ".join(members)
+            raise ValueError(
+                f"Unknown diffusion attention backend: '{name}'. Valid options are: {valid_backends}"
+            ) from None
+
+
+class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
+    """Enumeration of all supported diffusion attention backends.
+
+    The enum value is the default class path, but this can be overridden
+    at runtime using register_backend().
+
+    To get the actual backend class (respecting overrides), use:
+        backend.get_class()
+
+    Example:
+        # Get backend class
+        backend = DiffusionAttentionBackendEnum.FLASH_ATTN
+        backend_cls = backend.get_class()
+
+        # Register custom backend
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.CUSTOM)
+        class MyCustomBackend:
+            ...
+    """
+
+    # Common backends (available on most platforms)
+    FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
+    TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
+    SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+
+    def get_path(self, include_classname: bool = True) -> str:
+        """Get the class path for this backend (respects overrides).
+
+        Returns:
+            The fully qualified class path string
+
+        Raises:
+            ValueError: If backend has empty path and is not registered
+        """
+        path = _DIFFUSION_ATTN_OVERRIDES.get(self, self.value)
+        if not path:
+            raise ValueError(
+                f"Backend {self.name} must be registered before use. "
+                f"Use register_diffusion_backend(DiffusionAttentionBackendEnum.{self.name}, "
+                f"'your.module.YourClass')"
+            )
+        if not include_classname:
+            path = path.rsplit(".", 1)[0]
+        return path
+
+    def get_class(self) -> "type[AttentionBackend]":
+        """Get the backend class (respects overrides).
+
+        Returns:
+            The backend class
+
+        Raises:
+            ImportError: If the backend class cannot be imported
+            ValueError: If backend has empty path and is not registered
+        """
+        return resolve_obj_by_qualname(self.get_path())
+
+    def is_overridden(self) -> bool:
+        """Check if this backend has been overridden.
+
+        Returns:
+            True if the backend has a registered override
+        """
+        return self in _DIFFUSION_ATTN_OVERRIDES
+
+    def clear_override(self) -> None:
+        """Clear any override for this backend, reverting to the default."""
+        _DIFFUSION_ATTN_OVERRIDES.pop(self, None)
+
+
+# Override registry
+_DIFFUSION_ATTN_OVERRIDES: dict[DiffusionAttentionBackendEnum, str] = {}
+
+
+def register_diffusion_backend(
+    backend: DiffusionAttentionBackendEnum,
+    class_path: str | None = None,
+) -> Callable[[type], type]:
+    """Register or override a diffusion backend implementation.
+
+    Args:
+        backend: The DiffusionAttentionBackendEnum member to register
+        class_path: Optional class path. If not provided and used as
+            decorator, will be auto-generated from the class.
+
+    Returns:
+        Decorator function if class_path is None, otherwise a no-op
+
+    Examples:
+        # Override an existing backend
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.FLASH_ATTN)
+        class MyCustomFlashAttn:
+            ...
+
+        # Override an existing backend (e.g., ASCEND_ATTN)
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.ASCEND_ATTN)
+        class CustomAscendAttentionBackend:
+            ...
+
+        # Direct registration
+        register_diffusion_backend(
+            DiffusionAttentionBackendEnum.CUSTOM,
+            "my.module.MyCustomBackend"
+        )
+    """
+
+    def decorator(cls: type) -> type:
+        _DIFFUSION_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
+        return cls
+
+    if class_path is not None:
+        _DIFFUSION_ATTN_OVERRIDES[backend] = class_path
+        return lambda x: x
+
+    return decorator

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -54,7 +54,7 @@ class SageAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward(
+    def forward_cuda(
         self,
         query: torch.Tensor,
         key: torch.Tensor,

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -16,6 +16,10 @@ logger = init_logger(__name__)
 class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [x for x in range(1024)]  # todo

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -15,9 +15,9 @@
 import torch
 import torch.nn.functional as F
 
-from vllm_omni.utils.platform_utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
-if is_rocm():
+if current_omni_platform.is_rocm():
     from vllm._aiter_ops import is_aiter_found_and_supported
 
     # Choose to enable this by default on ROCm

vllm_omni/diffusion/attention/selector.py

@@ -1,52 +1,54 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Diffusion attention backend selector.
+
+This module provides the interface for selecting diffusion attention backends.
+The actual backend selection logic is delegated to the platform layer
+(vllm_omni.platforms), similar to how vLLM handles attention backend selection.
+
+Usage:
+    from vllm_omni.diffusion.attention.selector import get_attn_backend
+
+    # Get the appropriate backend for current platform
+    backend_cls = get_attn_backend(head_size=64)
+
+    # Or override via environment variable
+    # export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+"""
 
 import importlib
 import os
 from functools import cache
 
-import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
-from vllm_omni.utils.platform_utils import detect_device_type, is_rocm
 
 logger = init_logger(__name__)
 
-# environment variable value -> backend module and class
-_BACKEND_CONFIG = {
-    "FLASH_ATTN": {
-        "module": "vllm_omni.diffusion.attention.backends.flash_attn",
-        "class": "FlashAttentionBackend",
-    },
-    "TORCH_SDPA": {
-        "module": "vllm_omni.diffusion.attention.backends.sdpa",
-        "class": "SDPABackend",
-    },
-    "SAGE_ATTN": {
-        "module": "vllm_omni.diffusion.attention.backends.sage_attn",
-        "class": "SageAttentionBackend",
-    },
-    "ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"},
-}
 
-_BACKENDS_SUPPORT_ATTENTION_MASK = ["SDPA", "ASCEND", "FLASH_ATTN"]
+def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
+    """Load a backend class from its fully qualified path.
 
+    Args:
+        cls_path: Fully qualified class path (e.g.,
+            "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend")
 
-def load_backend(backend_name: str) -> type[AttentionBackend]:
-    config = _BACKEND_CONFIG[backend_name]
-
+    Returns:
+        The loaded backend class
+    """
+    module_path, class_name = cls_path.rsplit(".", 1)
     try:
-        module = importlib.import_module(config["module"])
-        backend_class = getattr(module, config["class"])
+        module = importlib.import_module(module_path)
+        backend_class = getattr(module, class_name)
         return backend_class
     except ImportError as e:
-        raise ImportError(f"Failed to import module {config['module']}: {e}")
+        raise ImportError(f"Failed to import module {module_path}: {e}")
     except AttributeError as e:
-        raise AttributeError(f"Class {config['class']} not found in module: {e}")
+        raise AttributeError(f"Class {class_name} not found in module: {e}")
 
 
 @cache
@@ -54,65 +56,30 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
     """
     Get attention backend for diffusion models.
 
-    The backend is selected based on the following priority:
-    1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
-    2. Default backend (SDPA)
+    The backend selection is delegated to the current platform
+    (vllm_omni.platforms.current_omni_platform), which selects the
+    appropriate backend based on:
+    1. User override via DIFFUSION_ATTENTION_BACKEND environment variable
+    2. Platform-specific defaults and capabilities
+
+    This is similar to how vLLM's get_attn_backend_cls works, where the
+    platform layer decides which backend to use based on hardware capabilities.
 
     Args:
-        head_size: Head size (currently not used for selection, but kept for API compatibility)
+        head_size: Head size for attention computation (may affect backend selection)
 
     Returns:
         The selected attention backend class
     """
-    # Check environment variable
+    from vllm_omni.platforms import current_omni_platform
 
-    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None)
+    # Check environment variable for user override
+    selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
 
-    if detect_device_type() == "cuda" and not is_rocm():
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-        if 80 <= major * 10 + minor < 100:
-            if backend_name is None:
-                backend_name = "FLASH_ATTN"
-        else:
-            if backend_name == "FLASH_ATTN":
-                logger.warning(
-                    """Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. "
-                               "Falling back to TORCH_SDPA backend."""
-                )
-                backend_name = "TORCH_SDPA"
-    elif detect_device_type() == "cuda" and is_rocm():
-        from vllm._aiter_ops import is_aiter_found_and_supported
+    # Delegate to platform for backend selection
+    backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
+        selected_backend=selected_backend,
+        head_size=head_size,
+    )
 
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-
-        # Choose to enable this by default on ROCm
-        # Whenever possible as it is the fastest backend
-        # is_aiter_found_and_supported() checks if aiter library is found
-        # and is aiter supported on the current platform
-        # aiter currently only is supported on gfx942 and gfx950
-        # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py
-        if is_aiter_found_and_supported() and 90 < major * 10 + minor < 100:
-            if backend_name is None:
-                backend_name = "FLASH_ATTN"
-        else:
-            if backend_name == "FLASH_ATTN":
-                logger.warning(
-                    "Flash Attention requires `aiter` library which is only supported "
-                    "on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
-                )
-                backend_name = "TORCH_SDPA"
-
-    if backend_name is not None:
-        backend_name_upper = backend_name.upper()
-        if backend_name_upper not in _BACKEND_CONFIG:
-            valid_backends = list(_BACKEND_CONFIG.keys())
-            raise ValueError(
-                f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
-            )
-        logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
-        return load_backend(backend_name_upper)
-
-    # Default to SDPA
-    return SDPABackend
+    return _load_backend_cls(backend_cls_path)

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)

vllm_omni/diffusion/envs.py

@@ -5,20 +5,15 @@ import os
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
-import torch
-from packaging import version
 from vllm.logger import init_logger
 
-logger = init_logger(__name__)
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     MASTER_ADDR: str = ""
     MASTER_PORT: int | None = None
     CUDA_HOME: str | None = None
     LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: str | None = None
-    CUDA_VERSION: version.Version
-    TORCH_VERSION: version.Version
 
 environment_variables: dict[str, Callable[[], Any]] = {
     # ================== Runtime Env Vars ==================
@@ -32,127 +27,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # local rank of the process in the distributed setting, used to determine
     # the GPU device id
     "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
-    # used to control the visible devices in the distributed setting
-    "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
 }
 
-
-def _is_hip():
-    has_rocm = torch.version.hip is not None
-    return has_rocm
-
-
-def _is_cuda():
-    has_cuda = torch.version.cuda is not None
-    return has_cuda
-
-
-def _is_musa():
-    try:
-        if hasattr(torch, "musa") and torch.musa.is_available():
-            return True
-    except ModuleNotFoundError:
-        return False
-
-
-def _is_mps():
-    return torch.backends.mps.is_available()
-
-
-def _is_npu():
-    try:
-        if hasattr(torch, "npu") and torch.npu.is_available():
-            return True
-    except ModuleNotFoundError:
-        return False
-
-
-def get_device(local_rank: int) -> torch.device:
-    if _is_cuda() or _is_hip():
-        return torch.device("cuda", local_rank)
-    elif _is_musa():
-        return torch.device("musa", local_rank)
-    elif _is_mps():
-        return torch.device("mps")
-    elif _is_npu():
-        return torch.device("npu", local_rank)
-    else:
-        return torch.device("cpu")
-
-
-def get_device_name() -> str:
-    if _is_cuda() or _is_hip():
-        return "cuda"
-    elif _is_musa():
-        return "musa"
-    elif _is_mps():
-        return "mps"
-    elif _is_npu():
-        return "npu"
-    else:
-        return "cpu"
-
-
-def get_device_version():
-    if _is_hip():
-        hip_version = torch.version.hip
-        hip_version = hip_version.split("-")[0]
-        return hip_version
-    elif _is_cuda():
-        return torch.version.cuda
-    elif _is_musa():
-        return torch.version.musa
-    elif _is_mps():
-        return None
-    elif _is_npu():
-        return None
-    else:
-        raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
-
-
-def get_torch_distributed_backend() -> str:
-    if _is_cuda() or _is_hip():
-        return "nccl"
-    elif _is_musa():
-        return "mccl"
-    elif _is_mps():
-        return "gloo"
-    elif _is_npu():
-        return "hccl"
-    else:
-        raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
-
-
-variables: dict[str, Callable[[], Any]] = {
-    # ================== Other Vars ==================
-    # used in version checking
-    "CUDA_VERSION": lambda: version.parse(get_device_version() or "0.0"),
-    "TORCH_VERSION": lambda: version.parse(version.parse(torch.__version__).base_version),
-}
-
-
-def _setup_musa(environment_variables, variables):
-    musa = getattr(torch, "musa", None)
-    if musa is None:
-        return
-    try:
-        if musa.is_available():
-            environment_variables["MUSA_HOME"] = lambda: os.environ.get("MUSA_HOME", None)
-            environment_variables["MUSA_VISIBLE_DEVICES"] = lambda: os.environ.get("MUSA_VISIBLE_DEVICES", None)
-            musa_ver = getattr(getattr(torch, "version", None), "musa", None)
-            if musa_ver:
-                variables["MUSA_VERSION"] = lambda: version.parse(musa_ver)
-    except Exception:
-        pass
-
-
-try:
-    _setup_musa(environment_variables, variables)
-except (AttributeError, ModuleNotFoundError):
-    pass
+logger = init_logger(__name__)
 
 
 class PackagesEnvChecker:
+    """Singleton class for checking package availability."""
+
     _instance = None
 
     def __new__(cls):
@@ -163,38 +45,39 @@ class PackagesEnvChecker:
 
     def initialize(self):
         packages_info = {}
-        packages_info["has_flash_attn"] = self.check_flash_attn(packages_info)
+        packages_info["has_flash_attn"] = self._check_flash_attn(packages_info)
         self.packages_info = packages_info
 
-    def check_flash_attn(self, packages_info):
-        if not torch.cuda.is_available():
+    def _check_flash_attn(self, packages_info) -> bool:
+        """Check if flash attention is available and compatible."""
+        platform = current_omni_platform
+
+        # Flash attention requires CUDA-like platforms (CUDA or ROCm)
+        if not platform.is_cuda_alike():
             return False
 
-        # Check if torch_npu is available
-        if _is_npu():
-            logger.info("`flash_attn` is not ready on torch_npu for now")
+        # Check if devices are available
+        if platform.get_device_count() == 0:
             return False
 
-        if _is_musa():
-            logger.info("Flash Attention library is not supported on MUSA for the moment.")
-            return False
         try:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            gpu_name = torch.cuda.get_device_name(device)
+            gpu_name = platform.get_device_name()
+            # Turing/Tesla/T4 GPUs don't support flash attention well
             if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
                 return False
-            else:
-                from flash_attn import __version__
 
-                if __version__ < "2.6.0":
-                    raise ImportError("install flash_attn >= 2.6.0")
-                return True
+            from flash_attn import __version__
+
+            if __version__ < "2.6.0":
+                raise ImportError("install flash_attn >= 2.6.0")
+            return True
         except ImportError:
             if not packages_info.get("has_aiter", False):
                 logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation')
             return False
 
-    def get_packages_info(self):
+    def get_packages_info(self) -> dict:
+        """Get the packages info dictionary."""
         return self.packages_info
 
 
@@ -205,8 +88,6 @@ def __getattr__(name):
     # lazy evaluation of environment variables
     if name in environment_variables:
         return environment_variables[name]()
-    if name in variables:
-        return variables[name]()
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -10,7 +10,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import Scheduler
-from vllm_omni.utils.platform_utils import get_diffusion_worker_class
+from vllm_omni.diffusion.worker import WorkerProc
 
 logger = init_logger(__name__)
 
@@ -77,9 +77,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         mp.set_start_method("spawn", force=True)
         processes = []
 
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
         # Launch all worker processes
         scheduler_pipe_readers = []
         scheduler_pipe_writers = []
@@ -88,7 +85,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             reader, writer = mp.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
             process = mp.Process(
-                target=worker_proc.worker_main,
+                target=WorkerProc.worker_main,
                 args=(
                     i,  # rank
                     od_config,

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -335,10 +335,7 @@ class SequenceParallelSplitHook(ModelHook):
         2. Creates an attention mask indicating valid vs padding positions
         3. Stores the mask and padding info in ForwardContext
         """
-        from vllm_omni.diffusion.attention.selector import (
-            _BACKENDS_SUPPORT_ATTENTION_MASK,
-            get_attn_backend,
-        )
+        from vllm_omni.diffusion.attention.selector import get_attn_backend
         from vllm_omni.diffusion.distributed.parallel_state import (
             get_ring_parallel_world_size,
             get_sequence_parallel_rank,
@@ -359,7 +356,7 @@ class SequenceParallelSplitHook(ModelHook):
 
         # Check backend compatibility
         attn_backend = get_attn_backend(-1)
-        if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
+        if not attn_backend.supports_attention_mask:
             raise ValueError(
                 f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
                 f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "

vllm_omni/diffusion/layers/adalayernorm.py

@@ -91,6 +91,14 @@ class AdaLayerNorm(CustomOp):
 
         return output, gate_result
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        mod_params: torch.Tensor,
+        index: torch.Tensor = None,
+    ) -> torch.Tensor:
+        return self.forward_native(x, mod_params, index)
+
     def forward_native(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/layers/custom_op.py

@@ -3,7 +3,7 @@ from typing import Any
 
 import torch.nn as nn
 
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 
 class CustomOp(nn.Module):
@@ -14,16 +14,17 @@ class CustomOp(nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
-        self.is_cuda = detect_device_type() == "cuda"
         self._forward_method = self.dispatch_forward()
 
     def dispatch_forward(self) -> Callable:
-        if is_rocm():
+        if current_omni_platform.is_rocm():
             return self.forward_hip
-        elif self.is_cuda:
+        elif current_omni_platform.is_cuda():
             return self.forward_cuda
-        elif is_npu():
+        elif current_omni_platform.is_npu():
             return self.forward_npu
+        elif current_omni_platform.is_xpu():
+            return self.forward_xpu
         else:
             return self.forward_native
 
@@ -44,6 +45,9 @@ class CustomOp(nn.Module):
     def forward_npu(self, *args, **kwargs):
         raise NotImplementedError
 
+    def forward_xpu(self, *args, **kwargs):
+        raise NotImplementedError
+
     def forward_hip(self, *args, **kwargs):
         # By default, we assume that HIP ops are compatible with CUDA ops.
         return self.forward_cuda(*args, **kwargs)

vllm_omni/diffusion/layers/rope.py

@@ -137,6 +137,14 @@ class RotaryEmbedding(CustomOp):
         else:
             return self.forward_native(x, cos, sin)
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x, cos, sin)
+
     def forward_native(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py

@@ -24,7 +24,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -610,7 +610,7 @@ class LongCatImageTransformer2DModel(nn.Module):
 
         ids = torch.cat((txt_ids, img_ids), dim=0)
 
-        if is_npu():
+        if current_omni_platform.is_npu():
             freqs_cos, freqs_sin = self.pos_embed(ids.cpu())
             image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
         else:

vllm_omni/diffusion/worker/__init__.py

@@ -2,14 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Worker classes for diffusion models."""
 
-from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import (
-    GPUDiffusionWorker,
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import (
+    DiffusionWorker,
     WorkerProc,
 )
 
 __all__ = [
-    "GPUDiffusionModelRunner",
-    "GPUDiffusionWorker",
+    "DiffusionModelRunner",
+    "DiffusionWorker",
     "WorkerProc",
 ]

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -34,7 +34,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 logger = init_logger(__name__)
 
 
-class GPUDiffusionModelRunner:
+class DiffusionModelRunner:
     """
     Model runner that handles model loading and execution for diffusion models.
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -4,7 +4,7 @@
 Diffusion Worker for vLLM-Omni.
 
 Handles GPU infrastructure initialization and delegates model operations
-to GPUDiffusionModelRunner.
+to DiffusionModelRunner.
 """
 
 import multiprocessing as mp
@@ -31,13 +31,14 @@ from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.profiler import CurrentProfiler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.lora.request import LoRARequest
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class GPUDiffusionWorker:
+class DiffusionWorker:
     """
     A worker that manages GPU infrastructure and delegates to the model runner.
 
@@ -47,7 +48,7 @@ class GPUDiffusionWorker:
     - Memory management (sleep/wake)
 
     All model-related operations (loading, compilation, execution) are
-    delegated to GPUDiffusionModelRunner.
+    delegated to DiffusionModelRunner.
     """
 
     def __init__(
@@ -61,7 +62,7 @@ class GPUDiffusionWorker:
         self.od_config = od_config
         self.device: torch.device | None = None
         self.vllm_config: VllmConfig | None = None
-        self.model_runner: GPUDiffusionModelRunner | None = None
+        self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
         self.init_device()
@@ -79,8 +80,8 @@ class GPUDiffusionWorker:
         os.environ["WORLD_SIZE"] = str(world_size)
 
         # Setup device
-        self.device = torch.device(f"cuda:{rank}")
-        torch.cuda.set_device(self.device)
+        self.device = current_omni_platform.get_torch_device(rank)
+        current_omni_platform.set_device(self.device)
 
         # Create vllm_config for parallel configuration
         vllm_config = VllmConfig()
@@ -105,7 +106,7 @@ class GPUDiffusionWorker:
             )
 
         # Create model runner and load model
-        self.model_runner = GPUDiffusionModelRunner(
+        self.model_runner = DiffusionModelRunner(
             vllm_config=self.vllm_config,
             od_config=self.od_config,
             device=self.device,
@@ -176,7 +177,7 @@ class GPUDiffusionWorker:
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
-        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
+        free_bytes_before_sleep = current_omni_platform.get_free_memory()
 
         # Save the buffers before level 2 sleep
         if level == 2 and self.model_runner is not None:
@@ -185,7 +186,9 @@ class GPUDiffusionWorker:
 
         allocator = CuMemAllocator.get_instance()
         allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
+        free_bytes_after_sleep = current_omni_platform.get_free_memory()
+        device_id = self.device.index if self.device.index is not None else 0
+        total = current_omni_platform.get_device_total_memory(device_id)
         freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
         used_bytes = total - free_bytes_after_sleep
         assert freed_bytes >= 0, "Memory usage increased after sleeping."
@@ -270,9 +273,9 @@ class WorkerProc:
         self.gpu_id = gpu_id
         self._running = True
 
-    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> GPUDiffusionWorker:
+    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> DiffusionWorker:
         """Create a worker instance. Override in subclasses for different worker types."""
-        return GPUDiffusionWorker(
+        return DiffusionWorker(
             local_rank=gpu_id,
             rank=gpu_id,
             od_config=od_config,
@@ -379,6 +382,9 @@ class WorkerProc:
         broadcast_handle,
     ) -> None:
         """Worker initialization and execution loops."""
+        from vllm_omni.plugins import load_omni_general_plugins
+
+        load_omni_general_plugins()
         worker_proc = WorkerProc(
             od_config,
             gpu_id=rank,

vllm_omni/diffusion/worker/npu/__init__.py

@@ -1,7 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""NPU worker classes for diffusion models."""
-
-from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorker, NPUWorkerProc
-
-__all__ = ["NPUWorker", "NPUWorkerProc"]

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -1,255 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import multiprocessing as mp
-import os
-import time
-from collections.abc import Iterable
-from contextlib import AbstractContextManager, nullcontext
-
-import torch
-from vllm.config import LoadConfig, VllmConfig
-from vllm.logger import init_logger
-from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
-
-from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.data import (
-    DiffusionOutput,
-    OmniDiffusionConfig,
-)
-from vllm_omni.diffusion.distributed.parallel_state import (
-    destroy_distributed_env,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc
-
-logger = init_logger(__name__)
-
-
-class NPUWorker:
-    """
-    A worker that executes the model on a single NPU.
-    Inherits from GPUWorker and overrides device-specific initialization.
-    """
-
-    def __init__(
-        self,
-        local_rank: int,
-        rank: int,
-        od_config: OmniDiffusionConfig,
-    ):
-        self.local_rank = local_rank
-        self.rank = rank
-        self.od_config = od_config
-        self.pipeline = None
-        self.device = None
-        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
-        self.init_device_and_model()
-
-    def init_device_and_model(self) -> None:
-        """Initialize the NPU device and load the model."""
-        world_size = self.od_config.num_gpus
-        rank = self.rank
-        # Set environment variables for distributed initialization
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = str(self.od_config.master_port)
-        os.environ["LOCAL_RANK"] = str(self.local_rank)
-        os.environ["RANK"] = str(rank)
-        os.environ["WORLD_SIZE"] = str(world_size)
-
-        device = torch.device(f"npu:{rank}")
-        torch.npu.set_device(device)
-
-        # hack
-        vllm_config = VllmConfig()
-        vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
-        vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
-        self.vllm_config = vllm_config
-        with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
-            init_distributed_environment(world_size=world_size, rank=rank)
-            logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
-            parallel_config = self.od_config.parallel_config
-            initialize_model_parallel(
-                data_parallel_size=parallel_config.data_parallel_size,
-                cfg_parallel_size=parallel_config.cfg_parallel_size,
-                sequence_parallel_size=parallel_config.sequence_parallel_size,
-                ulysses_degree=parallel_config.ulysses_degree,
-                ring_degree=parallel_config.ring_degree,
-                tensor_parallel_size=parallel_config.tensor_parallel_size,
-                pipeline_parallel_size=parallel_config.pipeline_parallel_size,
-            )
-
-            load_config = LoadConfig()
-            model_loader = DiffusersPipelineLoader(load_config)
-            time_before_load = time.perf_counter()
-            with DeviceMemoryProfiler() as m:
-                self.pipeline = model_loader.load_model(
-                    od_config=self.od_config,
-                    load_device=f"npu:{rank}",
-                )
-            time_after_load = time.perf_counter()
-
-        logger.info(
-            "Model loading took %.4f GiB and %.6f seconds",
-            m.consumed_memory / GiB_bytes,
-            time_after_load - time_before_load,
-        )
-        logger.info(f"Worker {self.rank}: Model loaded successfully.")
-
-        # Setup cache backend based on type (both backends use enable()/reset() interface)
-        self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config)
-
-        if self.cache_backend is not None:
-            self.cache_backend.enable(self.pipeline)
-
-    def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
-        """
-        Generate output for the given requests.
-
-        Args:
-            requests: List of diffusion requests
-
-        Returns:
-            DiffusionOutput with generated results
-        """
-        return self.execute_model(requests, self.od_config)
-
-    @torch.inference_mode()
-    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
-        """
-        Execute a forward pass.
-        """
-        assert self.pipeline is not None
-        if not reqs or len(reqs) == 0:
-            raise ValueError("Cannot execute model with empty request list")
-        # TODO: dealing with first req for now
-        req = reqs[0]
-
-        if req.generator is None and req.seed is not None:
-            req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
-
-        # Refresh cache context if needed
-        if (
-            not getattr(req, "skip_cache_refresh", False)
-            and self.cache_backend is not None
-            and self.cache_backend.is_enabled()
-        ):
-            self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
-        with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
-            output = self.pipeline.forward(req)
-        return output
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return self.pipeline.load_weights(weights)
-
-    def sleep(self, level: int = 1) -> bool:
-        """
-        Put the worker to sleep. The worker should not process any requests.
-        The caller should guarantee that no requests are being processed
-        during the sleep period, before `wake_up` is called.
-
-        Args:
-            level: The sleep level. Level 1 sleep will offload the model
-                weights and discard the kv cache.
-                Currently only support level 1.
-        """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
-        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
-
-        # Save the buffers before level 2 sleep
-        if level == 2:
-            model = self.pipeline
-            self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
-
-        allocator = CuMemAllocator.get_instance()
-        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
-        freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
-        used_bytes = total - free_bytes_after_sleep
-        assert freed_bytes >= 0, "Memory usage increased after sleeping."
-        logger.info(
-            "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.",
-            freed_bytes / GiB_bytes,
-            used_bytes / GiB_bytes,
-        )
-        return True
-
-    def wake_up(self, tags: list[str] | None = None) -> bool:
-        """
-        Wake up the worker from sleep mode. See the sleep function
-        method for more details.
-
-        Args:
-            tags: An optional list of tags to reallocate the worker memory
-                for specific memory allocations. Values must be in
-                `("weights")`. If None, all memory is reallocated.
-                wake_up should be called with all tags (or None) before the
-                worker is used again.
-        """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
-        allocator = CuMemAllocator.get_instance()
-        allocator.wake_up(tags)
-
-        # Restore the buffers after level 2 sleep
-        if len(self._sleep_saved_buffers):
-            model = self.pipeline
-            for name, buffer in model.named_buffers():
-                if name in self._sleep_saved_buffers:
-                    buffer.data.copy_(self._sleep_saved_buffers[name].data)
-            self._sleep_saved_buffers = {}
-        return True
-
-    def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
-        if self.od_config.enable_sleep_mode:
-            from vllm.device_allocator.cumem import CuMemAllocator
-
-            allocator = CuMemAllocator.get_instance()
-            if tag == "weights":
-                assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
-            return allocator.use_memory_pool(tag=tag)
-        else:
-            return nullcontext()
-
-    def shutdown(self) -> None:
-        destroy_distributed_env()
-
-
-class NPUWorkerProc(WorkerProc):
-    """Wrapper that runs one NPUWorker in a separate process."""
-
-    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> NPUWorker:
-        """Create an NPUWorker instead of GPUWorker."""
-        return NPUWorker(
-            local_rank=gpu_id,
-            rank=gpu_id,
-            od_config=od_config,
-        )
-
-    @staticmethod
-    def worker_main(
-        rank: int,
-        od_config: OmniDiffusionConfig,
-        pipe_writer: mp.connection.Connection,
-        broadcast_handle,
-    ) -> None:
-        """Worker initialization and execution loops."""
-
-        worker_proc = NPUWorkerProc(
-            od_config,
-            gpu_id=rank,
-            broadcast_handle=broadcast_handle,
-        )
-        logger.info(f"Worker {rank}: Scheduler loop started.")
-        pipe_writer.send(
-            {
-                "status": "ready",
-                "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
-            }
-        )
-        worker_proc.worker_busy_loop()
-        logger.info(f"Worker {rank}: Shutdown complete.")

vllm_omni/engine/arg_utils.py

@@ -7,6 +7,7 @@ from vllm.transformers_utils.config import get_hf_text_config
 from vllm.v1.engine.async_llm import AsyncEngineArgs
 
 from vllm_omni.config import OmniModelConfig
+from vllm_omni.plugins import load_omni_general_plugins
 
 logger = init_logger(__name__)
 
@@ -89,7 +90,11 @@ class OmniEngineArgs(EngineArgs):
                 f"Config attribute '{hf_config_name}' not found in hf_config, "
                 "falling back to default get_hf_text_config"
             )
-            return get_hf_text_config(hf_config)
+        return get_hf_text_config(hf_config)
+
+    def __post_init__(self) -> None:
+        load_omni_general_plugins()
+        super().__post_init__()
 
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
@@ -187,7 +192,11 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
                 f"Config attribute '{hf_config_name}' not found in hf_config, "
                 "falling back to default get_hf_text_config"
             )
-            return get_hf_text_config(hf_config)
+        return get_hf_text_config(hf_config)
+
+    def __post_init__(self) -> None:
+        load_omni_general_plugins()
+        super().__post_init__()
 
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):

vllm_omni/entrypoints/omni_stage.py

@@ -51,11 +51,27 @@ from vllm_omni.entrypoints.stage_utils import (
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils import detect_device_type
 
 logger = init_logger(__name__)
 
 
+def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
+    worker_type = engine_args.pop("worker_type", None)
+    if not worker_type:
+        return
+    if engine_args.get("worker_cls"):
+        return
+    from vllm_omni.platforms import current_omni_platform
+
+    worker_type = str(worker_type).lower()
+    if worker_type == "ar":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_ar_worker_cls()
+    elif worker_type == "generation":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_generation_worker_cls()
+    else:
+        raise ValueError(f"Unknown worker_type: {worker_type}")
+
+
 def _build_od_config(engine_args: dict[str, Any], model: str) -> dict[str, Any]:
     """Build OmniDiffusionConfig kwargs from engine args."""
     od_config = engine_args.get("od_config", {})
@@ -483,6 +499,9 @@ def _stage_worker(
     import os as _os
     import time as _time
 
+    from vllm_omni.plugins import load_omni_general_plugins
+
+    load_omni_general_plugins()
     # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
     # GPUARModelRunner) are spawned with a fork-safe method.
     # Mooncake / gRPC / RDMA and CUDA/NCCL can deadlock under fork-with-threads.
@@ -503,6 +522,9 @@ def _stage_worker(
     connectors_config = stage_payload.get("connectors_config", {})
     stage_type: Literal["llm", "diffusion"] = stage_payload.get("stage_type", "llm")
 
+    if stage_type != "diffusion":
+        _resolve_worker_cls(engine_args)
+
     # Aggregates for running average
     _agg_total_tokens = 0
     _agg_total_gen_time_ms = 0.0
@@ -512,7 +534,9 @@ def _stage_worker(
     # Device mapping
     device_type = None
     try:
-        device_type = detect_device_type()
+        from vllm_omni.platforms import current_omni_platform
+
+        device_type = current_omni_platform.device_type
         set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         logger.warning("Device setup failed: %s", e)
@@ -521,139 +545,137 @@ def _stage_worker(
     # when multiple instances start simultaneously
     # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage
     lock_files = []
-    if device_type == "cuda":
-        try:
-            import torch
+    try:
+        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
+        if "parallel_config" in engine_args:
+            parallel_config = engine_args["parallel_config"]
+            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
+            data_parallel_size = parallel_config.get("data_parallel_size", 1)
+            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
+            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+        else:
+            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+            data_parallel_size = engine_args.get("data_parallel_size", 1)
+            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = 1  # not use in omni model
+            cfg_parallel_size = 1  # not used in omni model
 
-            if torch.cuda.is_available():
-                # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
-                if "parallel_config" in engine_args:
-                    parallel_config = engine_args["parallel_config"]
-                    tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
-                    pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
-                    data_parallel_size = parallel_config.get("data_parallel_size", 1)
-                    prefill_context_parallel_size = 1  # not used for diffusion
-                    sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
-                    cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
-                else:
-                    tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-                    pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-                    data_parallel_size = engine_args.get("data_parallel_size", 1)
-                    prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
-                    sequence_parallel_size = 1  # not use in omni model
-                    cfg_parallel_size = 1  # not used in omni model
+        # Calculate total number of devices needed for this stage
+        # For a single stage worker:
+        # - TP: splits model across devices (always needed)
+        # - PP: splits layers across pipeline stages, but each stage uses TP devices
+        # - DP: replicates model, but each replica uses TP devices
+        # - PCP: context parallelism, typically uses TP devices
+        # - SP: sequence parallelism, typically uses TP devices
+        # - CFG: Classifier-Free Guidance parallelism for diffusion models
+        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
+        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+        num_devices_per_stage = (
+            tensor_parallel_size
+            * pipeline_parallel_size
+            * data_parallel_size
+            * prefill_context_parallel_size
+            * sequence_parallel_size
+            * cfg_parallel_size
+        )
 
-                # Calculate total number of devices needed for this stage
-                # For a single stage worker:
-                # - TP: splits model across GPUs (always needed)
-                # - PP: splits layers across pipelinestages, but each stage uses TP devices
-                # - DP: replicates model, but each replica uses TP devices
-                # - PCP: context parallelism, typically uses TP devices
-                # - SP: sequence parallelism, typically uses TP devices
-                # - CFG: Classifier-Free Guidance parallelism for diffusion models
-                # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
-                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-                num_devices_per_stage = (
-                    tensor_parallel_size
-                    * pipeline_parallel_size
-                    * data_parallel_size
-                    * prefill_context_parallel_size
-                    * sequence_parallel_size
-                    * cfg_parallel_size
-                )
+        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
+        # After set_stage_devices, this env var is set to physical device(s)
+        device_control_env = current_omni_platform.device_control_env_var
+        visible_devices_str = _os.environ.get(device_control_env)
+        physical_devices = []
 
-                # Get physical device IDs from CUDA_VISIBLE_DEVICES
-                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
-                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
-                physical_devices = []
+        if visible_devices_str:
+            try:
+                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+            except (ValueError, IndexError):
+                pass
 
-                if cuda_visible_devices:
+        if not physical_devices:
+            # Fallback: use logical device count if device control env var not set
+            num_devices = current_omni_platform.get_device_count()
+            physical_devices = list(range(num_devices))
+
+        # Determine which devices will be used (min of devices per stage and available devices)
+        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        devices_to_lock = physical_devices[:num_devices_to_lock]
+
+        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+        devices_to_lock = sorted(devices_to_lock)
+
+        logger.debug(
+            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+            tensor_parallel_size,
+            pipeline_parallel_size,
+            data_parallel_size,
+            prefill_context_parallel_size,
+            sequence_parallel_size,
+            cfg_parallel_size,
+            num_devices_to_lock,
+            devices_to_lock,
+        )
+
+        # Acquire exclusive locks for all devices using fcntl.flock
+        # Locks are automatically released when process dies
+        wait_start = _time.time()
+        acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+        for device_id in devices_to_lock:
+            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_acquired = False
+
+            while not lock_acquired:
+                try:
+                    # Open or create the lock file
+                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                    # Try to acquire exclusive lock (non-blocking first)
                     try:
-                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
-                    except (ValueError, IndexError):
-                        pass
+                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                        # Successfully acquired lock - write PID
+                        _os.ftruncate(lock_fd, 0)  # Clear file
+                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                        _os.fsync(lock_fd)  # Ensure written to disk
+                        lock_acquired = True
+                        acquired_lock_fds.append(lock_fd)
+                        logger.debug("Acquired exclusive lock for device %s", device_id)
+                    except BlockingIOError:
+                        # Lock is held by another process
+                        _os.close(lock_fd)
 
-                if not physical_devices:
-                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
-                    num_devices = torch.cuda.device_count()
-                    physical_devices = list(range(num_devices))
-
-                # Determine which devices will be used (min of devices per stage and available devices)
-                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-                devices_to_lock = physical_devices[:num_devices_to_lock]
-
-                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-                devices_to_lock = sorted(devices_to_lock)
-
-                logger.debug(
-                    "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d; will lock %d devices: %s",
-                    tensor_parallel_size,
-                    pipeline_parallel_size,
-                    data_parallel_size,
-                    prefill_context_parallel_size,
-                    sequence_parallel_size,
-                    num_devices_to_lock,
-                    devices_to_lock,
-                )
-
-                # Acquire exclusive locks for all devices using fcntl.flock
-                # Locks are automatically released when process dies
-                wait_start = _time.time()
-                acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-                for device_id in devices_to_lock:
-                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-                    lock_acquired = False
-
-                    while not lock_acquired:
-                        try:
-                            # Open or create the lock file
-                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                            # Try to acquire exclusive lock (non-blocking first)
-                            try:
-                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                                # Successfully acquired lock - write PID
-                                _os.ftruncate(lock_fd, 0)  # Clear file
-                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                                _os.fsync(lock_fd)  # Ensure written to disk
-                                lock_acquired = True
-                                acquired_lock_fds.append(lock_fd)
-                                logger.debug("Acquired exclusive lock for device %s", device_id)
-                            except BlockingIOError:
-                                # Lock is held by another process
-                                _os.close(lock_fd)
-
-                                # Check if we've been waiting too long
-                                if _time.time() - wait_start > stage_init_timeout:
-                                    logger.warning(
-                                        "Timeout waiting for device %s initialization lock, proceeding anyway",
-                                        device_id,
-                                    )
-                                    break
-
-                                # Wait a bit before retrying
-                                _time.sleep(0.1)
-                        except OSError as e:
-                            # Other error - log and continue without lock
-                            logger.debug(
-                                "Failed to acquire lock for device %s: %s, continuing anyway",
+                        # Check if we've been waiting too long
+                        if _time.time() - wait_start > stage_init_timeout:
+                            logger.warning(
+                                "Timeout waiting for device %s initialization lock, proceeding anyway",
                                 device_id,
-                                e,
                             )
-                            try:
-                                _os.close(lock_fd)
-                            except (OSError, NameError):
-                                pass
                             break
 
-                lock_files = acquired_lock_fds
-        except Exception as e:
-            logger.debug(
-                "[Stage-%s] Failed to set up sequential initialization lock: %s",
-                stage_id,
-                e,
-            )
+                        # Wait a bit before retrying
+                        _time.sleep(0.1)
+                except OSError as e:
+                    # Other error - log and continue without lock
+                    logger.debug(
+                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                        device_id,
+                        e,
+                    )
+                    try:
+                        _os.close(lock_fd)
+                    except (OSError, NameError):
+                        pass
+                    break
+
+        lock_files = acquired_lock_fds
+    except Exception as e:
+        logger.debug(
+            "[Stage-%s] Failed to set up sequential initialization lock: %s",
+            stage_id,
+            e,
+        )
     # Init engine based on stage_type
     logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()))
     if engine_args.get("async_chunk", False):
@@ -1001,6 +1023,9 @@ async def _stage_worker_async(
     import os as _os
     import time as _time
 
+    from vllm_omni.plugins import load_omni_general_plugins
+
+    load_omni_general_plugins()
     # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
     # GPUARModelRunner) are spawned with a fork-safe method.
     if _os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
@@ -1018,6 +1043,9 @@ async def _stage_worker_async(
     connectors_config = stage_payload.get("connectors_config", {})
     stage_type = stage_payload.get("stage_type", "llm")
 
+    if stage_type != "diffusion":
+        _resolve_worker_cls(engine_args)
+
     in_q = omni_stage._in_q
     out_q = omni_stage._out_q
 
@@ -1031,9 +1059,9 @@ async def _stage_worker_async(
     # Device mapping
     device_type = None
     try:
-        from vllm_omni.utils import detect_device_type
+        from vllm_omni.platforms import current_omni_platform
 
-        device_type = detect_device_type()
+        device_type = current_omni_platform.device_type
         set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         logger.warning("Device setup failed: %s", e)
@@ -1051,120 +1079,136 @@ async def _stage_worker_async(
 
     # Sequential initialization on the same device to avoid memory calculation errors
     # when multiple instances start simultaneously
-    # For TP, we need to lock ALL devices that will be used by this stage
+    # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage
     lock_files = []
-    if device_type == "cuda":
-        try:
-            import torch
+    try:
+        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
+        if "parallel_config" in engine_args:
+            parallel_config = engine_args["parallel_config"]
+            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
+            data_parallel_size = parallel_config.get("data_parallel_size", 1)
+            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
+            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+        else:
+            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+            data_parallel_size = engine_args.get("data_parallel_size", 1)
+            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = 1  # not use in omni model
+            cfg_parallel_size = 1  # not used in omni model
 
-            if torch.cuda.is_available():
-                # Get all parallel sizes from engine_args (defaults to 1)
-                tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-                pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-                data_parallel_size = engine_args.get("data_parallel_size", 1)
-                prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+        # Calculate total number of devices needed for this stage
+        # For a single stage worker:
+        # - TP: splits model across devices (always needed)
+        # - PP: splits layers across pipeline stages, but each stage uses TP devices
+        # - DP: replicates model, but each replica uses TP devices
+        # - PCP: context parallelism, typically uses TP devices
+        # - SP: sequence parallelism, typically uses TP devices
+        # - CFG: Classifier-Free Guidance parallelism for diffusion models
+        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
+        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+        num_devices_per_stage = (
+            tensor_parallel_size
+            * pipeline_parallel_size
+            * data_parallel_size
+            * prefill_context_parallel_size
+            * sequence_parallel_size
+            * cfg_parallel_size
+        )
 
-                # Calculate total number of devices needed for this stage
-                # For a single stage worker in omni:
-                # - TP: splits model across GPUs (always needed)
-                # - PP: splits layers across stages, but each stage uses TP devices
-                # - DP: replicates model, but each replica uses TP devices
-                # - PCP: context parallelism, typically uses TP devices
-                # The number of devices per stage is determined by TP * PP * DP * PCP size
-                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-                num_devices_per_stage = (
-                    tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size
-                )
+        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
+        # After set_stage_devices, this env var is set to physical device(s)
+        device_control_env = current_omni_platform.device_control_env_var
+        visible_devices_str = _os.environ.get(device_control_env)
+        physical_devices = []
 
-                # Get physical device IDs from CUDA_VISIBLE_DEVICES
-                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
-                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
-                physical_devices = []
+        if visible_devices_str:
+            try:
+                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+            except (ValueError, IndexError):
+                pass
 
-                if cuda_visible_devices:
+        if not physical_devices:
+            # Fallback: use logical device count if device control env var not set
+            num_devices = current_omni_platform.get_device_count()
+            physical_devices = list(range(num_devices))
+
+        # Determine which devices will be used (min of devices per stage and available devices)
+        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        devices_to_lock = physical_devices[:num_devices_to_lock]
+
+        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+        devices_to_lock = sorted(devices_to_lock)
+
+        logger.debug(
+            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+            tensor_parallel_size,
+            pipeline_parallel_size,
+            data_parallel_size,
+            prefill_context_parallel_size,
+            sequence_parallel_size,
+            cfg_parallel_size,
+            num_devices_to_lock,
+            devices_to_lock,
+        )
+
+        # Acquire exclusive locks for all devices using fcntl.flock
+        # Locks are automatically released when process dies
+        wait_start = _time.time()
+        acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+        for device_id in devices_to_lock:
+            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_acquired = False
+
+            while not lock_acquired:
+                try:
+                    # Open or create the lock file
+                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                    # Try to acquire exclusive lock (non-blocking first)
                     try:
-                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
-                    except (ValueError, IndexError):
-                        pass
+                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                        # Successfully acquired lock - write PID
+                        _os.ftruncate(lock_fd, 0)  # Clear file
+                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                        _os.fsync(lock_fd)  # Ensure written to disk
+                        lock_acquired = True
+                        acquired_lock_fds.append(lock_fd)
+                        logger.debug("Acquired exclusive lock for device %s", device_id)
+                    except BlockingIOError:
+                        # Lock is held by another process
+                        _os.close(lock_fd)
 
-                if not physical_devices:
-                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
-                    num_devices = torch.cuda.device_count()
-                    physical_devices = list(range(num_devices))
-
-                # Determine which devices will be used (min of devices per stage and available devices)
-                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-                devices_to_lock = physical_devices[:num_devices_to_lock]
-
-                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-                devices_to_lock = sorted(devices_to_lock)
-
-                logger.debug(
-                    "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
-                    tensor_parallel_size,
-                    pipeline_parallel_size,
-                    data_parallel_size,
-                    prefill_context_parallel_size,
-                    num_devices_to_lock,
-                    devices_to_lock,
-                )
-
-                # Acquire exclusive locks for all devices using fcntl.flock
-                # Locks are automatically released when process dies
-                wait_start = _time.time()
-                acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-                for device_id in devices_to_lock:
-                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-                    lock_acquired = False
-
-                    while not lock_acquired:
-                        try:
-                            # Open or create the lock file
-                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                            # Try to acquire exclusive lock (non-blocking first)
-                            try:
-                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                                # Successfully acquired lock - write PID
-                                _os.ftruncate(lock_fd, 0)  # Clear file
-                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                                _os.fsync(lock_fd)  # Ensure written to disk
-                                lock_acquired = True
-                                acquired_lock_fds.append(lock_fd)
-                                logger.debug("Acquired exclusive lock for device %s", device_id)
-                            except BlockingIOError:
-                                # Lock is held by another process
-                                _os.close(lock_fd)
-
-                                # Check if we've been waiting too long
-                                if _time.time() - wait_start > stage_init_timeout:
-                                    logger.warning(
-                                        "Timeout waiting for device %s initialization lock, "
-                                        "proceeding anyway with timeout %s",
-                                        device_id,
-                                        stage_init_timeout,
-                                    )
-                                    break
-
-                                # Wait a bit before retrying
-                                _time.sleep(0.1)
-                        except OSError as e:
-                            # Other error - log and continue without lock
-                            logger.debug(
-                                "Failed to acquire lock for device %s: %s, continuing anyway",
+                        # Check if we've been waiting too long
+                        if _time.time() - wait_start > stage_init_timeout:
+                            logger.warning(
+                                "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s",
                                 device_id,
-                                e,
+                                stage_init_timeout,
                             )
-                            try:
-                                _os.close(lock_fd)
-                            except (OSError, NameError):
-                                pass
                             break
 
-                lock_files = acquired_lock_fds
-        except Exception as e:
-            logger.debug("Failed to set up sequential initialization lock: %s", e)
+                        # Wait a bit before retrying
+                        _time.sleep(0.1)
+                except OSError as e:
+                    # Other error - log and continue without lock
+                    logger.debug(
+                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                        device_id,
+                        e,
+                    )
+                    try:
+                        _os.close(lock_fd)
+                    except (OSError, NameError):
+                        pass
+                    break
+
+        lock_files = acquired_lock_fds
+    except Exception as e:
+        logger.debug("Failed to set up sequential initialization lock: %s", e)
 
     # Init engine based on stage_type
     logger.debug(

vllm_omni/entrypoints/stage_utils.py

@@ -57,12 +57,12 @@ def set_stage_devices(
         - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device()
         - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device()
     """
-    from vllm_omni.utils import detect_device_type, get_device_control_env_var
+    from vllm_omni.platforms import current_omni_platform
 
     if device_type is None:
-        device_type = detect_device_type()
+        device_type = current_omni_platform.device_type
 
-    env_var = get_device_control_env_var()
+    env_var = current_omni_platform.device_control_env_var
 
     try:
         selected_physical: int | None = None

vllm_omni/entrypoints/utils.py

@@ -10,7 +10,7 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
 from vllm_omni.entrypoints.stage_utils import _to_dict
-from vllm_omni.utils import detect_device_type, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
@@ -164,16 +164,12 @@ def resolve_model_config_path(model: str) -> str:
                 f"Model is not in standard transformers format and does not have model_index.json. "
                 f"Please ensure the model has proper configuration files with 'model_type' field"
             )
-    device_type = detect_device_type()
 
-    # Try device-specific config first
-    if device_type != "cuda" or is_rocm():
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        if is_rocm():
-            device_config_file = f"vllm_omni/model_executor/stage_configs/rocm/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            return str(device_config_path)
+    default_config_path = current_omni_platform.get_default_stage_config_path()
+    model_type_str = f"{model_type}.yaml"
+    complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
+    if os.path.exists(complete_config_path):
+        return str(complete_config_path)
 
     # Fall back to default config
     stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -36,7 +36,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -247,7 +247,7 @@ class Qwen2_5OmniForConditionalGeneration(
             if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
                 inputs_embeds = inputs_embeds.to(thinker_dev)
 
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # TODO: remove this hack when NPU supports batched inputs properly
                 thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
                 # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py

@@ -35,7 +35,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
@@ -732,7 +732,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
         beta = 0.0
 
     # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
-    if is_npu():
+    if current_omni_platform.is_npu():
         kaiser_window = torch.kaiser_window(
             kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
         ).to("npu")
@@ -796,7 +796,7 @@ class UpSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
+        if current_omni_platform.is_npu():
             # TODO: When F.pad supports replicate mode on NPU, remove this branch
             input_dtype = hidden_states.dtype
             # F.pad in NPU doesn't support BF16 when mode is replicate.
@@ -843,7 +843,7 @@ class DownSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
+        if current_omni_platform.is_npu():
             input_dtype = hidden_states.dtype
             # F.pad in NPU doesn't support BF16 when mode is replicate.
             # To ensure the accuracy, manually pad the input tensor.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -35,7 +35,7 @@ from vllm.v1.sample.sampler import Sampler
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 # Special token IDs for Qwen3 Omni MoE
 # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -269,7 +269,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 1: Thinker ==========
         if self.model_stage == "thinker":
             thinker_dev = self._module_device(self.thinker)
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # Normalize to batched inputs if needed
                 _added_batch_dim = False
                 if input_ids is not None and input_ids.ndim == 1:
@@ -308,7 +308,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     "capture_layer_indices": [0, int(accept_layer)],
                     "return_hidden_states": True,
                 }
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # TODO: remove this hack when NPU supports batched inputs properly
                 thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids
                 thinker_inputs_embeds = (

vllm_omni/model_executor/stage_configs/bagel.yaml

@@ -9,7 +9,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: BagelForConditionalGeneration
-      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.4
       enforce_eager: true

vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml

@@ -9,7 +9,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: BagelForConditionalGeneration
-      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.4
       enforce_eager: true

vllm_omni/model_executor/stage_configs/qwen2_5_omni.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
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -40,7 +40,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
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -69,7 +69,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

vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.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
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -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
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -73,7 +73,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.3
       enforce_eager: true

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -14,7 +14,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
@@ -45,7 +45,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
@@ -76,7 +76,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

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml

@@ -14,7 +14,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
@@ -46,7 +46,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
@@ -77,7 +77,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

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml

@@ -13,7 +13,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
@@ -46,7 +46,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
@@ -82,7 +82,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

vllm_omni/model_executor/stage_configs/qwen3_tts.yaml

@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSForConditionalGeneration
-      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

vllm_omni/platforms/__init__.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import traceback
+from itertools import chain
+from typing import TYPE_CHECKING
+
+from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.utils.torch_utils import supports_xccl
+
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.plugins import (
+    OMNI_PLATFORM_PLUGINS_GROUP,
+    load_omni_plugins_by_group,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def cuda_omni_platform_plugin() -> str | None:
+    """Check if CUDA OmniPlatform should be activated."""
+    is_cuda = False
+    logger.debug("Checking if CUDA OmniPlatform is available.")
+    try:
+        from vllm.utils.import_utils import import_pynvml
+
+        pynvml = import_pynvml()
+        pynvml.nvmlInit()
+        try:
+            if pynvml.nvmlDeviceGetCount() > 0:
+                is_cuda = True
+                logger.debug("Confirmed CUDA OmniPlatform is available.")
+            else:
+                logger.debug("CUDA OmniPlatform is not available because no GPU is found.")
+        finally:
+            pynvml.nvmlShutdown()
+    except Exception as e:
+        logger.debug("CUDA OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.cuda.platform.CudaOmniPlatform" if is_cuda else None
+
+
+def rocm_omni_platform_plugin() -> str | None:
+    """Check if ROCm OmniPlatform should be activated."""
+    is_rocm = False
+    logger.debug("Checking if ROCm OmniPlatform is available.")
+    try:
+        import amdsmi
+
+        amdsmi.amdsmi_init()
+        try:
+            if len(amdsmi.amdsmi_get_processor_handles()) > 0:
+                is_rocm = True
+                logger.debug("Confirmed ROCm OmniPlatform is available.")
+            else:
+                logger.debug("ROCm OmniPlatform is not available because no GPU is found.")
+        finally:
+            amdsmi.amdsmi_shut_down()
+    except Exception as e:
+        logger.debug("ROCm OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.rocm.platform.RocmOmniPlatform" if is_rocm else None
+
+
+def npu_omni_platform_plugin() -> str | None:
+    """Check if NPU OmniPlatform should be activated."""
+    is_npu = False
+    logger.debug("Checking if NPU OmniPlatform is available.")
+    try:
+        import torch
+
+        if hasattr(torch, "npu") and torch.npu.is_available():
+            is_npu = True
+            logger.debug("Confirmed NPU OmniPlatform is available.")
+    except Exception as e:
+        logger.debug("NPU OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.npu.platform.NPUOmniPlatform" if is_npu else None
+
+
+def xpu_omni_platform_plugin() -> str | None:
+    """Check if XPU OmniPlatform should be activated."""
+    is_xpu = False
+    logger.debug("Checking if XPU OmniPlatform is available.")
+    try:
+        # installed IPEX if the machine has XPUs.
+        import intel_extension_for_pytorch  # noqa: F401
+        import torch
+
+        if supports_xccl():
+            dist_backend = "xccl"
+        else:
+            dist_backend = "ccl"
+            import oneccl_bindings_for_pytorch  # noqa: F401
+
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            is_xpu = True
+            from vllm_omni.platforms.xpu import XPUOmniPlatform
+
+            XPUOmniPlatform.dist_backend = dist_backend
+            logger.debug("Confirmed %s backend is available.", XPUOmniPlatform.dist_backend)
+            logger.debug("Confirmed XPU platform is available.")
+    except Exception as e:
+        logger.debug("XPU omni platform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.xpu.platform.XPUOmniPlatform" if is_xpu else None
+
+
+builtin_omni_platform_plugins = {
+    "cuda": cuda_omni_platform_plugin,
+    "rocm": rocm_omni_platform_plugin,
+    "npu": npu_omni_platform_plugin,
+    "xpu": xpu_omni_platform_plugin,
+}
+
+
+def resolve_current_omni_platform_cls_qualname() -> str:
+    """Resolve the current OmniPlatform class qualified name."""
+    platform_plugins = load_omni_plugins_by_group(OMNI_PLATFORM_PLUGINS_GROUP)
+
+    activated_plugins = []
+
+    for name, func in chain(builtin_omni_platform_plugins.items(), platform_plugins.items()):
+        try:
+            assert callable(func)
+            platform_cls_qualname = func()
+            if platform_cls_qualname is not None:
+                activated_plugins.append(name)
+        except Exception:
+            pass
+
+    activated_builtin_plugins = list(set(activated_plugins) & set(builtin_omni_platform_plugins.keys()))
+    activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys()))
+
+    if len(activated_oot_plugins) >= 2:
+        raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_oot_plugins}")
+    elif len(activated_oot_plugins) == 1:
+        platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]()
+        logger.info("OmniPlatform plugin %s is activated", activated_oot_plugins[0])
+    elif len(activated_builtin_plugins) >= 2:
+        raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_builtin_plugins}")
+    elif len(activated_builtin_plugins) == 1:
+        platform_cls_qualname = builtin_omni_platform_plugins[activated_builtin_plugins[0]]()
+        logger.debug("Automatically detected OmniPlatform %s.", activated_builtin_plugins[0])
+    else:
+        platform_cls_qualname = "vllm_omni.platforms.interface.UnspecifiedOmniPlatform"
+        logger.debug("No platform detected, vLLM-Omni is running on UnspecifiedOmniPlatform")
+
+    return platform_cls_qualname
+
+
+_current_omni_platform = None
+_init_trace: str = ""
+
+if TYPE_CHECKING:
+    current_omni_platform: OmniPlatform
+
+
+def __getattr__(name: str):
+    if name == "current_omni_platform":
+        # Lazy init current_omni_platform
+        global _current_omni_platform
+        if _current_omni_platform is None:
+            platform_cls_qualname = resolve_current_omni_platform_cls_qualname()
+            _current_omni_platform = resolve_obj_by_qualname(platform_cls_qualname)()
+            global _init_trace
+            _init_trace = "".join(traceback.format_stack())
+        return _current_omni_platform
+    elif name in globals():
+        return globals()[name]
+    else:
+        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
+
+
+def __setattr__(name: str, value):  # noqa: N807
+    if name == "current_omni_platform":
+        global _current_omni_platform
+        _current_omni_platform = value
+    elif name in globals():
+        globals()[name] = value
+    else:
+        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
+
+
+__all__ = [
+    "OmniPlatform",
+    "OmniPlatformEnum",
+    "current_omni_platform",
+    "_init_trace",
+]

vllm_omni/platforms/cuda/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+__all__ = ["CudaOmniPlatform"]

vllm_omni/platforms/cuda/platform.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.cuda import CudaPlatformBase
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
+    """CUDA/GPU implementation of OmniPlatform (default).
+
+    Inherits all CUDA-specific implementations from vLLM's CudaPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.CUDA
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/model_executor/stage_configs"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        # Check compute capability for Flash Attention support
+        # Flash Attention requires compute capability >= 8.0 and < 10.0
+        compute_capability = cls.get_device_capability()
+        flash_attn_supported = False
+        if compute_capability is not None:
+            major, minor = compute_capability
+            capability = major * 10 + minor
+            flash_attn_supported = 80 <= capability < 100
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
+                logger.warning(
+                    "Flash Attention requires GPU with compute capability >= 8.0 "
+                    "and < 10.0. Falling back to TORCH_SDPA backend."
+                )
+                logger.info("Defaulting to diffusion attention backend SDPA")
+                return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        if flash_attn_supported:
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("cuda")
+        return torch.device("cuda", local_rank)
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.cuda.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        return torch.version.cuda
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.cuda.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.cuda.mem_get_info(device)
+        return free
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)

vllm_omni/platforms/interface.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+
+import torch
+from vllm.platforms import Platform
+
+
+class OmniPlatformEnum(Enum):
+    """Enum for supported Omni platforms."""
+
+    CUDA = "cuda"
+    ROCM = "rocm"
+    NPU = "npu"
+    XPU = "xpu"
+    UNSPECIFIED = "unspecified"
+
+
+class OmniPlatform(Platform):
+    """
+    Abstract base class for vllm-omni Platform.
+
+    Inherits from vLLM's Platform and adds Omni-specific interfaces.
+    This gives OmniPlatform all vLLM Platform capabilities plus
+    Omni-specific methods.
+    """
+
+    _omni_enum: OmniPlatformEnum
+
+    def is_npu(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.NPU
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        """Get the diffusion attention backend class path for this platform.
+
+        This method selects the appropriate attention backend for diffusion
+        models based on platform capabilities and user preferences.
+
+        Args:
+            selected_backend: User-selected backend name (e.g., "FLASH_ATTN",
+                "TORCH_SDPA", "SAGE_ATTN"). If None, uses platform default.
+            head_size: Attention head size.
+
+        Returns:
+            Fully qualified class path of the selected backend.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        raise NotImplementedError
+
+    @classmethod
+    def synchronize(cls) -> None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        raise NotImplementedError
+
+
+class UnspecifiedOmniPlatform(OmniPlatform):
+    _omni_enum = OmniPlatformEnum.UNSPECIFIED
+    device_type = ""

vllm_omni/platforms/npu/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+__all__ = ["NPUOmniPlatform"]

vllm_omni/platforms/npu/platform.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm_ascend.platform import NPUPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class NPUOmniPlatform(OmniPlatform, NPUPlatform):
+    """NPU/Ascend implementation of OmniPlatform.
+
+    Inherits all NPU-specific implementations from vllm-ascend's NPUPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.NPU
+    dist_backend: str = "hccl"
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker"
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/npu/stage_configs"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        from importlib.util import find_spec
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        # Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA
+        if find_spec("mindiesd"):
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Falling back to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("npu")
+        return torch.device("npu", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.npu.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.npu.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.npu.mem_get_info(device)
+        return free
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.npu.get_device_properties(device_id)
+        return device_props.total_memory

vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml

@@ -9,10 +9,10 @@ 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
       gpu_memory_utilization: 0.8
-      enforce_eager: true  # Now we only support eager mode
+      enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
@@ -36,10 +36,10 @@ 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
       gpu_memory_utilization: 0.8
-      enforce_eager: true
+      enforce_eager: false
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: latent
@@ -58,12 +58,12 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       process: true
-      devices: "0"            # Example: use a different NPU than the previous stage; use "0" if single NPU
+      devices: "2"            # Example: use a different NPU than the previous stage; use "0" if single NPU
       max_batch_size: 1
     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

vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml

@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.2
        enforce_eager: false
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml

@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -37,7 +37,7 @@ from vllm_ascend.utils import ProfileExecuteDuration
 
 from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 
 
 class ExecuteModelState(NamedTuple):

vllm_omni/platforms/npu/worker/npu_ar_worker.py

@@ -4,10 +4,11 @@
 from vllm.v1.worker.workspace import init_workspace_manager
 from vllm_ascend.worker.worker import NPUWorker
 
-from vllm_omni.worker.npu.npu_ar_model_runner import NPUARModelRunner
+from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUARWorker(NPUWorker):
+class NPUARWorker(OmniWorkerMixin, NPUWorker):
     """NPU AR worker for thinker/talker stages in Omni model."""
 
     def init_device(self):

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -27,8 +27,8 @@ from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable
 
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.npu.npu_ar_model_runner import ExecuteModelState
-from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 
 
 class NPUGenerationModelRunner(OmniNPUModelRunner):

vllm_omni/platforms/npu/worker/npu_generation_worker.py

@@ -4,10 +4,11 @@
 from vllm.v1.worker.workspace import init_workspace_manager
 from vllm_ascend.worker.worker import NPUWorker
 
-from vllm_omni.worker.npu.npu_generation_model_runner import NPUGenerationModelRunner
+from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUGenerationWorker(NPUWorker):
+class NPUGenerationWorker(OmniWorkerMixin, NPUWorker):
     """NPU generation worker for code2wav stage in Omni model."""
 
     def init_device(self):

vllm_omni/platforms/rocm/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.rocm.platform import RocmOmniPlatform
+
+__all__ = ["RocmOmniPlatform"]

vllm_omni/platforms/rocm/platform.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.rocm import RocmPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class RocmOmniPlatform(OmniPlatform, RocmPlatform):
+    """ROCm/AMD GPU implementation of OmniPlatform.
+
+    Inherits all ROCm-specific implementations from vLLM's RocmPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.ROCM
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        from vllm._aiter_ops import is_aiter_found_and_supported
+
+        # Check if aiter is available for Flash Attention support
+        # aiter currently only is supported on gfx942 and gfx950
+        # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py
+        compute_capability = torch.cuda.get_device_capability()
+        major, minor = compute_capability
+        capability = major * 10 + minor
+        aiter_supported = is_aiter_found_and_supported() and 90 < capability < 100
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            if backend_upper == "FLASH_ATTN" and not aiter_supported:
+                logger.warning(
+                    "Flash Attention requires `aiter` library which is only supported "
+                    "on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
+                )
+                logger.info("Defaulting to diffusion attention backend SDPA")
+                return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        # Choose to enable Flash Attention by default on ROCm
+        # whenever possible as it is the fastest backend
+        if aiter_supported:
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/rocm/stage_configs"
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("cuda")
+        return torch.device("cuda", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.cuda.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        if torch.version.hip is not None:
+            hip_version = torch.version.hip
+            return hip_version.split("-")[0]
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.cuda.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.cuda.mem_get_info(device)
+        return free

vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml

@@ -10,7 +10,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
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -38,7 +38,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
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -66,7 +66,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

vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.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: true
@@ -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

vllm_omni/platforms/xpu/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+__all__ = ["XPUOmniPlatform"]

vllm_omni/platforms/xpu/platform.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.xpu import XPUPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class XPUOmniPlatform(OmniPlatform, XPUPlatform):
+    """XPU/Intel GPU implementation of OmniPlatform.
+
+    Inherits all XPU-specific implementations from vLLM's XPUPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.XPU
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/xpu/stage_configs"
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("xpu")
+        return torch.device("xpu", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.xpu.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        # XPU does not have a version string like CUDA
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.xpu.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        if device is None:
+            device_id = 0
+        else:
+            device_id = device.index if device.index is not None else 0
+        props = torch.xpu.get_device_properties(device_id)
+        return props.total_memory

vllm_omni/plugins/__init__.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+import vllm.envs as envs
+
+logger = logging.getLogger(__name__)
+
+# Default plugins group will be loaded in all processes (process0, engine core
+# process and worker processes).
+OMNI_DEFAULT_PLUGINS_GROUP = "vllm_omni.general_plugins"
+# Platform plugins group will be loaded in all processes when
+# `vllm_omni.platforms.current_omni_platform` is called and the value not
+# initialized.
+OMNI_PLATFORM_PLUGINS_GROUP = "vllm_omni.platform_plugins"
+
+# Make sure one process only loads plugins once.
+omni_plugins_loaded = False
+
+
+def load_omni_plugins_by_group(group: str) -> dict[str, Callable[[], Any]]:
+    from importlib.metadata import entry_points
+
+    allowed_plugins = envs.VLLM_PLUGINS
+
+    discovered_plugins = entry_points(group=group)
+    if len(discovered_plugins) == 0:
+        logger.debug("No plugins for group %s found.", group)
+        return {}
+
+    # Check if the only discovered plugin is the default one.
+    is_default_group = group == OMNI_DEFAULT_PLUGINS_GROUP
+    # Use INFO for non-default groups and DEBUG for the default group.
+    log_level = logger.debug if is_default_group else logger.info
+
+    log_level("Available plugins for group %s:", group)
+    for plugin in discovered_plugins:
+        log_level("- %s -> %s", plugin.name, plugin.value)
+
+    if allowed_plugins is None:
+        log_level("All plugins in this group will be loaded. Set `VLLM_PLUGINS` to control which plugins to load.")
+
+    plugins: dict[str, Callable[[], Any]] = {}
+    for plugin in discovered_plugins:
+        if allowed_plugins is None or plugin.name in allowed_plugins:
+            if allowed_plugins is not None:
+                log_level("Loading plugin %s", plugin.name)
+
+            try:
+                func = plugin.load()
+                plugins[plugin.name] = func
+            except Exception:
+                logger.exception("Failed to load plugin %s", plugin.name)
+
+    return plugins
+
+
+def load_omni_general_plugins() -> None:
+    """WARNING: plugins can be loaded for multiple times in different
+    processes. They should be designed in a way that they can be loaded
+    multiple times without causing issues.
+    """
+    global omni_plugins_loaded
+    if omni_plugins_loaded:
+        return
+    omni_plugins_loaded = True
+
+    plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP)
+    # General plugins: we only need to execute the loaded functions.
+    for func in plugins.values():
+        func()
+
+
+__all__ = [
+    "load_omni_general_plugins",
+    "OMNI_DEFAULT_PLUGINS_GROUP",
+    "OMNI_PLATFORM_PLUGINS_GROUP",
+]

vllm_omni/utils/__init__.py

@@ -1,13 +0,0 @@
-from vllm_omni.utils.platform_utils import (
-    detect_device_type,
-    get_device_control_env_var,
-    is_npu,
-    is_rocm,
-)
-
-__all__ = [
-    "detect_device_type",
-    "get_device_control_env_var",
-    "is_npu",
-    "is_rocm",
-]

vllm_omni/utils/platform_utils.py

@@ -1,58 +0,0 @@
-from __future__ import annotations
-
-import torch
-from vllm.platforms import current_platform
-
-
-def detect_device_type() -> str:
-    device_type = getattr(current_platform, "device_type", None)
-    if isinstance(device_type, str) and device_type:
-        return device_type.lower()
-    if torch.cuda.is_available():
-        return "cuda"
-    if hasattr(torch, "npu") and torch.npu.is_available():  # type: ignore[attr-defined]
-        return "npu"
-    return "cpu"
-
-
-def is_npu() -> bool:
-    return detect_device_type() == "npu"
-
-
-def is_rocm() -> bool:
-    return current_platform.is_rocm()
-
-
-def get_device_control_env_var() -> str:
-    """Return the environment variable name for device visibility control."""
-    if hasattr(current_platform, "device_control_env_var"):
-        env_var = getattr(current_platform, "device_control_env_var", None)
-        if isinstance(env_var, str) and env_var:
-            return env_var
-
-    device_type = detect_device_type()
-    if device_type == "npu":
-        return "ASCEND_RT_VISIBLE_DEVICES"
-    return "CUDA_VISIBLE_DEVICES"  # fallback
-
-
-def get_diffusion_worker_class() -> type:
-    """Get the appropriate diffusion WorkerProc class based on current device type.
-
-    Returns:
-        The WorkerProc class for the detected device type (either NPUWorkerProc or WorkerProc).
-
-    Raises:
-        ImportError: If the worker module for the detected device is not available.
-    """
-    device_type = detect_device_type()
-
-    if device_type == "npu":
-        from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorkerProc
-
-        return NPUWorkerProc
-    else:
-        # Default to GPU worker for cuda and other devices
-        from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc
-
-        return WorkerProc