[Diffusion][Attention] sage attention backend (#243)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -22,7 +22,7 @@ except ImportError:
         "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
         " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
     )
-    flash_attn_3_func = None
+    raise ImportError
 
 
 class FlashAttentionBackend(AttentionBackend):

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from sageattention import sageattn
+except ImportError:
+    logger.warning(
+        "SageAttentionBackend is not available. You may install sage-attention"
+        " by pip install git+https://github.com/thu-ml/SageAttention.git"
+    )
+    raise ImportError
+
+# TODO add sage3 attention backend
+
+
+class SageAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+
+class SageAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        output = sageattn(
+            query,
+            key,
+            value,
+            tensor_layout="NHD",
+            is_causal=self.causal,
+            sm_scale=self.softmax_scale,
+        )
+        return output

vllm_omni/diffusion/attention/selector.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib
 import os
 from functools import cache
 
@@ -9,18 +10,40 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
-from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 
 logger = init_logger(__name__)
 
-# environment variable value -> backend class
-SUPPORTED_BACKENDS = {
-    "FLASH_ATTN": FlashAttentionBackend,
-    "TORCH_SDPA": SDPABackend,
+# environment variable value -> backend module and class
+_BACKEND_CONFIG = {
+    "FLASH_ATTN": {
+        "module": "vllm_omni.diffusion.attention.backends.flash_attn",
+        "class": "FlashAttentionBackend",
+    },
+    "TORCH_SDPA": {
+        "module": "vllm_omni.diffusion.attention.backends.sdpa",
+        "class": "SDPABackend",
+    },
+    "SAGE_ATTN": {
+        "module": "vllm_omni.diffusion.attention.backends.sage_attn",
+        "class": "SageAttentionBackend",
+    },
 }
 
 
+def load_backend(backend_name: str) -> type[AttentionBackend]:
+    config = _BACKEND_CONFIG[backend_name]
+
+    try:
+        module = importlib.import_module(config["module"])
+        backend_class = getattr(module, config["class"])
+        return backend_class
+    except ImportError as e:
+        raise ImportError(f"Failed to import module {config['module']}: {e}")
+    except AttributeError as e:
+        raise AttributeError(f"Class {config['class']} not found in module: {e}")
+
+
 @cache
 def get_attn_backend(head_size: int) -> type[AttentionBackend]:
     """
@@ -41,13 +64,13 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
 
     if backend_name is not None:
         backend_name_upper = backend_name.upper()
-        if backend_name_upper not in SUPPORTED_BACKENDS:
-            valid_backends = list(SUPPORTED_BACKENDS.keys())
+        if backend_name_upper not in _BACKEND_CONFIG:
+            valid_backends = list(_BACKEND_CONFIG.keys())
             raise ValueError(
                 f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
             )
         logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
-        return SUPPORTED_BACKENDS[backend_name_upper]
+        return load_backend(backend_name_upper)
 
     # Default to SDPA
     return SDPABackend