86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Diffusion attention backend selector.
|
|
|
|
This module provides the interface for selecting diffusion attention backends.
|
|
The actual backend selection logic is delegated to the platform layer
|
|
(vllm_omni.platforms), similar to how vLLM handles attention backend selection.
|
|
|
|
Usage:
|
|
from vllm_omni.diffusion.attention.selector import get_attn_backend
|
|
|
|
# Get the appropriate backend for current platform
|
|
backend_cls = get_attn_backend(head_size=64)
|
|
|
|
# Or override via environment variable
|
|
# export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
|
"""
|
|
|
|
import importlib
|
|
import os
|
|
from functools import cache
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import (
|
|
AttentionBackend,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
|
|
"""Load a backend class from its fully qualified path.
|
|
|
|
Args:
|
|
cls_path: Fully qualified class path (e.g.,
|
|
"vllm_omni.diffusion.attention.backends.sdpa.SDPABackend")
|
|
|
|
Returns:
|
|
The loaded backend class
|
|
"""
|
|
module_path, class_name = cls_path.rsplit(".", 1)
|
|
try:
|
|
module = importlib.import_module(module_path)
|
|
backend_class = getattr(module, class_name)
|
|
return backend_class
|
|
except ImportError as e:
|
|
raise ImportError(f"Failed to import module {module_path}: {e}")
|
|
except AttributeError as e:
|
|
raise AttributeError(f"Class {class_name} not found in module: {e}")
|
|
|
|
|
|
@cache
|
|
def get_attn_backend(head_size: int) -> type[AttentionBackend]:
|
|
"""
|
|
Get attention backend for diffusion models.
|
|
|
|
The backend selection is delegated to the current platform
|
|
(vllm_omni.platforms.current_omni_platform), which selects the
|
|
appropriate backend based on:
|
|
1. User override via DIFFUSION_ATTENTION_BACKEND environment variable
|
|
2. Platform-specific defaults and capabilities
|
|
|
|
This is similar to how vLLM's get_attn_backend_cls works, where the
|
|
platform layer decides which backend to use based on hardware capabilities.
|
|
|
|
Args:
|
|
head_size: Head size for attention computation (may affect backend selection)
|
|
|
|
Returns:
|
|
The selected attention backend class
|
|
"""
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
# Check environment variable for user override
|
|
selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
|
|
|
|
# Delegate to platform for backend selection
|
|
backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
|
|
selected_backend=selected_backend,
|
|
head_size=head_size,
|
|
)
|
|
|
|
return _load_backend_cls(backend_cls_path)
|