[diffusion] refactor: refactor cuda attention backend resolver (#29852)

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(
+    def _resolve_default_attn_backend(cls) -> AttentionBackendEnum:
-        cls,
+        if cls.is_sm120():
-        selected_backend: AttentionBackendEnum | None,
+            # On SM12.x, the sgl-kernel FlashAttention wheels may not include
-        head_size: int,
+            # support yet. Default to Torch SDPA for correctness.
-        dtype: torch.dtype,
+            logger.info("Defaulting to Torch SDPA backend on SM12.x")
-    ) -> str:
+            return AttentionBackendEnum.TORCH_SDPA
-        target_backend: AttentionBackendEnum | None = None
+        return AttentionBackendEnum.FA
-        # 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
 
-                from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import (  # noqa: F401
+    @classmethod
-                    SlidingTileAttentionBackend,
+    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"
+        try:
-            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():
             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

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