Skip to content

vllm_omni.diffusion.attention.backends.registry

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().

logger module-attribute 

logger = init_logger(__name__)

DiffusionAttentionBackendEnum

Bases: Enum

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

CUDNN_ATTN class-attribute instance-attribute 

CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"

FLASHINFER_ATTN class-attribute instance-attribute 

FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"

FLASH_ATTN class-attribute instance-attribute 

FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"

SAGE_ATTN class-attribute instance-attribute 

SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"

SAGE_ATTN_3 class-attribute instance-attribute 

SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"

TORCH_SDPA class-attribute instance-attribute 

TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"

clear_override 

clear_override() -> None

Clear any override for this backend, reverting to the default.

get_class 

get_class() -> type[AttentionBackend]

Get the backend class (respects overrides).

Returns:

Type Description
type[AttentionBackend]

The backend class

Raises:

Type Description
ImportError

If the backend class cannot be imported
ValueError

If backend has empty path and is not registered

get_path 

get_path(include_classname: bool = True) -> str

Get the class path for this backend (respects overrides).

Returns:

Type Description
str

The fully qualified class path string

Raises:

Type Description
ValueError

If backend has empty path and is not registered

is_overridden 

is_overridden() -> bool

Check if this backend has been overridden.

Returns:

Type Description
bool

True if the backend has a registered override

register_diffusion_backend 

register_diffusion_backend(
    backend: DiffusionAttentionBackendEnum,
    class_path: str | None = None,
) -> Callable[[type], type]

Register or override a diffusion backend implementation.

Parameters:

Name Type Description Default
backend DiffusionAttentionBackendEnum

The DiffusionAttentionBackendEnum member to register

 required
class_path str | None

Optional class path. If not provided and used as decorator, will be auto-generated from the class.

 None

Returns:

Type Description
Callable[[type], type]

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" )