[diffusion] fix: fix RuntimeError in SageAttention3 on Blackwell with Qwen-Image (#16335)
Co-authored-by: qimcis <qimcis@users.noreply.github.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
co-authored by
qimcis
Mick
qimcis
Mick
parent
87ef05e2e1
commit
dcacc492d0
@@ -5,6 +5,7 @@
|
||||
"""Code inside this file can safely assume cuda platform, e.g. importing
|
||||
pynvml. However, it should not initialize cuda context.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from functools import lru_cache, wraps
|
||||
@@ -130,7 +131,6 @@ class CudaPlatformBase(Platform):
|
||||
sm = capability.to_int() if capability else 0
|
||||
|
||||
if sm in SHARED_SYSMEM_DEVICE_MEM_SMS:
|
||||
|
||||
free_gpu_memory = psutil.virtual_memory().available
|
||||
else:
|
||||
free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
|
||||
@@ -151,6 +151,7 @@ class CudaPlatformBase(Platform):
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> str:
|
||||
target_backend: AttentionBackendEnum | None = None
|
||||
# TODO(will): maybe come up with a more general interface for local attention
|
||||
# if distributed is False, we always try to use Flash attn
|
||||
if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
|
||||
@@ -187,6 +188,7 @@ class CudaPlatformBase(Platform):
|
||||
logger.info(
|
||||
"Sage Attention backend is not installed (To install it, run `pip install sageattention==2.2.0 --no-build-isolation`). Falling back to Flash Attention."
|
||||
)
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
elif selected_backend == AttentionBackendEnum.SAGE_ATTN_3:
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401
|
||||
@@ -198,8 +200,9 @@ class CudaPlatformBase(Platform):
|
||||
except ImportError as e:
|
||||
logger.info(e)
|
||||
logger.info(
|
||||
"Sage Attention 3 backend is not installed (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation). Falling back to Flash Attention."
|
||||
"Sage Attention 3 backend is not installed (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation). Falling back to Torch SDPA."
|
||||
)
|
||||
target_backend = AttentionBackendEnum.TORCH_SDPA
|
||||
elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
|
||||
try:
|
||||
from vsa import block_sparse_attn # noqa: F401
|
||||
@@ -245,43 +248,51 @@ class CudaPlatformBase(Platform):
|
||||
elif selected_backend in [
|
||||
AttentionBackendEnum.FA,
|
||||
]:
|
||||
if cls.is_blackwell():
|
||||
if cls.is_sm120():
|
||||
logger.info(
|
||||
"FlashAttention is not supported on SM12.x in this build; falling back to Torch SDPA."
|
||||
)
|
||||
target_backend = AttentionBackendEnum.TORCH_SDPA
|
||||
elif cls.is_blackwell():
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
|
||||
set_fa_ver,
|
||||
)
|
||||
|
||||
set_fa_ver(4)
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
else:
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
elif selected_backend:
|
||||
raise ValueError(f"Invalid attention backend for {cls.device_name}")
|
||||
else:
|
||||
|
||||
if cls.is_blackwell():
|
||||
if cls.is_sm120():
|
||||
# On SM12.x, the sgl-kernel FlashAttention wheels may not include
|
||||
# support yet. Default to Torch SDPA for correctness.
|
||||
logger.info("Defaulting to Torch SDPA backend on SM12.x")
|
||||
target_backend = AttentionBackendEnum.TORCH_SDPA
|
||||
elif cls.is_blackwell():
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
|
||||
set_fa_ver,
|
||||
)
|
||||
|
||||
set_fa_ver(4)
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
if cls.is_sm120():
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401
|
||||
SageAttention3Backend,
|
||||
)
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
else:
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
|
||||
logger.info("Using Sage Attention 3 backend")
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3.SageAttention3Backend"
|
||||
except ImportError as e:
|
||||
logger.info(e)
|
||||
logger.info(
|
||||
"Sage Attention 3 backend is not installed, Falling back to Torch SDPA (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation)"
|
||||
)
|
||||
target_backend = AttentionBackendEnum.TORCH_SDPA
|
||||
# Ensure we have a target backend selected before validation/fallback.
|
||||
if target_backend is None:
|
||||
target_backend = AttentionBackendEnum.FA
|
||||
|
||||
if target_backend == AttentionBackendEnum.FA and cls.is_blackwell():
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
|
||||
set_fa_ver,
|
||||
)
|
||||
|
||||
set_fa_ver(4)
|
||||
|
||||
if not cls.has_device_capability(80):
|
||||
logger.info(
|
||||
"Cannot use FlashAttention backend for Volta and Turing " "GPUs."
|
||||
)
|
||||
logger.info("Cannot use FlashAttention backend for Volta and Turing GPUs.")
|
||||
target_backend = AttentionBackendEnum.TORCH_SDPA
|
||||
elif dtype not in (torch.float16, torch.bfloat16):
|
||||
logger.info(
|
||||
@@ -332,7 +343,6 @@ class CudaPlatformBase(Platform):
|
||||
# all the related functions work on real physical device ids.
|
||||
# the major benefit of using NVML is that it will not initialize CUDA
|
||||
class NvmlCudaPlatform(CudaPlatformBase):
|
||||
|
||||
@classmethod
|
||||
@lru_cache(maxsize=8)
|
||||
@with_nvml_context
|
||||
@@ -431,7 +441,6 @@ class NvmlCudaPlatform(CudaPlatformBase):
|
||||
|
||||
|
||||
class NonNvmlCudaPlatform(CudaPlatformBase):
|
||||
|
||||
@classmethod
|
||||
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
|
||||
major, minor = torch.cuda.get_device_capability(device_id)
|
||||
|
||||
Reference in New Issue
Block a user