minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[diffusion] fix: fix RuntimeError in SageAttention3 on Blackwell with Qwen-Image (#16335)

Browse Source 
Co-authored-by: qimcis <qimcis@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
2026-01-03 22:12:53 +08:00
committed by GitHub
co-authored by qimcis Mick
parent 87ef05e2e1
commit dcacc492d0
2 changed files with 60 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py
+26 -2
View File
@@ -3,6 +3,7 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
import torch import torch
import torch.nn.functional as F
from sageattn3 import sageattn3_blackwell from sageattn3 import sageattn3_blackwell
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
@@ -17,7 +18,6 @@ logger = init_logger(__name__)
class SageAttention3Backend(AttentionBackend): class SageAttention3Backend(AttentionBackend):
accept_output_buffer: bool = True accept_output_buffer: bool = True
@staticmethod @staticmethod
@@ -38,6 +38,7 @@ class SageAttention3Backend(AttentionBackend):
class SageAttention3Impl(AttentionImpl): class SageAttention3Impl(AttentionImpl):
_warned_gqa_fallback_global: bool = False
def __init__( def __init__(
self, self,
@@ -63,6 +64,29 @@ class SageAttention3Impl(AttentionImpl):
query = query.transpose(1, 2) query = query.transpose(1, 2)
key = key.transpose(1, 2) key = key.transpose(1, 2)
value = value.transpose(1, 2) value = value.transpose(1, 2)
output = sageattn3_blackwell(query, key, value, is_causal=self.causal) # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA
# (Hq != Hkv), fall back to torch SDPA which supports GQA.
if key.shape[1] != query.shape[1]:
if query.shape[1] % key.shape[1] != 0:
raise ValueError(
"GQA/MQA requires query heads to be a multiple of KV heads, "
f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
)
if not type(self)._warned_gqa_fallback_global:
logger.warning(
"SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA."
)
type(self)._warned_gqa_fallback_global = True
output = F.scaled_dot_product_attention(
query,
key,
value,
is_causal=self.causal,
dropout_p=self.dropout,
scale=self.softmax_scale,
enable_gqa=True,
)
else:
output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
output = output.transpose(1, 2) output = output.transpose(1, 2)
return output return output
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+34 -25
View File
@@ -5,6 +5,7 @@
"""Code inside this file can safely assume cuda platform, e.g. importing """Code inside this file can safely assume cuda platform, e.g. importing
pynvml. However, it should not initialize cuda context. pynvml. However, it should not initialize cuda context.
""" """
import os import os
from collections.abc import Callable from collections.abc import Callable
from functools import lru_cache, wraps from functools import lru_cache, wraps
@@ -130,7 +131,6 @@ class CudaPlatformBase(Platform):
sm = capability.to_int() if capability else 0 sm = capability.to_int() if capability else 0
if sm in SHARED_SYSMEM_DEVICE_MEM_SMS: if sm in SHARED_SYSMEM_DEVICE_MEM_SMS:
free_gpu_memory = psutil.virtual_memory().available free_gpu_memory = psutil.virtual_memory().available
else: else:
free_gpu_memory, _ = torch.cuda.mem_get_info(device_id) free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
@@ -151,6 +151,7 @@ class CudaPlatformBase(Platform):
head_size: int, head_size: int,
dtype: torch.dtype, dtype: torch.dtype,
) -> str: ) -> str:
target_backend: AttentionBackendEnum | None = None
# TODO(will): maybe come up with a more general interface for local attention # 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 distributed is False, we always try to use Flash attn
if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
@@ -187,6 +188,7 @@ class CudaPlatformBase(Platform):
logger.info( 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." "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: elif selected_backend == AttentionBackendEnum.SAGE_ATTN_3:
try: try:
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401
@@ -198,8 +200,9 @@ class CudaPlatformBase(Platform):
except ImportError as e: except ImportError as e:
logger.info(e) logger.info(e)
logger.info( 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: elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
try: try:
from vsa import block_sparse_attn # noqa: F401 from vsa import block_sparse_attn # noqa: F401
@@ -245,43 +248,51 @@ class CudaPlatformBase(Platform):
elif selected_backend in [ elif selected_backend in [
AttentionBackendEnum.FA, 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 ( from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
set_fa_ver, set_fa_ver,
) )
set_fa_ver(4) set_fa_ver(4)
target_backend = AttentionBackendEnum.FA target_backend = AttentionBackendEnum.FA
else:
target_backend = AttentionBackendEnum.FA
elif selected_backend: elif selected_backend:
raise ValueError(f"Invalid attention backend for {cls.device_name}") raise ValueError(f"Invalid attention backend for {cls.device_name}")
else: else:
if cls.is_sm120():
if cls.is_blackwell(): # 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 ( from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
set_fa_ver, set_fa_ver,
) )
set_fa_ver(4) set_fa_ver(4)
target_backend = AttentionBackendEnum.FA target_backend = AttentionBackendEnum.FA
if cls.is_sm120(): else:
try: target_backend = AttentionBackendEnum.FA
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401
SageAttention3Backend,
)
logger.info("Using Sage Attention 3 backend") # Ensure we have a target backend selected before validation/fallback.
return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3.SageAttention3Backend" if target_backend is None:
except ImportError as e: target_backend = AttentionBackendEnum.FA
logger.info(e)
logger.info( if target_backend == AttentionBackendEnum.FA and cls.is_blackwell():
"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)" from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
) set_fa_ver,
target_backend = AttentionBackendEnum.TORCH_SDPA )
set_fa_ver(4)
if not cls.has_device_capability(80): if not cls.has_device_capability(80):
logger.info( logger.info("Cannot use FlashAttention backend for Volta and Turing GPUs.")
"Cannot use FlashAttention backend for Volta and Turing " "GPUs."
)
target_backend = AttentionBackendEnum.TORCH_SDPA target_backend = AttentionBackendEnum.TORCH_SDPA
elif dtype not in (torch.float16, torch.bfloat16): elif dtype not in (torch.float16, torch.bfloat16):
logger.info( logger.info(
@@ -332,7 +343,6 @@ class CudaPlatformBase(Platform):
# all the related functions work on real physical device ids. # all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA # the major benefit of using NVML is that it will not initialize CUDA
class NvmlCudaPlatform(CudaPlatformBase): class NvmlCudaPlatform(CudaPlatformBase):
@classmethod @classmethod
@lru_cache(maxsize=8) @lru_cache(maxsize=8)
@with_nvml_context @with_nvml_context
@@ -431,7 +441,6 @@ class NvmlCudaPlatform(CudaPlatformBase):
class NonNvmlCudaPlatform(CudaPlatformBase): class NonNvmlCudaPlatform(CudaPlatformBase):
@classmethod @classmethod
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
major, minor = torch.cuda.get_device_capability(device_id) major, minor = torch.cuda.get_device_capability(device_id)