Skip to content

vllm_omni.platforms

Modules:

Name Description
cuda
interface
musa
npu
rocm
xpu

builtin_omni_platform_plugins module-attribute 

builtin_omni_platform_plugins = {
    "cuda": cuda_omni_platform_plugin,
    "rocm": rocm_omni_platform_plugin,
    "npu": npu_omni_platform_plugin,
    "xpu": xpu_omni_platform_plugin,
    "musa": musa_omni_platform_plugin,
}

current_omni_platform module-attribute 

current_omni_platform: OmniPlatform

logger module-attribute 

logger = logging.getLogger(__name__)

OmniPlatform

Bases: 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.

build_diffusion_kv_attn_metadata classmethod 

build_diffusion_kv_attn_metadata(
    **kwargs: Any,
) -> dict[str, Any]

Build native attention metadata for the diffusion paged path.

This default is the GPU/common path and uses vLLM's builder. NPU overrides it to build Ascend attention metadata without making the shared diffusion adapter import vllm_ascend. seq_lens_cpu is an adapter-only convenience value and is removed before calling the upstream builder.

configure_diffusion_vllm_config classmethod 

configure_diffusion_vllm_config(
    vllm_config: Any, od_config: Any
) -> None

Apply platform-specific native cache geometry for diffusion.

create_autocast_context classmethod 

create_autocast_context(
    *, device_type: str, dtype: dtype, enabled: bool = True
)

get_default_stage_config_path classmethod 

get_default_stage_config_path() -> str

get_device_count classmethod 

get_device_count() -> int

get_device_memory classmethod 

get_device_memory(
    device: device | None = None,
) -> tuple[int, int]

get_device_version classmethod 

get_device_version() -> str | None

get_diffusion_attn_backend_cls classmethod 

get_diffusion_attn_backend_cls(
    selected_backend: str | None,
    head_size: int,
    allow_trtllm_default: bool = False,
) -> 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.

Parameters:

Name Type Description Default
selected_backend str | None

User-selected backend name (e.g., "FLASH_ATTN", "TORCH_SDPA", "SAGE_ATTN"). If None, uses platform default.

 required
head_size int

Attention head size.

 required
allow_trtllm_default bool

Whether TRTLLM may be chosen as the default.

 False

Returns:

Type Description
str

Fully qualified class path of the selected backend.

get_diffusion_kv_block_tables_cls classmethod 

get_diffusion_kv_block_tables_cls() -> type

Return the platform's native paged-KV BlockTables implementation.

get_diffusion_model_runner_cls classmethod 

get_diffusion_model_runner_cls() -> str

Get the diffusion model runner class path for this platform.

Returns a fully qualified class path string. The class must be compatible with the DiffusionModelRunner interface.

get_diffusion_packed_modules_mapping classmethod 

get_diffusion_packed_modules_mapping(
    model_class: type[Module],
) -> dict[str, list[str]] | None

get_diffusion_paged_kv_attn_backend classmethod 

get_diffusion_paged_kv_attn_backend(
    attn_backend: type, *, ulysses_degree: int
) -> type

Specialize a native paged backend for diffusion execution.

get_diffusion_worker_cls classmethod 

get_diffusion_worker_cls() -> str

Get the diffusion worker class path for this platform.

Returns a fully qualified class path string that will be resolved and instantiated by WorkerWrapperBase. The class must be compatible with the DiffusionWorker interface.

get_free_memory classmethod 

get_free_memory(device: device | None = None) -> int

get_graph_wrapper_cls classmethod 

get_graph_wrapper_cls() -> type

Return the platform's full-graph wrapper class.

Defaults to vLLM's CUDAGraphWrapper; NPU overrides with ACLGraphWrapper.

get_omni_ar_worker_cls classmethod 

get_omni_ar_worker_cls() -> str

get_omni_generation_worker_cls classmethod 

get_omni_generation_worker_cls() -> str

get_profiler_cls classmethod 

get_profiler_cls() -> str

Get the profiler class for this platform.

Returns:

Type Description
str

Fully qualified class path of the profiler.
str

Default returns the base OmniTorchProfilerWrapper.

get_torch_device classmethod 

get_torch_device(local_rank: int | None = None) -> device

has_flash_attn_package classmethod 

has_flash_attn_package() -> bool

Check if a Flash Attention package is available and usable on this platform.

init_diffusion_model_runner_runtime classmethod 

init_diffusion_model_runner_runtime(
    vllm_config: Any, od_config: Any, device: device
) -> None

Initialize platform-specific runtime state for diffusion model runners.

init_diffusion_worker_vllm_config classmethod 

init_diffusion_worker_vllm_config(vllm_config: Any) -> None

Initialize platform-specific state for diffusion worker VllmConfig.

is_cuda 

is_cuda() -> bool

is_musa 

is_musa() -> bool

is_npu 

is_npu() -> bool

is_out_of_tree 

is_out_of_tree() -> bool

is_rocm 

is_rocm() -> bool

is_xpu 

is_xpu() -> bool

memory_reserved classmethod 

memory_reserved(device: device | int | None = None) -> int

Bytes reserved by this process's caching allocator on device.

Device-wide get_device_memory counts every process; this is the portable equivalent of torch.cuda.memory_reserved.

prepare_diffusion_op_runtime classmethod 

prepare_diffusion_op_runtime(
    op_name: str, **kwargs: Any
) -> None

record_device_event classmethod 

record_device_event()

Record a device event on the default stream to mark tensor readiness.

On platforms where distributed communication (e.g. HCCL) may use internal streams not visible to the default stream, this method should synchronize the default stream before recording the event to ensure the event captures all completed work including cross-device communication results.

Returns None by default so that platforms without a native implementation (ROCm, XPU, MUSA) fall through to a safe no-op. Override in platform subclasses to provide real event support.

register_additional_diffusion_fused_moe_hooks classmethod 

register_additional_diffusion_fused_moe_hooks(
    moe_runner: Any,
) -> None

requires_diffusion_paged_kv_prewrite classmethod 

requires_diffusion_paged_kv_prewrite() -> bool

Whether paged attention must write K/V before native execution.

The default GPU path keeps cache-update ownership in its native attention call. Ascend overrides this because piecewise FIA should write the complete K/V span once, then read it from cache for every segment.

reset_diffusion_fused_moe_forward_context classmethod 

reset_diffusion_fused_moe_forward_context() -> None

set_device_control_env_var classmethod 

set_device_control_env_var(
    devices: str | int | None,
) -> None

set_forward_context classmethod 

set_forward_context(
    attn_metadata: Any,
    vllm_config: VllmConfig,
    *,
    cudagraph_runtime_mode: CUDAGraphMode,
    batch_descriptor: BatchDescriptor,
)

Platform-neutral wrapper around the device's set_forward_context.

Defaults to vLLM's set_forward_context; NPU overrides to dispatch to set_ascend_forward_context (renaming cudagraph_runtime_mode to aclgraph_runtime_mode).

supports_cpu_offload classmethod 

supports_cpu_offload() -> bool

supports_diffusion_dense_flash_attention classmethod 

supports_diffusion_dense_flash_attention() -> bool

Whether the platform's dense FLASH_ATTN dependencies exist.

supports_float64 classmethod 

supports_float64() -> bool

supports_talker_mtp_graph_capture classmethod 

supports_talker_mtp_graph_capture() -> bool

Whether a model may capture its dedicated talker MTP graph.

supports_torch_inductor classmethod 

supports_torch_inductor() -> bool

Check if the platform supports torch.compile with inductor backend.

synchronize classmethod 

synchronize() -> None

unset_device_control_env_var classmethod 

unset_device_control_env_var() -> None

validate_diffusion_attn_backend classmethod 

validate_diffusion_attn_backend(
    selected_backend: str,
) -> None

Reject an explicitly selected backend this platform cannot run.

Platforms call this from get_diffusion_attn_backend_cls so that a backend restricted to other hardware, or one whose kernel package is missing, fails during resolution rather than at the first forward.

OmniPlatformEnum

Bases: Enum

Enum for supported Omni platforms.

CUDA class-attribute instance-attribute 

CUDA = 'cuda'

MUSA class-attribute instance-attribute 

MUSA = 'musa'

NPU class-attribute instance-attribute 

NPU = 'npu'

OOT class-attribute instance-attribute 

OOT = 'oot'

ROCM class-attribute instance-attribute 

ROCM = 'rocm'

UNSPECIFIED class-attribute instance-attribute 

UNSPECIFIED = 'unspecified'

XPU class-attribute instance-attribute 

XPU = 'xpu'

cuda_omni_platform_plugin 

cuda_omni_platform_plugin() -> str | None

Check if CUDA OmniPlatform should be activated.

musa_omni_platform_plugin 

musa_omni_platform_plugin() -> str | None

Check if MUSA OmniPlatform should be activated.

npu_omni_platform_plugin 

npu_omni_platform_plugin() -> str | None

Check if NPU OmniPlatform should be activated.

resolve_current_omni_platform_cls_qualname 

resolve_current_omni_platform_cls_qualname() -> str

Resolve the current OmniPlatform class qualified name.

rocm_omni_platform_plugin 

rocm_omni_platform_plugin() -> str | None

Check if ROCm OmniPlatform should be activated.

xpu_omni_platform_plugin 

xpu_omni_platform_plugin() -> str | None

Check if XPU OmniPlatform should be activated.