453 lines
21 KiB
Python
453 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import importlib
|
|
|
|
import torch
|
|
import vllm.envs as envs
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.kernel import IrOpPriorityConfig
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms.cuda import CudaPlatformBase
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
|
from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
|
|
"""CUDA/GPU implementation of OmniPlatform (default).
|
|
|
|
Inherits all CUDA-specific implementations from vLLM's CudaPlatform,
|
|
and adds Omni-specific interfaces from OmniPlatform.
|
|
"""
|
|
|
|
_omni_enum = OmniPlatformEnum.CUDA
|
|
|
|
@classmethod
|
|
def get_omni_ar_worker_cls(cls) -> str:
|
|
return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
|
|
|
|
@classmethod
|
|
def get_omni_generation_worker_cls(cls) -> str:
|
|
return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
|
|
|
|
@classmethod
|
|
def get_default_stage_config_path(cls) -> str:
|
|
return "vllm_omni/deploy"
|
|
|
|
@classmethod
|
|
def has_flash_attn_package(cls) -> bool:
|
|
from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed
|
|
|
|
# Turing/Tesla/T4 GPUs don't support flash attention well
|
|
gpu_name = cls.get_device_name()
|
|
if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
|
|
return False
|
|
|
|
if not is_flash_attn_installed():
|
|
return False
|
|
|
|
return True
|
|
|
|
@classmethod
|
|
def has_flash_attn_4(cls) -> bool:
|
|
"""Return whether CuTe FA4 is importable (Blackwell-capable FLASH_ATTN)."""
|
|
from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_4_available
|
|
|
|
return is_flash_attn_4_available()
|
|
|
|
@classmethod
|
|
def supports_diffusion_dense_flash_attention(cls) -> bool:
|
|
"""Dense FLASH_ATTN is FA4 on Blackwell; FA2/FA3 on older CUDA GPUs."""
|
|
capability = cls.get_device_capability()
|
|
if capability is not None and capability.major in (10, 11, 12):
|
|
return cls.has_flash_attn_4()
|
|
return cls.has_flash_attn_package()
|
|
|
|
@classmethod
|
|
def get_diffusion_attn_backend_cls(
|
|
cls,
|
|
selected_backend: str | None,
|
|
head_size: int,
|
|
allow_trtllm_default: bool = True,
|
|
) -> str:
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
|
|
# Check compute capability for Flash Attention support.
|
|
# FA requires sm_80+. Blackwell (sm_10x/sm_12x) only works with FA builds
|
|
# that include the Blackwell CUTE kernel — plain FA2 will crash there.
|
|
#
|
|
# Known Blackwell SKUs:
|
|
# sm_100 = B200 / GB200 (datacenter)
|
|
# sm_103 = B300 / GB300 (Blackwell Ultra)
|
|
# sm_120 = RTX Pro 6000, RTX 50-series (consumer)
|
|
# sm_121 = consumer Blackwell refresh
|
|
_known_blackwell_sms = {(10, 0), (10, 3), (12, 0), (12, 1)}
|
|
compute_capability = cls.get_device_capability()
|
|
compute_supported = False
|
|
is_blackwell = False
|
|
sm_str = ""
|
|
if compute_capability is not None:
|
|
major, minor = compute_capability
|
|
capability = major * 10 + minor
|
|
compute_supported = capability >= 80
|
|
sm_str = f"sm_{major}{minor}"
|
|
# Accept major in {10, 11, 12} to cover future Blackwell refreshes.
|
|
is_blackwell = major in (10, 11, 12)
|
|
if is_blackwell and (major, minor) not in _known_blackwell_sms:
|
|
logger.info(
|
|
"Detected Blackwell-class GPU %s (untested variant); routing to CUDNN_ATTN.",
|
|
sm_str,
|
|
)
|
|
|
|
# Check if FA packages are available
|
|
packages_info = PACKAGES_CHECKER.get_packages_info()
|
|
packages_available = packages_info.get("has_flash_attn", False)
|
|
|
|
# FA2/FA3 wheels are often importable on Blackwell but only ship
|
|
# Ampere/Ada/Hopper kernels; the first forward then dies with
|
|
# "no kernel image". FA4 (flash_attn.cute / vllm-omni[fa4]) is the
|
|
# FLASH_ATTN implementation that actually runs on Blackwell.
|
|
if is_blackwell:
|
|
flash_attn_supported = compute_supported and cls.has_flash_attn_4()
|
|
else:
|
|
flash_attn_supported = compute_supported and packages_available
|
|
|
|
# cuDNN 9.5+ ships Blackwell FMHA kernels. If the runtime is older,
|
|
# skip the CUDNN_ATTN default rather than selecting a backend whose
|
|
# kernel selector may reject the runtime shape.
|
|
cudnn_version = torch.backends.cudnn.version() or 0
|
|
cudnn_blackwell_ready = cudnn_version >= 90500
|
|
|
|
# FlashInfer edges cuDNN by ~4% at the kernel level on sm_120 but
|
|
# regresses ~2x at e2e on HV-1.5 because its dense-prefill path cannot
|
|
# represent every diffusion attention mask without changing semantics.
|
|
# CUDNN_ATTN pins sdpa_kernel([CUDNN_ATTENTION]) directly so masked
|
|
# calls keep the cuDNN path. Blackwell default prefers CUDNN_ATTN;
|
|
# users can opt into FLASHINFER_ATTN explicitly for no-mask workloads.
|
|
flashinfer_available = False
|
|
try:
|
|
import flashinfer # noqa: F401
|
|
|
|
flashinfer_available = True
|
|
except Exception as e:
|
|
# A partially installed / ABI-mismatched wheel can raise OSError or
|
|
# RuntimeError from extension loading, not just ImportError. This
|
|
# runs during default backend selection, so a probe failure must
|
|
# not abort startup — just treat FlashInfer as unavailable.
|
|
logger.debug("FlashInfer probe failed (%s); treating as unavailable", e)
|
|
|
|
# FLASHINFER_ATTN needs BatchPrefillWithRaggedKVCacheWrapper, not just a
|
|
# top-level flashinfer package. Probe the actual symbol so a stub /
|
|
# partial wheel is not auto-routed then crash at layer init.
|
|
flashinfer_prefill_available = False
|
|
if flashinfer_available:
|
|
try:
|
|
from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper # noqa: F401
|
|
|
|
flashinfer_prefill_available = True
|
|
except Exception as e:
|
|
logger.debug("FlashInfer prefill wrapper probe failed (%s); treating as unavailable", e)
|
|
|
|
# TRTLLM_ATTN needs the trtllm-gen kernel specifically, not just any FlashInfer
|
|
# wheel. Probe the actual symbol so a released wheel lacking it does not get
|
|
# auto-routed to TRTLLM_ATTN and then crash on the first forward.
|
|
trtllm_gen_available = False
|
|
if flashinfer_available:
|
|
try:
|
|
from flashinfer.prefill import trtllm_ragged_attention_deepseek # noqa: F401
|
|
|
|
trtllm_gen_available = True
|
|
except Exception as e:
|
|
logger.debug("trtllm-gen kernel probe failed (%s); treating as unavailable", e)
|
|
|
|
if selected_backend is not None:
|
|
backend_upper = selected_backend.upper()
|
|
# Architecture / package probes for optional kernels must run before
|
|
# validate_diffusion_attn_backend, which imports the backend module
|
|
# and would otherwise load sageattention / sageattn3 on unsupported GPUs.
|
|
if backend_upper == "SAGE_ATTN":
|
|
sage_supported_sms = {(8, 0), (8, 6), (8, 9), (9, 0), (12, 0), (12, 1)}
|
|
if compute_capability is None or tuple(compute_capability) not in sage_supported_sms:
|
|
raise ValueError(
|
|
f"SAGE_ATTN was explicitly selected but does not provide a kernel for {sm_str or 'this GPU'}. "
|
|
"Select a compatible backend."
|
|
)
|
|
try:
|
|
importlib.import_module("sageattention")
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"SAGE_ATTN was explicitly selected, but the sageattention package is not available."
|
|
) from e
|
|
if backend_upper == "SAGE_ATTN_3":
|
|
sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10
|
|
if not sage_attn3_supported:
|
|
raise ValueError(
|
|
"SAGE_ATTN_3 was explicitly selected but requires a Blackwell-class GPU "
|
|
"with compute capability >= 10.0. Select a compatible backend."
|
|
)
|
|
try:
|
|
importlib.import_module("sageattn3")
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"SAGE_ATTN_3 was explicitly selected, but the sageattn3 package is not available. "
|
|
"Install SageAttention/sageattention3_blackwell or select a different backend."
|
|
) from e
|
|
cls.validate_diffusion_attn_backend(backend_upper)
|
|
if backend_upper == "SAGE_ATTN_3" and head_size > 0:
|
|
sage3_cls = DiffusionAttentionBackendEnum.SAGE_ATTN_3.get_class()
|
|
if not sage3_cls.supports_head_size(head_size):
|
|
raise ValueError(
|
|
f"SAGE_ATTN_3 was explicitly selected but head_size={head_size} is unsupported. "
|
|
f"Supported head sizes: {sage3_cls.get_supported_head_sizes()}. "
|
|
"Select TORCH_SDPA or another backend."
|
|
)
|
|
if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
|
|
try:
|
|
importlib.import_module("kernels")
|
|
logger.info("Using HuggingFace kernels-backed attention backend '%s'", backend_upper)
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
f"{backend_upper} was explicitly selected, but the HuggingFace `kernels` "
|
|
"library is not available. Install `kernels` or select a different backend."
|
|
) from e
|
|
|
|
if backend_upper == "FLASH_ATTN_HUB":
|
|
fa2_hub_supported = compute_capability is not None and compute_capability.major in (8, 9)
|
|
if not fa2_hub_supported:
|
|
raise ValueError(
|
|
"FLASH_ATTN_HUB was explicitly selected but its current FA2 kernels require "
|
|
"an Ampere, Ada, or Hopper GPU (compute capability 8.x/9.x)."
|
|
)
|
|
|
|
if backend_upper == "FLASH_ATTN_3_HUB":
|
|
fa3_hub_supported = compute_capability is not None and compute_capability.major == 9
|
|
if not fa3_hub_supported:
|
|
raise ValueError(
|
|
"FLASH_ATTN_3_HUB was explicitly selected but its current kernels require a Hopper GPU "
|
|
"with compute capability 9.x. Select a compatible backend."
|
|
)
|
|
|
|
if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
|
|
# Same loader/cache as execution (Hub versions 1 then 2). A missing
|
|
# torch/CUDA variant used to degrade silently to native/SDPA (#6971).
|
|
hub_repo = (
|
|
"kernels-community/flash-attn3"
|
|
if backend_upper == "FLASH_ATTN_3_HUB"
|
|
else "kernels-community/flash-attn2"
|
|
)
|
|
try:
|
|
from vllm_omni.diffusion.attention.backends.flash_attn_hub import _get_hub_module
|
|
|
|
_get_hub_module(hub_repo)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"{backend_upper} was explicitly selected, but {hub_repo} has no "
|
|
"compatible build for this torch/CUDA (or failed to load). "
|
|
"Install kernels>=0.16.1 so stable-ABI variants resolve, or select "
|
|
"a different attention backend."
|
|
) from e
|
|
|
|
if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
|
|
if is_blackwell:
|
|
reason = (
|
|
"Blackwell requires CuTe FlashAttention-4 "
|
|
"(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only"
|
|
)
|
|
elif not compute_supported:
|
|
reason = "compute capability < 8.0"
|
|
else:
|
|
reason = "Flash Attention package unavailable"
|
|
raise ValueError(
|
|
f"FLASH_ATTN was explicitly selected but is unsupported ({reason}). Select a compatible backend."
|
|
)
|
|
if backend_upper == "FLASHINFER_ATTN" and not flashinfer_available:
|
|
raise ValueError(
|
|
"FLASHINFER_ATTN was explicitly selected, but FlashInfer is unavailable. "
|
|
"Install a compatible FlashInfer build or select a different backend."
|
|
)
|
|
if backend_upper == "FLASHINFER_ATTN" and not flashinfer_prefill_available:
|
|
raise ValueError(
|
|
"FLASHINFER_ATTN was explicitly selected, but the installed FlashInfer build does not "
|
|
"provide BatchPrefillWithRaggedKVCacheWrapper. Install a compatible build."
|
|
)
|
|
if backend_upper == "FLASHINFER_ATTN":
|
|
from vllm_omni.diffusion.attention.backends.flashinfer_attn import FlashInferAttentionBackend
|
|
|
|
# head_size <= 0 is the capability-probe sentinel. Skip geometry
|
|
# validation; Attention construction still checks the real size.
|
|
if head_size > 0 and not FlashInferAttentionBackend.supports_head_size(head_size):
|
|
raise ValueError(
|
|
f"FLASHINFER_ATTN was explicitly selected but head_size={head_size} is unsupported. "
|
|
f"Supported head sizes: {FlashInferAttentionBackend.get_supported_head_sizes()}. "
|
|
"Select TORCH_SDPA or another backend."
|
|
)
|
|
if backend_upper == "CUDNN_ATTN":
|
|
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend
|
|
|
|
# head_size <= 0 is the capability-probe sentinel (auto-pad does
|
|
# not know head_dim). Skip geometry validation; the real
|
|
# Attention construction still checks the configured size.
|
|
if head_size > 0 and not CuDNNAttentionBackend.supports_head_size(head_size):
|
|
raise ValueError(
|
|
f"CUDNN_ATTN was explicitly selected but head_size={head_size} is unsupported. "
|
|
"Blackwell cuDNN FMHA requires head_dim divisible by 8 and no larger than 256. "
|
|
"Select FLASHINFER_ATTN or TORCH_SDPA."
|
|
)
|
|
if backend_upper == "TRTLLM_ATTN":
|
|
trtllm_attn_supported = compute_capability is not None and compute_capability.major == 10
|
|
if not trtllm_attn_supported:
|
|
raise ValueError(
|
|
"TRTLLM_ATTN diffusion attention backend requires a datacenter "
|
|
"Blackwell GPU (SM100 / SM103, compute capability 10.x). Select a "
|
|
"different --diffusion-attention-backend."
|
|
)
|
|
if not trtllm_gen_available:
|
|
raise ValueError(
|
|
"TRTLLM_ATTN was explicitly selected, but the installed FlashInfer build does not "
|
|
"provide trtllm_ragged_attention_deepseek. Install a compatible build."
|
|
)
|
|
backend = DiffusionAttentionBackendEnum[backend_upper]
|
|
logger.debug("Using diffusion attention backend '%s'", backend_upper)
|
|
return backend.get_path()
|
|
|
|
trtllm_attn_default_ok = (
|
|
allow_trtllm_default
|
|
and compute_capability is not None
|
|
and compute_capability.major == 10
|
|
and head_size == 128
|
|
and trtllm_gen_available
|
|
)
|
|
if trtllm_attn_default_ok:
|
|
logger.info(
|
|
"Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell %s, head_dim %d)",
|
|
sm_str,
|
|
head_size,
|
|
)
|
|
return DiffusionAttentionBackendEnum.TRTLLM_ATTN.get_path()
|
|
|
|
if is_blackwell and cudnn_blackwell_ready:
|
|
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend
|
|
|
|
# Unknown head_size is a capability probe (SP auto-pad), not a
|
|
# real DiT geometry. Keep CUDNN so mask support is reported.
|
|
if head_size <= 0 or CuDNNAttentionBackend.supports_head_size(head_size):
|
|
logger.info(
|
|
"Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d, head_dim %d)",
|
|
sm_str,
|
|
cudnn_version,
|
|
head_size,
|
|
)
|
|
return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
|
|
logger.info(
|
|
"Skipping CUDNN_ATTN on Blackwell %s: head_dim %d is outside cuDNN FMHA "
|
|
"(multiples of 8, 8-256). Automatic selection will use FlashInfer when available, otherwise SDPA.",
|
|
sm_str,
|
|
head_size,
|
|
)
|
|
|
|
if is_blackwell and flashinfer_prefill_available:
|
|
from vllm_omni.diffusion.attention.backends.flashinfer_attn import FlashInferAttentionBackend
|
|
|
|
if FlashInferAttentionBackend.supports_head_size(head_size):
|
|
logger.info(
|
|
"Defaulting to diffusion attention backend FLASHINFER_ATTN (Blackwell %s, head_dim %d)",
|
|
sm_str,
|
|
head_size,
|
|
)
|
|
return DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
|
|
|
|
if is_blackwell and not cudnn_blackwell_ready:
|
|
logger.warning(
|
|
"Detected Blackwell %s but cuDNN %d < 9.5 — no tuned Blackwell FMHA. "
|
|
"Automatic backend selection will use FlashInfer when available, otherwise SDPA.",
|
|
sm_str,
|
|
cudnn_version,
|
|
)
|
|
|
|
if flash_attn_supported:
|
|
logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
|
|
return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
|
|
|
logger.debug("Defaulting to diffusion attention backend SDPA")
|
|
return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
|
|
|
@classmethod
|
|
def supports_torch_inductor(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
|
|
if local_rank is None:
|
|
return torch.device("cuda")
|
|
return torch.device("cuda", local_rank)
|
|
|
|
@classmethod
|
|
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
|
|
major, minor = torch.cuda.get_device_capability(device_id)
|
|
return DeviceCapability(major=major, minor=minor)
|
|
|
|
@classmethod
|
|
def get_device_count(cls) -> int:
|
|
return torch.accelerator.device_count()
|
|
|
|
@classmethod
|
|
def get_device_version(cls) -> str | None:
|
|
return torch.version.cuda
|
|
|
|
@classmethod
|
|
def synchronize(cls) -> None:
|
|
torch.accelerator.synchronize()
|
|
|
|
@classmethod
|
|
def record_device_event(cls) -> torch.Event | None:
|
|
"""Record a device event on the default stream to mark tensor readiness."""
|
|
try:
|
|
event = torch.Event()
|
|
event.record()
|
|
return event
|
|
except Exception:
|
|
logger.warning("Failed to record device event for cross-stream sync")
|
|
return None
|
|
|
|
@classmethod
|
|
def get_free_memory(cls, device: torch.device | None = None) -> int:
|
|
free, _ = torch.cuda.mem_get_info(device)
|
|
return free
|
|
|
|
@classmethod
|
|
def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
|
|
free, total = torch.cuda.mem_get_info(device)
|
|
return free, total
|
|
|
|
@classmethod
|
|
def get_device_name(cls, device_id: int = 0) -> str:
|
|
return torch.cuda.get_device_name(device_id)
|
|
|
|
@classmethod
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
device_props = torch.cuda.get_device_properties(device_id)
|
|
return device_props.total_memory
|
|
|
|
@classmethod
|
|
def is_fully_connected(cls, device_ids: list[int]) -> bool:
|
|
logger.debug("NVLink detection not available on CudaOmniPlatform; assuming no NVLink.")
|
|
return False
|
|
|
|
@classmethod
|
|
def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
|
|
"""Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops."""
|
|
default = ["vllm_c", "native"]
|
|
|
|
# Use oink if enabled for rms_norm
|
|
# TODO(Laurawly/luka): remove this env var,
|
|
# users can just use IR op priority directly
|
|
rms_norm = default
|
|
if envs.VLLM_USE_OINK_OPS:
|
|
rms_norm = ["oink"] + default
|
|
|
|
return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)
|