diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py index e23588f88..ef78e80fe 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 917d013dd..3deff525d 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -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)