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

2026-03-11 12:17:45 +08:00
parent 21a0015aa3
commit c8bbe5010a
7 changed files with 98 additions and 1 deletions
@@ -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()
@@ -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,