Skip to content

vllm_omni.platforms.musa.platform

logger module-attribute 

logger = init_logger(__name__)

MUSAOmniPlatform

Bases: OmniPlatform, MUSAPlatformBase

MUSA/Moore Threads GPU implementation of OmniPlatform.

Inherits all MUSA-specific implementations from vllm-musa's MUSAPlatformBase, and adds Omni-specific interfaces from OmniPlatform.

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 the compute capability of the MUSA device.

get_device_count classmethod 

get_device_count() -> int

Get the number of available MUSA devices.

get_device_name classmethod 

get_device_name(device_id: int = 0) -> str

get_device_version classmethod 

get_device_version() -> str | None

Get the MUSA runtime version.

get_diffusion_attn_backend_cls classmethod 

get_diffusion_attn_backend_cls(
    selected_backend: str | None, head_size: int
) -> str

Get the diffusion attention backend class path for MUSA platform.

MUSA supports FLASH_ATTN via the mate package, and SDPA as fallback.

Parameters:

Name Type Description Default
selected_backend str | None

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

 required
head_size int

Attention head size.

 required

Returns:

Type Description
str

Fully qualified class path of the selected backend.

get_free_memory classmethod 

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

Get the free memory on the MUSA device.

Parameters:

Name Type Description Default
device device | None

Optional device to query. If None, uses current device.

 None

Returns:

Type Description
int

Free memory in bytes.

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

Get the torch device for MUSA platform.

Parameters:

Name Type Description Default
local_rank int | None

Optional local rank for multi-GPU setups.

 None

Returns:

Type Description
device

torch.device for MUSA GPU.

has_flash_attn_package classmethod 

has_flash_attn_package() -> bool

supports_float64 classmethod 

supports_float64() -> bool

MUSA does not support float64 yet.

supports_torch_inductor classmethod 

supports_torch_inductor() -> bool

MUSA supports torch.compile with inductor backend.

synchronize classmethod 

synchronize() -> None

Synchronize all MUSA operations.