chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 09:12:53 -05:00
committed by GitHub
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

28
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py
View File

@@ -3,6 +3,7 @@
# SPDX-License-Identifier: Apache-2.0
import torch
import torch.nn.functional as F
from sageattn3 import sageattn3_blackwell
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
@@ -17,7 +18,6 @@ logger = init_logger(__name__)
class SageAttention3Backend(AttentionBackend):
accept_output_buffer: bool = True
@staticmethod
@@ -38,6 +38,7 @@ class SageAttention3Backend(AttentionBackend):
class SageAttention3Impl(AttentionImpl):
_warned_gqa_fallback_global: bool = False
def __init__(
self,
@@ -63,6 +64,29 @@ class SageAttention3Impl(AttentionImpl):
query = query.transpose(1, 2)
key = key.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)
return output

59
python/sglang/multimodal_gen/runtime/platforms/cuda.py
View File

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