From c8bbe5010a7ced112b6b848be64eee54613bacee Mon Sep 17 00:00:00 2001 From: Aleksi Vesanto Date: Wed, 11 Mar 2026 06:17:45 +0200 Subject: [PATCH] [diffusion] feat: add AITER Sage attention backend (#20178) --- .../performance/attention_backends.md | 4 +- .../configs/models/adapter/base.py | 1 + .../configs/models/dits/base.py | 1 + .../layers/attention/backends/aiter_sage.py | 81 +++++++++++++++++++ .../runtime/models/dits/qwen_image.py | 1 + .../runtime/platforms/interface.py | 1 + .../multimodal_gen/runtime/platforms/rocm.py | 10 +++ 7 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py diff --git a/docs/diffusion/performance/attention_backends.md b/docs/diffusion/performance/attention_backends.md index 9113d5bb1..ebb19e111 100644 --- a/docs/diffusion/performance/attention_backends.md +++ b/docs/diffusion/performance/attention_backends.md @@ -30,6 +30,7 @@ For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBa | `video_sparse_attn` | `VIDEO_SPARSE_ATTN` | Requires `vsa`. Configure `sparsity` via `--attention-backend-config`. | | `vmoba_attn` | `VMOBA_ATTN` | Requires `kernel.attn.vmoba_attn.vmoba`. Configure via `--attention-backend-config`. | | `aiter` | `AITER` | Requires `aiter`. | +| `aiter_sage` | `AITER_SAGE` | Requires `aiter`. | | `sparse_video_gen_2_attn` | `SPARSE_VIDEO_GEN_2_ATTN` | Requires `svg`. See installation instructions at https://github.com/svg-project/Sparse-VideoGen. | ## Selection priority @@ -93,7 +94,8 @@ Some backends require additional configuration. You can pass these parameters vi | `sage_attn_3` | ✅ | ❌ | ❌ | ❌ | CUDA-only (optional dependency). | | `video_sparse_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `vsa`. Configure `sparsity` via `--attention-backend-config`. | | `vmoba_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `kernel.attn.vmoba_attn.vmoba`. Configure via `--attention-backend-config`. | -| `aiter` | ✅ | ❌ | ❌ | ❌ | Requires `aiter`. | +| `aiter` | ❌ | ✅ | ❌ | ❌ | Requires `aiter`. | +| `aiter_sage` | ❌ | ✅ | ❌ | ❌ | Requires `aiter`. | | `sparse_video_gen_2_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `svg`. | ## Usage diff --git a/python/sglang/multimodal_gen/configs/models/adapter/base.py b/python/sglang/multimodal_gen/configs/models/adapter/base.py index abd22f115..b0ac5efac 100644 --- a/python/sglang/multimodal_gen/configs/models/adapter/base.py +++ b/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, diff --git a/python/sglang/multimodal_gen/configs/models/dits/base.py b/python/sglang/multimodal_gen/configs/models/dits/base.py index 71ad7c663..44dacbbfe 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/base.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py new file mode 100644 index 000000000..56c274c42 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 37afc2458..de2aa35ba 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 8c03fb324..6feb3b71f 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index d36c76dc0..34ddf020e 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/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,