[RFC]: Hardware Abstraction Layer & Plugin System for Unified Backend Support

[faaany]

Motivation

Currently,  vllm-omni supports multiple hardware backends (CUDA, NPU, XPU) through hardcoded paths and hardware-specific  if/else checks scattered throughout the codebase (e.g., stage_configs and  worker directories). This creates several issues:

• Coupling: Core modeling logic (like qwen2_5_omni_token2wav.py ) is polluted with device-specific code (e.g.,  is_npu()  checks for torch.kaiser_window kernels).
• Maintenance: Adding a new backend requires invasive changes to the core repository.
• Redundancy: We are reimplementing platform detection logic that already exists in upstream vLLM.

This proposal transitions vllm-omni into a modular architecture where the core engine remains hardware-blind, delegating device-specific logic to an OmniPlatform layer.

Proposed Architecture

A. OmniPlatform Hardware Abstraction Layer
Instead of ad-hoc detection, we unify all hardware-aware implementations into a specialized platform layer that inherits from upstream vLLM.

• Upstream Alignment: OmniPlatform inherits directly from vllm.platforms.Platform. Concrete implementations (e.g. XPUOmniPlatform, NPUOmniPlatform) extends vLLM Platform behavior with Omni-specific APIs
• Generic Device Handling: remove is_xpu()/is_npu() checks and device-specific strings from core modeling files. Instead, it will use current_omni_platform APIs to dynamically retrieve device metadata, worker classes, attention backends etc
• Custom Op Dispatch: Operators will dispatch via platform-owned selectors (e.g. current_omni_platform.is_npu()) to keep device knowledge centralized

B. The Plugin System
We follow the vLLM plugin structure but introduce Omni-specific groups to handle both in-tree and out-of-tree backends.

• Modular Repositories: Platform-specific code (Workers, ModelRunners, and specialized kernels) will be encapsulated in dedicated platform directories.
• Platform-Agnostic Configuration: YAML stage configs will be simplified. Instead of explicitly specifying  vllm_omni.worker.xpu.xpu_ar_worker.XPUARWorker , current_omni_platform dynamically resolves the correct hardware-specific implementation at runtime.
• Registration via Entry Points: We will utilize two distinct Python entry point groups to mirror vLLM's loading behavior: vllm_omni.general_plugins and vllm_omni.platform_plugins.

Implementation Details

see @gcanlin 's post below.

Feedback Period.

No response

CC List.

No response

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[gcanlin]

Hi! I'm focusing on developing for adapting NPU hardware now and also interested in implementing plugin systems like vLLM. In the past, we first introduced a plugin-based mechanism through vLLM-Ascend to integrate with vLLM. After a year of iteration, we believe that the successful experience from vllm-ascend can also benefit vllm-omni. Therefore, we would like to participate in the implementation of the plugin system. Would it be possible for us to work on this together?

[hsliuustc0106]

I think you can work together

[faaany]

Hi! I'm focusing on developing for adapting NPU hardware now and also interested in implementing plugin systems like vLLM. In the past, we first introduced a plugin-based mechanism through vLLM-Ascend to integrate with vLLM. After a year of iteration, we believe that the successful experience from vllm-ascend can also benefit vllm-omni. Therefore, we would like to participate in the implementation of the plugin system. Would it be possible for us to work on this together?

sure! Let's co-work on it. I think I found you in Slack. Let's talk offline.

[gcanlin]

vLLM-Omni plugin system design(WIP)

In general, we aim to provide a unified interface. When implementations vary across hardware, the differences should be encapsulated in the platform layer. Each hardware platform defines its own hardware-aware implementations, allowing developers to stay hardware-agnostic and rely on calls like current_omni_platform.get_xxx() to automatically dispatch to the correct platform-specific implementation.

Overview

vLLM-Omni’s platform inherits from vLLM’s platform because Omni includes AR models. We directly inherit it to avoid code duplication. However, we still need a new platform to unify hardware-specific code for the diffusion module.

Directories Structure

vllm_omni/

├── platforms/                    
│  ├── __init__.py             
│  ├── cuda/
│   │   ├── __init__.py
│   │   └── platform.py
│  ├── rocm/
│   │   ├── __init__.py
│   │   └── platform.py
│   │   └── stage_configs
│   ├── npu/
│   │   ├── __init__.py
│   │   └── platform.py
│   │   └── stage_configs
│   │   └── worker
│   ├── xpu/
│   │   ├── __init__.py
│   │   └── platform.py
│   │   └── stage_configs
│   │   └── worker
│   └── interface.py   

├── plugins/                       
│   ├── __init__.py 

Platform

from abc import abstractmethod
from enum import Enum

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):
    _omni_enum: OmniPlatformEnum = OmniPlatformEnum.UNSPECIFIED

    # Omni-specific abstract interfaces.
    @classmethod
    @abstractmethod
    def get_omni_ar_worker_cls(cls) -> str:
        ...

    @classmethod
    @abstractmethod
    def get_omni_generation_worker_cls(cls) -> str:
        ...

    @classmethod
    @abstractmethod
    def get_default_stage_config_path(cls) -> str:
        raise NotImplementedError

    @classmethod
    @abstractmethod
    def get_diffusion_attn_backend_cls(cls) -> str:
        ...

    def is_npu(self) -> bool:
        return self._omni_enum == OmniPlatformEnum.NPU
    # Other hardware has been collected by vLLM's Platform

Concrete platforms combine Omni APIs with vLLM platform behavior:

from vllm.platforms.cuda import CudaPlatform

class CudaOmniPlatform(OmniPlatform, CudaPlatform):
    _omni_enum = OmniPlatformEnum.CUDA
    ...

from vllm_ascend.platform import NPUPlatform
class NPUOmniPlatform(OmniPlatform, NPUPlatform):
    _omni_enum = OmniPlatformEnum.NPU
    ...

This means NPUOmniPlatform is both an OmniPlatform (Omni APIs) and an NPUPlatform (device_type, env vars, compiler backend, memory ops).

Plugin

We will follow the design of plugin system in vLLM. But have the following key point to mention:

• vLLM-Omni defines its own entry point group (vllm_omni.platform_plugins) and adds Omni-only APIs for AR workers, diffusion workers, and stage configs.
• NPU is first-class in vLLM-Omni (in-tree) rather than an out-of-tree plugin, because Omni needs to ship NPU-specific model runners and workers.
• The detection/selection flow intentionally mirrors vLLM (NVML, amdsmi, torch availability) to stay consistent with vLLM behavior.
• vLLM-Omni keeps the vLLM Platform interface to stay compatible with AR models while adding Omni-specific hooks.

vLLM-Omni now follows the vLLM plugin loader structure with two groups:

• OMNI_DEFAULT_PLUGINS_GROUP (vllm_omni.general_plugins)
  • Loaded explicitly via load_omni_general_plugins().
  • Uses the same allowlist env var as vLLM: VLLM_PLUGINS.
    • None: load all plugins.
    • Empty string: load none.
• OMNI_PLATFORM_PLUGINS_GROUP (vllm_omni.platform_plugins)
  • Loaded lazily when current_omni_platform is first accessed.

Proposal Changes

Device detection

All hardware decisions go through current_omni_platform, and call sites use explicit platform APIs rather than re-detecting devices:

• From: ad-hoc detection and scattered helpers, To: current_omni_platform.device_type, current_omni_platform.device_control_env_var
• Use current_omni_platform.is_cuda() / is_npu() / is_rocm() for branching instead of introducing new helpers like get_device_name().

This reduces duplicated detection logic, aligns behavior across AR and diffusion pipelines, and keeps platform-specific behavior centralized in one place.

Plugin load points

vLLM loads plugins in several processes to ensure global patches and extensions apply everywhere:

• Process0 / CLI & config:
  • vllm/engine/arg_utils.py (EngineArgs.__post_init__, AsyncEngineArgs.add_cli_args) calls load_general_plugins().
• Engine core / scheduler process:
  • vllm/v1/engine/core.py (EngineCore.__init__) calls load_general_plugins().
• Worker processes:
  • vllm/v1/worker/worker_base.py (WorkerBase.init_worker) calls load_general_plugins().
• Model registry subprocess:
  • vllm/model_executor/models/registry.py (_run) calls load_general_plugins() before importing models.
• Platform plugins:
  • vllm/platforms/__init__.py loads PLATFORM_PLUGINS_GROUP lazily when current_platform is resolved.
• IO processor plugins (process0 only):
  • vllm/plugins/io_processors/__init__.py loads IO_PROCESSOR_PLUGINS_GROUP when get_io_processor() is called.
• Stat logger plugins (process0 async only):
  • vllm/v1/metrics/loggers.py loads STAT_LOGGER_PLUGINS_GROUP when custom stat loggers are requested.

For vLLM-Omni, would follow the similar loading point(need to test more):

• Process0 / CLI & config:

  • vllm_omni/engine/arg_utils.py
class OmniEngineArgs(EngineArgs) and
    class AsyncOmniEngineArgs(AsyncEngineArgs)

    def __post_init__(self):
        from vllm_omni.plugins import load_omni_general_plugins
    
        load_omni_general_plugins()
        super().__post_init__()

• Stage worker processes (LLM):

  • vllm_omni/entrypoints/omni_stage.py
_stage_worker(...) and _stage_worker_async(...)

    def _stage_worker(...):
        from vllm_omni.plugins import load_omni_general_plugins
    
        load_omni_general_plugins()
        ...

• Diffusion worker processes:

  • vllm_omni/diffusion/worker/gpu_worker.py and
    vllm_omni/diffusion/worker/npu/npu_worker.py
    (inside the worker entrypoint before model import or init)

    def worker_main(...):
        from vllm_omni.plugins import load_omni_general_plugins
    
        load_omni_general_plugins()
        ...

• Model registry subprocess (if used):

  • If Omni ever uses a model-registry subprocess like vLLM does, call load_omni_general_plugins() before importing models.

• LLM worker classes (AR / Generation):

  • vllm_omni/worker/mixins.py defines OmniWorkerMixin, which loads plugins in __init__.
  • Omni LLM workers inherit this mixin:
    • vllm_omni/worker/gpu_ar_worker.py
    • vllm_omni/worker/gpu_generation_worker.py
    • vllm_omni/worker/npu/npu_ar_worker.py
    • vllm_omni/worker/npu/npu_generation_worker.py

This keeps Omni behavior aligned with vLLM while remaining minimal: only OMNI_DEFAULT_PLUGINS_GROUP and OMNI_PLATFORM_PLUGINS_GROUP are needed.

Stage config

We want stage configs to avoid hardcoding platform-specific worker class names such as:

worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker

Instead, stage configs should specify a worker type (e.g., ar orgeneration), and the runtime will resolve the concrete worker class using the current platform:

worker_type: ar  # or generation

Resolution logic:

if worker_type == "ar":
    worker_cls = current_omni_platform.get_omni_ar_worker_cls()
elif worker_type == "generation":
    worker_cls = current_omni_platform.get_omni_generation_worker_cls()

Benefits:

• Stage configs become platform-agnostic and easier to maintain.
• Adding a new platform does not require editing every stage config file.
• Platform-specific worker mapping stays centralized in OmniPlatform.

Stage config lookup can also be simplified: get config from current_omni_platform.get_default_stage_config_path first. Otherwise, fallback to vllm_omni/model_executor/stage_configs/{model_type}.yaml.

class NPUOmniPlatform(OmniPlatform, NPUPlatform):

    @classmethod
    def get_default_stage_config_path(cls) -> str:
        return "vllm_omni/platforms/npu/stage_configs"

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)

Layer Backend and custom ops

We want hardware-specific selection logic to live in the platform layer, so model and operator code stays hardware-agnostic. This has two parts.

1) Platform-owned hardware selectors

We will gradually identify and lift hardware selectors into OmniPlatform APIs. Each platform provides the default backend or strategy, and call sites use the platform API instead of hardcoding device checks.

Concrete examples to standardize:

• Diffusion ring attention:

  • Today selection is scattered across ring attention code paths.
  • Proposed API: get_ring_attn_backend().
  • The ring selector (e.g., vllm_omni/diffusion/attention/parallel/ring) should call the platform API and instantiate the returned backend.

• Diffusion attention backend:

  • Today selection mixes checks in attention selector logic.
  • Add get_diffusion_attn_backend_cls interfaces. Extract the selecting logic in diffusion/attention/selector.py into get_diffusion_attn_backend_cls
  • Every hardware platform will implement their custom diffusion attention layer backend. The diffusion attention selector should resolve the backend via platform without checking CUDA/NPU/ROCm directly.

• more hardware-related dispatcher....

They will be like:

class OmniPlatform(Platform):
    @classmethod
    def get_ring_attn_backend(cls) -> str:
        ...

    @classmethod
    def get_diffusion_attn_backend(cls) -> str:
        ...

Each platform returns a fully qualified class name or backend identifier appropriate for its hardware.

2) Custom op dispatch via platform

Custom ops should dispatch via current_omni_platform.is_xxx() rather than doing their own hardware probing (torch.cuda.is_available, string compares, etc.). This keeps all device knowledge in one place.

Example pattern:

from vllm_omni.platforms import current_omni_platform

def dispatch_forward(self) -> Callable:
  ...
  if current_platform.is_rocm():
      return self.forward_hip
  elif current_platform.is_npu():
      return self.forward_npu
  else:
      return self.forward_cuda

Diffusion Worker

• Remove diffusion/npu/npu_worker.py.
• Extract the API about device into platform.

class DiffusionWorker:
    def init_device_and_model(self) -> None:
        """Initialize the device and load the model."""
        ....
        # Use platform API for device initialization
        self.device = current_omni_platform.get_torch_device(rank)
        current_omni_platform.set_device(self.device)