diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index d3d605170..4139d20c8 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -27,6 +27,10 @@ from sglang.multimodal_gen.utils import import_pynvml logger = init_logger(__name__) +_SDPA_BACKEND_CLS_STR = ( + "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" +) + _P = ParamSpec("_P") _R = TypeVar("_R") @@ -68,6 +72,218 @@ def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: return wrapper +class _CudaAttentionBackendResolver: + backend: AttentionBackendEnum + + @classmethod + def resolve(cls, platform) -> str | AttentionBackendEnum: + raise NotImplementedError + + +class _DirectCudaAttentionBackendResolver(_CudaAttentionBackendResolver): + backend_cls_str: str + + @classmethod + def resolve(cls, platform) -> str: + return cls.backend_cls_str + + +class _AITerAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.AITER + backend_cls_str = ( + "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend" + ) + + +class _TorchSDPAAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.TORCH_SDPA + backend_cls_str = _SDPA_BACKEND_CLS_STR + + +class _SparseLinearAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.SLA_ATTN + backend_cls_str = "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend" + + +class _SageSparseLinearAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.SAGE_SLA_ATTN + backend_cls_str = "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend" + + +class _SlidingTileAttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.SLIDING_TILE_ATTN + + @classmethod + def resolve(cls, platform) -> str: + try: + from st_attn import sliding_tile_attention # noqa: F401 + + from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( # noqa: F401 + SlidingTileAttentionBackend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" + except ImportError as e: + logger.error("Failed to import Sliding Tile Attention backend: %s", str(e)) + raise ImportError( + "Sliding Tile Attention backend is not installed. " + ) from e + + +class _SageAttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.SAGE_ATTN + + @classmethod + def resolve(cls, platform) -> str | AttentionBackendEnum: + try: + from sageattention import sageattn # noqa: F401 + + from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import ( # noqa: F401 + SageAttentionBackend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend" + except ImportError as e: + logger.info(e) + 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." + ) + return AttentionBackendEnum.FA + + +class _SageAttention3BackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.SAGE_ATTN_3 + + @classmethod + def resolve(cls, platform) -> str | AttentionBackendEnum: + try: + from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 + SageAttention3Backend, + ) + + 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 (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation). Falling back to Torch SDPA." + ) + return AttentionBackendEnum.TORCH_SDPA + + +class _VideoSparseAttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN + + @classmethod + def resolve(cls, platform) -> str: + try: + from vsa import block_sparse_attn # noqa: F401 + + from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( # noqa: F401 + VideoSparseAttentionBackend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" + except ImportError as e: + logger.error("Failed to import Video Sparse Attention backend: %s", str(e)) + raise ImportError("Video Sparse Attention backend is not installed.") from e + + +class _SparseVideoGen2AttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN + + @classmethod + def resolve(cls, platform) -> str: + try: + from svg.kernels.triton.permute import ( # noqa: F401 + apply_inverse_permutation_triton, + permute_tensor_by_labels_triton, + ) + from svg.kmeans_utils import ( # noqa: F401 + batch_kmeans_Euclid, + density_calculation, + dynamic_block_sparse_fwd_flashinfer, + identify_dynamic_map, + ) + + from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import ( # noqa: F401 + SparseVideoGen2AttentionBackend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn.SparseVideoGen2AttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Sparse Video Gen 2 (SAP) Attention backend: %s", + str(e), + ) + raise ImportError( + "Sparse Video Gen 2 (SAP) Attention backend is not installed. " + "Please install it by following the instructions at " + "https://github.com/svg-project/Sparse-VideoGen" + ) from e + + +class _VMOBAAttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.VMOBA_ATTN + + @classmethod + def resolve(cls, platform) -> str: + try: + from kernel.attn.vmoba_attn.vmoba import moba_attn_varlen # noqa: F401 + + from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import ( # noqa: F401 + VMOBAAttentionBackend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.vmoba.VMOBAAttentionBackend" + except ImportError as e: + logger.error("Failed to import Video MoBA Attention backend: %s", str(e)) + raise ImportError("Video MoBA Attention backend is not installed. ") from e + + +class _FlashAttention2BackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.FA2 + + @classmethod + def resolve(cls, platform) -> str: + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import ( # noqa: F401 + FlashAttention2Backend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend" + + +class _FlashAttentionBackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.FA + + @classmethod + def resolve(cls, platform) -> AttentionBackendEnum: + if platform.is_sm120(): + logger.info( + "FlashAttention is not supported on SM12.x in this build; falling back to Torch SDPA." + ) + return AttentionBackendEnum.TORCH_SDPA + return AttentionBackendEnum.FA + + +_CUDA_ATTENTION_BACKEND_RESOLVERS = { + resolver.backend: resolver + for resolver in ( + _AITerAttentionBackendResolver, + _TorchSDPAAttentionBackendResolver, + _SparseLinearAttentionBackendResolver, + _SageSparseLinearAttentionBackendResolver, + _SlidingTileAttentionBackendResolver, + _SageAttentionBackendResolver, + _SageAttention3BackendResolver, + _VideoSparseAttentionBackendResolver, + _SparseVideoGen2AttentionBackendResolver, + _VMOBAAttentionBackendResolver, + _FlashAttention2BackendResolver, + _FlashAttentionBackendResolver, + ) +} + + class CudaPlatformBase(Platform): _enum = PlatformEnum.CUDA device_name: str = "cuda" @@ -216,179 +432,39 @@ class CudaPlatformBase(Platform): return free_gpu_memory / (1 << 30) @classmethod - def get_attn_backend_cls_str( - cls, - selected_backend: AttentionBackendEnum | None, - 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: - try: - from st_attn import sliding_tile_attention # noqa: F401 + def _resolve_default_attn_backend(cls) -> AttentionBackendEnum: + 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") + return AttentionBackendEnum.TORCH_SDPA + return AttentionBackendEnum.FA - from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( # noqa: F401 - SlidingTileAttentionBackend, - ) + @classmethod + def _prepare_flash_attention_for_blackwell(cls) -> bool: + if not cls.is_blackwell(): + return True - return "sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" - except ImportError as e: - logger.error( - "Failed to import Sliding Tile Attention backend: %s", str(e) - ) - raise ImportError( - "Sliding Tile Attention backend is not installed. " - ) from e - elif selected_backend == AttentionBackendEnum.SAGE_ATTN: - try: - from sageattention import sageattn # noqa: F401 - - from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import ( # noqa: F401 - SageAttentionBackend, - ) - - return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend" - except ImportError as e: - logger.info(e) - 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 - SageAttention3Backend, - ) - - 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 (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 - - from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( # noqa: F401 - VideoSparseAttentionBackend, - ) - - return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" - except ImportError as e: - logger.error( - "Failed to import Video Sparse Attention backend: %s", str(e) - ) - raise ImportError( - "Video Sparse Attention backend is not installed." - ) from e - elif selected_backend == AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN: - try: - from svg.kernels.triton.permute import ( # noqa: F401 - apply_inverse_permutation_triton, - permute_tensor_by_labels_triton, - ) - from svg.kmeans_utils import ( # noqa: F401 - batch_kmeans_Euclid, - density_calculation, - dynamic_block_sparse_fwd_flashinfer, - identify_dynamic_map, - ) - - from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import ( # noqa: F401 - SparseVideoGen2AttentionBackend, - ) - - return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn.SparseVideoGen2AttentionBackend" - except ImportError as e: - logger.error( - "Failed to import Sparse Video Gen 2 (SAP) Attention backend: %s", - str(e), - ) - raise ImportError( - "Sparse Video Gen 2 (SAP) Attention backend is not installed. " - "Please install it by following the instructions at " - "https://github.com/svg-project/Sparse-VideoGen" - ) from e - elif selected_backend == AttentionBackendEnum.VMOBA_ATTN: - try: - from kernel.attn.vmoba_attn.vmoba import moba_attn_varlen # noqa: F401 - - from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import ( # noqa: F401 - VMOBAAttentionBackend, - ) - - return "sglang.multimodal_gen.runtime.layers.attention.backends.vmoba.VMOBAAttentionBackend" - except ImportError as e: - logger.error( - "Failed to import Video MoBA Attention backend: %s", str(e) - ) - raise ImportError( - "Video MoBA Attention backend is not installed. " - ) from e - elif selected_backend == AttentionBackendEnum.AITER: - return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend" - elif selected_backend == AttentionBackendEnum.TORCH_SDPA: - return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" - elif selected_backend == AttentionBackendEnum.SLA_ATTN: - return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend" - elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN: - return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend" - elif selected_backend == AttentionBackendEnum.FA2: - from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import ( # noqa: F401 - FlashAttention2Backend, - ) - - return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend" - elif selected_backend in [ - AttentionBackendEnum.FA, - ]: - 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 - else: - target_backend = AttentionBackendEnum.FA - elif selected_backend: - raise ValueError(f"Invalid attention backend for {cls.device_name}") - else: - 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 - else: - target_backend = AttentionBackendEnum.FA - - # 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(): + try: from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) + except ImportError: + logger.info( + "Cannot use FlashAttention backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + return False - set_fa_ver(4) + set_fa_ver(4) + return True + @classmethod + def _resolve_flash_attention_backend_cls_str( + cls, target_backend: AttentionBackendEnum, head_size: int, dtype: torch.dtype + ) -> str: if not cls.has_device_capability(80): logger.info("Cannot use FlashAttention backend for Volta and Turing GPUs.") target_backend = AttentionBackendEnum.TORCH_SDPA @@ -398,8 +474,13 @@ class CudaPlatformBase(Platform): "torch.float16 or torch.bfloat16." ) target_backend = AttentionBackendEnum.TORCH_SDPA - # FlashAttn is valid for the model, checking if the package is - # installed. + + if ( + target_backend == AttentionBackendEnum.FA + and not cls._prepare_flash_attention_for_blackwell() + ): + target_backend = AttentionBackendEnum.TORCH_SDPA + if target_backend == AttentionBackendEnum.FA: try: from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( # noqa: F401 @@ -423,10 +504,33 @@ class CudaPlatformBase(Platform): target_backend = AttentionBackendEnum.TORCH_SDPA if target_backend == AttentionBackendEnum.TORCH_SDPA: - return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" + return _SDPA_BACKEND_CLS_STR return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend" + @classmethod + def get_attn_backend_cls_str( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + if selected_backend is None: + target_backend = cls._resolve_default_attn_backend() + else: + resolver = _CUDA_ATTENTION_BACKEND_RESOLVERS.get(selected_backend) + if resolver is None: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + + resolved_backend = resolver.resolve(cls) + if isinstance(resolved_backend, str): + return resolved_backend + target_backend = resolved_backend + + return cls._resolve_flash_attention_backend_cls_str( + target_backend, head_size, dtype + ) + @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py new file mode 100644 index 000000000..fb13bd493 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py @@ -0,0 +1,100 @@ +import unittest +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatformBase +from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum + +SDPA_BACKEND_CLS_STR = ( + "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" +) + + +class FakeCudaPlatform(CudaPlatformBase): + is_sm120_device = False + is_blackwell_device = False + supports_flash_attention = True + + @classmethod + def is_sm120(cls): + return cls.is_sm120_device + + @classmethod + def is_blackwell(cls): + return cls.is_blackwell_device + + @classmethod + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + return cls.supports_flash_attention + + +class TestCudaAttentionBackendSelection(unittest.TestCase): + def setUp(self): + FakeCudaPlatform.is_sm120_device = False + FakeCudaPlatform.is_blackwell_device = False + FakeCudaPlatform.supports_flash_attention = True + + def resolve( + self, + selected_backend: AttentionBackendEnum | None, + dtype: torch.dtype = torch.float16, + ) -> str: + return FakeCudaPlatform.get_attn_backend_cls_str( + selected_backend=selected_backend, + head_size=128, + dtype=dtype, + ) + + def test_direct_torch_sdpa_selection(self): + self.assertEqual( + self.resolve(AttentionBackendEnum.TORCH_SDPA), SDPA_BACKEND_CLS_STR + ) + + def test_direct_aiter_selection(self): + self.assertEqual( + self.resolve(AttentionBackendEnum.AITER), + "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend", + ) + + def test_default_backend_uses_torch_sdpa_on_sm120(self): + FakeCudaPlatform.is_sm120_device = True + + self.assertEqual(self.resolve(None), SDPA_BACKEND_CLS_STR) + + def test_requested_flash_attention_uses_torch_sdpa_on_sm120(self): + FakeCudaPlatform.is_sm120_device = True + + self.assertEqual(self.resolve(AttentionBackendEnum.FA), SDPA_BACKEND_CLS_STR) + + def test_default_backend_falls_back_for_non_flash_attention_dtype(self): + self.assertEqual(self.resolve(None, torch.float32), SDPA_BACKEND_CLS_STR) + + def test_default_backend_falls_back_without_flash_attention_capability(self): + FakeCudaPlatform.supports_flash_attention = False + + self.assertEqual(self.resolve(None), SDPA_BACKEND_CLS_STR) + + def test_blackwell_falls_back_when_flash_attention_prepare_fails(self): + FakeCudaPlatform.is_blackwell_device = True + + with patch.object( + FakeCudaPlatform, + "_prepare_flash_attention_for_blackwell", + return_value=False, + ) as prepare_flash_attention: + self.assertEqual(self.resolve(None), SDPA_BACKEND_CLS_STR) + + prepare_flash_attention.assert_called_once_with() + + def test_invalid_backend_raises(self): + with self.assertRaisesRegex(ValueError, "Invalid attention backend"): + self.resolve(AttentionBackendEnum.AITER_SAGE) + + +if __name__ == "__main__": + unittest.main()