Skip to content

vllm_omni.platforms.cuda

Modules:

Name Description
platform

CudaOmniPlatform

Bases: OmniPlatform, CudaPlatformBase

CUDA/GPU implementation of OmniPlatform (default).

Inherits all CUDA-specific implementations from vLLM's CudaPlatform, and adds Omni-specific interfaces from OmniPlatform.

get_default_ir_op_priority classmethod 

get_default_ir_op_priority(
    vllm_config: VllmConfig,
) -> IrOpPriorityConfig

Prefer vllm_c CUDA kernels over native for diffusion IR ops.

get_default_stage_config_path classmethod 

get_default_stage_config_path() -> str

get_device_capability classmethod 

get_device_capability(
    device_id: int = 0,
) -> DeviceCapability | None

get_device_count classmethod 

get_device_count() -> int

get_device_memory classmethod 

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

get_device_name classmethod 

get_device_name(device_id: int = 0) -> str

get_device_total_memory classmethod 

get_device_total_memory(device_id: int = 0) -> 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 = True,
) -> str

get_free_memory classmethod 

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

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_torch_device classmethod 

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

has_flash_attn_4 classmethod 

has_flash_attn_4() -> bool

Return whether CuTe FA4 is importable (Blackwell-capable FLASH_ATTN).

has_flash_attn_package classmethod 

has_flash_attn_package() -> bool

is_fully_connected classmethod 

is_fully_connected(device_ids: list[int]) -> bool

memory_reserved classmethod 

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

record_device_event classmethod 

record_device_event() -> Event | None

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

supports_diffusion_dense_flash_attention classmethod 

supports_diffusion_dense_flash_attention() -> bool

Dense FLASH_ATTN is FA4 on Blackwell; FA2/FA3 on older CUDA GPUs.

supports_torch_inductor classmethod 

supports_torch_inductor() -> bool

synchronize classmethod 

synchronize() -> None