[diffusion] feat: add AITER Sage attention backend (#20178)

python/sglang/multimodal_gen/configs/models/adapter/base.py

@@ -22,6 +22,7 @@ class AdapterArchConfig(ArchConfig):
             AttentionBackendEnum.SAGE_ATTN,
             AttentionBackendEnum.FA,
             AttentionBackendEnum.AITER,
+            AttentionBackendEnum.AITER_SAGE,
             AttentionBackendEnum.TORCH_SDPA,
             AttentionBackendEnum.VIDEO_SPARSE_ATTN,
             AttentionBackendEnum.VMOBA_ATTN,

python/sglang/multimodal_gen/configs/models/dits/base.py

@@ -29,6 +29,7 @@ class DiTArchConfig(ArchConfig):
             AttentionBackendEnum.SAGE_ATTN,
             AttentionBackendEnum.FA,
             AttentionBackendEnum.AITER,
+            AttentionBackendEnum.AITER_SAGE,
             AttentionBackendEnum.TORCH_SDPA,
             AttentionBackendEnum.VIDEO_SPARSE_ATTN,
             AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,

python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AITERSageBackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER_SAGE
+
+    @staticmethod
+    def get_impl_cls() -> type["AITERSageImpl"]:
+        return AITERSageImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITER Sage backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError(
+            "AITER Sage backend does not have a metadata builder."
+        )
+
+
+class AITERSageImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        dropout_p: float = 0.0,
+        **extra_impl_args,
+    ) -> None:
+
+        try:
+            from aiter.ops.triton.attention.fav3_sage import fav3_sage_wrapper_func
+
+            self.aiter_sage_attn_fn = fav3_sage_wrapper_func
+        except ImportError:
+            raise ImportError(
+                "AITER Sage attention is not available, please update AITER version."
+            )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using aiter sage backend.
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, head_num, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, head_num, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, head_num, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, head_num, head_dim]
+        """
+
+        output = self.aiter_sage_attn_fn(query, key, value)
+        return output

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -576,6 +576,7 @@ class QwenImageCrossAttention(nn.Module):
             supported_attention_backends={
                 AttentionBackendEnum.FA,
                 AttentionBackendEnum.AITER,
+                AttentionBackendEnum.AITER_SAGE,
                 AttentionBackendEnum.TORCH_SDPA,
                 AttentionBackendEnum.SAGE_ATTN,
                 AttentionBackendEnum.SAGE_ATTN_3,

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -34,6 +34,7 @@ class AttentionBackendEnum(enum.Enum):
     SPARSE_VIDEO_GEN_2_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()
     AITER = enum.auto()
+    AITER_SAGE = enum.auto()
     SLA_ATTN = enum.auto()
     SAGE_SLA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -115,6 +115,16 @@ class RocmPlatform(Platform):
             logger.info("Using AITer backend on ROCm.")
             return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend"
 
+        elif selected_backend == AttentionBackendEnum.AITER_SAGE:
+            if dtype in (torch.float16, torch.bfloat16):
+                logger.info("Using AITER Sage backend on ROCm.")
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter_sage.AITERSageBackend"
+            else:
+                logger.warning(
+                    "AITER Sage backend only supports bf16/fp16 inputs but got dtype=%s.",
+                    dtype,
+                )
+
         elif selected_backend in (
             AttentionBackendEnum.SLIDING_TILE_ATTN,
             AttentionBackendEnum.SAGE_ATTN,