chore: import upstream snapshot with attribution

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

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

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

@@ -0,0 +1,207 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+
+import aiter
+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
+from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER_GFX95
+
+logger = logging.getLogger(__name__)
+
+_use_fp8_attn = os.environ.get("SGLANG_DIFFUSION_AITER_FP8_ATTN", "0") == "1"
+_fp8_dtype = torch.float8_e4m3fn
+
+# fmha_fwd_hd128_fp8_gfx950 ASM kernel. Support full MHA with q/k/v head_dim == 128 -- e.g., Wan 2.2 self- and cross-attention.
+_FMHA_FP8_HEAD_DIM = 128
+
+
+if _use_fp8_attn:
+    logger.info("DiT FP8 attention enabled via SGLANG_DIFFUSION_AITER_FP8_ATTN=1")
+
+
+def _can_use_fmha_fp8_prefill(
+    q_head_dim: int,
+    k_head_dim: int,
+    v_head_dim: int,
+    num_heads: int,
+    num_kv_heads: int,
+) -> bool:
+    """True if MHA q/k/v head_dim==128 on a gfx950-class arch."""
+    if not USE_AITER_GFX95:
+        return False
+    if num_kv_heads != num_heads:
+        return False
+    return q_head_dim == k_head_dim == v_head_dim == _FMHA_FP8_HEAD_DIM
+
+
+def _fmha_fp8_prefill_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: float,
+    is_causal: bool,
+    q_scale: torch.Tensor,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> torch.Tensor:
+    """
+    FP8 FMHA prefill via aiter.flash_attn_fp8_pertensor_func.
+
+    Expects q, k, v as (batch, seqlen, nheads, 128) FP8, contiguous.
+    """
+
+    def _ensure_fp8_descale(scale: torch.Tensor) -> torch.Tensor:
+        """Per-tensor descale as shape (1,) float32 for flash_attn_fp8_pertensor_func."""
+        return scale.to(dtype=torch.float32).reshape(1).contiguous()
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    q_descale = _ensure_fp8_descale(q_scale)
+    k_descale = _ensure_fp8_descale(k_scale)
+    v_descale = _ensure_fp8_descale(v_scale)
+
+    return aiter.flash_attn_fp8_pertensor_func(
+        q,
+        k,
+        v,
+        q_descale,
+        k_descale,
+        v_descale,
+        causal=is_causal,
+        softmax_scale=softmax_scale,
+        window_size=(-1, -1, 0),
+    )
+
+
+class AITerBackend(AttentionBackend):
+    """
+    Backend for AITemplate attention implementation.
+    """
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER
+
+    @staticmethod
+    def get_impl_cls() -> type["AITerImpl"]:
+        return AITerImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITer backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError("AITer backend does not have a metadata builder.")
+
+
+class AITerImpl(AttentionImpl):
+    """
+    Implementation of attention using AITemplate.
+    """
+
+    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:
+        if num_kv_heads is not None and num_kv_heads != num_heads:
+            raise NotImplementedError(
+                "AITer backend does not support Grouped Query Attention yet."
+            )
+        self.causal = causal
+        self.dropout_p = dropout_p
+        self.softmax_scale = softmax_scale
+
+    @torch.compiler.disable
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using one of:
+          - _fmha_fp8_prefill_attention (FP8, SGLANG_DIFFUSION_AITER_FP8_ATTN=1 when eligible)
+          - flash_attn_func (BF16, default or FP8 fallback for unsupported shapes)
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, num_heads, head_dim]
+        """
+        if _use_fp8_attn:
+            if query.dtype != _fp8_dtype:
+                q_fp8, q_scale = aiter.per_tensor_quant(query, quant_dtype=_fp8_dtype)
+                k_fp8, k_scale = aiter.per_tensor_quant(key, quant_dtype=_fp8_dtype)
+                v_fp8, v_scale = aiter.per_tensor_quant(value, quant_dtype=_fp8_dtype)
+            else:
+                q_fp8, k_fp8, v_fp8 = query, key, value
+                one = torch.tensor(1.0, dtype=torch.float32, device=query.device)
+                q_scale = k_scale = v_scale = one
+
+            d_q = q_fp8.shape[-1]
+            d_k = k_fp8.shape[-1]
+            d_v = v_fp8.shape[-1]
+            h_q = q_fp8.shape[2]
+            h_kv = k_fp8.shape[2]
+
+            if _can_use_fmha_fp8_prefill(d_q, d_k, d_v, h_q, h_kv):
+                return _fmha_fp8_prefill_attention(
+                    q_fp8,
+                    k_fp8,
+                    v_fp8,
+                    softmax_scale=self.softmax_scale,
+                    is_causal=self.causal,
+                    q_scale=q_scale,
+                    k_scale=k_scale,
+                    v_scale=v_scale,
+                )
+
+            logger.warning_once(
+                "FP8 FMHA prefill unsupported for this shape (need gfx950-class AITER, "
+                "full MHA, q/k/v head_dim=%d; got q=%d, k=%d, v=%d, num_heads=%d, "
+                "num_kv_heads=%d). Falling back to BF16.",
+                _FMHA_FP8_HEAD_DIM,
+                d_q,
+                d_k,
+                d_v,
+                h_q,
+                h_kv,
+            )
+
+        # BF16 path
+        output, _ = aiter.flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=self.dropout_p,
+            causal=self.causal,
+            return_attn_probs=False,
+            return_lse=True,
+        )
+        return output

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/layers/attention/backends/ascend_fa.py

@@ -0,0 +1,104 @@
+from dataclasses import dataclass
+from typing import Any
+
+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
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class AscendFAMetadata:
+    pass
+
+
+class AscendFAMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        return AscendFAMetadata()
+
+
+class AscendFABackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["AscendFAImpl"]:
+        return AscendFAImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return AscendFAMetadataBuilder
+
+
+class AscendFAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        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,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        mask = None
+        num_heads, num_key_value_heads = query.shape[2], key.shape[2]
+        if self.causal:
+            seq_len = query.shape[1]
+            mask = torch.triu(
+                torch.ones(seq_len, seq_len, device=query.device), diagonal=1
+            ).bool()
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        output, lse = torch.ops.npu.npu_fused_infer_attention_score(
+            query,
+            key,
+            value,
+            num_heads=num_heads,
+            num_key_value_heads=num_key_value_heads,
+            scale=self.softmax_scale,
+            input_layout="BNSD",
+            softmax_lse_flag=return_softmax_lse,
+            atten_mask=mask,
+        )
+        output = output.transpose(1, 2)
+        if return_softmax_lse:
+            return output, lse
+        return output

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

@@ -0,0 +1,179 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    pass
+
+import torch
+
+from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+
+    # For some attention backends, we allocate an output tensor before
+    # calling the custom op. When piecewise cudagraph is enabled, this
+    # makes sure the output tensor is allocated inside the cudagraph.
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_enum() -> AttentionBackendEnum:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # @abstractmethod
+    # def get_state_cls() -> Type["AttentionState"]:
+    #     raise NotImplementedError
+
+    # @classmethod
+    # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
+    #     return cls.get_metadata_cls()(*args, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return None
+
+
+@dataclass
+class AttentionMetadata:
+    """Attention metadata for prefill and decode batched together."""
+
+    # Current step of diffusion process
+    current_timestep: int
+
+    def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]:
+        """Similar to dataclasses.asdict, but avoids deepcopying."""
+        if skip_fields is None:
+            skip_fields = set()
+        # Note that if we add dataclasses as fields, they will need
+        # similar handling.
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if field.name not in skip_fields
+        }
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class AttentionMetadataBuilder(ABC, Generic[T]):
+    """Abstract class for attention metadata builders."""
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """Create the builder, remember some configuration and parameters."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare(self) -> None:
+        """Prepare for one batch."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        """Build attention metadata with on-device tensors."""
+        raise NotImplementedError
+
+
+class AttentionLayer(Protocol):
+
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+    _k_scale_float: float
+    _v_scale_float: float
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor: ...
+
+
+class AttentionImpl(ABC, Generic[T]):
+
+    @abstractmethod
+    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:
+        raise NotImplementedError
+
+    def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor:
+        """Preprocess QKV tensor before performing attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom preprocessing
+        like reshaping, tiling, scaling, or other transformations.
+
+        Called AFTER all_to_all for distributed attention
+
+        """
+        return qkv
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        """Postprocess the output tensor after the attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom postprocessing
+        like untiling, scaling, or other transformations.
+
+        Called BEFORE all_to_all for distributed attention
+
+        """
+
+        return output
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
+    return wrap_method_with_debug_kernel_once(
+        attn_impl,
+        "forward",
+        op_name=f"diffusion.attn_impl.{attn_impl.__class__.__name__}.forward",
+    )

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

@@ -0,0 +1,279 @@
+from dataclasses import dataclass
+from typing import Any
+
+import attentions  # noqa: F401
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+BSA_BLOCK_SIZE = 128
+
+
+class BlockSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.BLOCK_SPARSE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["BlockSparseAttentionImpl"]:
+        return BlockSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["BlockSparseAttentionMetadata"]:
+        return BlockSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["BlockSparseAttentionMetadataBuilder"]:
+        return BlockSparseAttentionMetadataBuilder
+
+
+@dataclass
+class BlockSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    block_frame_stride: int
+
+
+class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> BlockSparseAttentionMetadata:
+        """
+        Builds BlockSparseAttention metadata.
+
+        Args:
+            current_timestep: The current diffusion timestep.
+            skip_first_steps: Number of initial timesteps to skip before applying
+                sparsity. Must be non‑negative.
+            sparsity: Fraction of tokens to drop (block‑wise) in the block sparse
+                attention mechanism. Must be in the range [0.0, 1.0).
+            raw_latent_shape: Shape of the latent tensor before patching.
+            patch_size: Patch size as (T, height, width). Only the height
+                and width components are used to divide the latent dimensions.
+            **kwargs: Additional keyword arguments (ignored, but accepted for
+                compatibility with base class or calling conventions).
+
+        Returns:
+            BlockSparseAttentionMetadata
+        Note:
+            The `block_frame_stride` is needed to set the first blocks to be non‑sparse.
+        """
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        if len(raw_latent_shape) >= 5:
+            latent_height, latent_width = raw_latent_shape[3:5]
+        else:
+            latent_height, latent_width = raw_latent_shape[-2:]
+
+        latent_height //= patch_size[1]
+        latent_width //= patch_size[2]
+
+        frame_stride = latent_height * latent_width
+        block_frame_stride = (frame_stride + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE
+
+        return BlockSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            block_frame_stride=block_frame_stride,
+        )
+
+
+class BlockSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.block_size = BSA_BLOCK_SIZE
+        self.stride = 8
+        self.default_tokens = 214748647
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _get_estimate_mask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.sparse_block_estimate(
+            query=query,
+            key=key,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+            input_layout="BNSD",
+            stride=self.stride,
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale / self.stride,
+            threshold=1.0,
+            causal=self.causal,
+            keep_sink=True,
+            keep_recent=True,
+            row_sparse=1.0 - sparsity,
+        )
+
+    def _block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        smask: torch.Tensor,
+        sct: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.ops.attentions.block_sparse_attention(
+            query=query,
+            key=key,
+            value=value,
+            sparse_mask=smask,
+            sparse_count_table=sct,
+            input_layout="BNSD",
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale,
+            causal=self.causal,
+            inner_precise=1,
+            pre_tokens=self.default_tokens,
+            next_tokens=self.default_tokens,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+        )
+
+    def _get_smask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        smask, sct = self._get_estimate_mask(
+            query,
+            key,
+            sparsity,
+        )
+
+        seq_len = smask.shape[2]
+
+        # Set the first blocks to be non-sparse to ensure the quality of the first few steps
+        smask[:, :, :block_frame_stride, :seq_len] = 1
+        smask[:, :, :seq_len, :block_frame_stride] = 1
+        smask = smask.to(torch.int8)
+        sct = smask.sum(dim=-1, dtype=torch.int32)
+        return smask, sct
+
+    def _adaptive_block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> torch.Tensor:
+        # TODO Currently implementation for BSND input layout has quality issues
+        # When the implementation is improved, transposes can be removed
+        q = query.permute(0, 2, 1, 3).contiguous()
+        k = key.permute(0, 2, 1, 3).contiguous()
+        v = value.permute(0, 2, 1, 3).contiguous()
+
+        smask, sct = self._get_smask(
+            q,
+            k,
+            block_frame_stride,
+            sparsity,
+        )
+        output = self._block_sparse_attention(q, k, v, smask, sct)
+        output = output.permute(0, 2, 1, 3).contiguous()
+
+        return output
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._adaptive_block_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.block_frame_stride,
+                attn_metadata.sparsity,
+            )
+
+        return output

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

@@ -0,0 +1,445 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
+from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+)
+
+
+def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+# -----------------------------
+# Fake implementations for schema / tracing
+# custom op schema requires FIXED return structure.
+# We provide TWO ops:
+# 1) out-only op: always returns Tensor
+# 2) out+lse op: always returns Tuple[Tensor, Tensor]
+# -----------------------------
+def flash_attn_varlen_func_fake_out(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (q.shape[0],)
+    )
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    return out
+
+
+def flash_attn_varlen_func_fake_out_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+        total_q = batch_size * seqlen_q
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+        total_q = q.shape[0]
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
+    )
+    lse_shape = (
+        (batch_size, num_head, seqlen_q)
+        if cu_seqlens_q is None
+        else (num_head, total_q)
+    )
+
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    lse = q.new_empty(lse_shape, dtype=torch.float32)
+    return out, lse
+
+
+# -----------------------------
+# Registered custom ops
+# NOTE: fixed return schemas to avoid:
+# "Object of type 'Tensor' is not an instance of 'sequence'"
+# -----------------------------
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out)
+def flash_attn_varlen_func_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    if window_size is None:
+        window_size = [-1, -1]
+    if return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. "
+            "Use flash_attn_varlen_func_op_lse for (out, lse)."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=False,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out_lse)
+def flash_attn_varlen_func_op_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if window_size is None:
+        window_size = [-1, -1]
+    if not return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. "
+            "Use flash_attn_varlen_func_op for out-only."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=True,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+fa_ver = 3
+
+
+def set_fa_ver(ver: int) -> None:
+    global fa_ver
+    fa_ver = ver
+
+
+@dataclass
+class FlashAttentionMetadata:
+    # Sequence lengths for the forward batch
+    # Maximum sequence length for query
+    max_seqlen_q: int = 1
+    # Maximum sequence length for key
+    max_seqlen_k: int = 0
+    # Cumulative sequence lengths for query
+    cu_seqlens_q: torch.Tensor = None
+    # Cumulative sequence lengths for key
+    cu_seqlens_k: torch.Tensor = None
+
+
+class FlashAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore
+        self,
+        raw_latent_shape=list,
+        **kwargs: dict[str, Any],
+    ) -> FlashAttentionMetadata:
+        # TODO: put empty values here to be set at first-run, since the q_len calculation can be complicated
+        return FlashAttentionMetadata(max_seqlen_q=None, max_seqlen_k=None)
+
+
+class FlashAttentionBackend(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_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.attention_metadata = FlashAttentionMetadata()
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        if attn_metadata is not None:
+            if attn_metadata.max_seqlen_q is None:
+                attn_metadata.max_seqlen_q = query.shape[1]
+            if attn_metadata.max_seqlen_k is None:
+                attn_metadata.max_seqlen_k = key.shape[1]
+            max_seqlen_q = attn_metadata.max_seqlen_q
+            max_seqlen_k = attn_metadata.max_seqlen_k
+        else:
+            max_seqlen_q = query.shape[1]
+            max_seqlen_k = key.shape[1]
+
+        # FA version selection:
+        # - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag)
+        # - fa_ver == 4: call custom ops with FIXED return schema
+        if fa_ver == 3:
+            flash_attn_op = flash_attn_varlen_func
+            output = flash_attn_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=return_softmax_lse,
+                ver=fa_ver,
+            )
+            return output
+
+        if fa_ver == 4:
+            if return_softmax_lse:
+                out_tensor, softmax_lse = flash_attn_varlen_func_op_lse(
+                    q=query,
+                    k=key,
+                    v=value,
+                    cu_seqlens_q=None,
+                    cu_seqlens_k=None,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    softmax_scale=self.softmax_scale,
+                    causal=self.causal,
+                    return_softmax_lse=True,
+                    ver=fa_ver,
+                )
+                return out_tensor, softmax_lse
+            out_tensor = flash_attn_varlen_func_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=False,
+                ver=fa_ver,
+            )
+            return out_tensor
+
+        raise ValueError(f"flash attention version {fa_ver} is not supported.")

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

@@ -0,0 +1,79 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# 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.layers.attention.backends.flash_attn import (
+    flash_attn_func,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class FlashAttention2Backend(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_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA2
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttention2Impl"]:
+        return FlashAttention2Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+class FlashAttention2Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        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,
+    ):
+        output = flash_attn_func(
+            q=query,  # type: ignore[no-untyped-call]
+            k=key,
+            v=value,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            max_seqlen_q=None,
+            max_seqlen_k=None,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return output

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

@@ -0,0 +1,191 @@
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# Import to use torch.ops.attentions, install package with sgl_kernel_npu
+try:
+    import attentions  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        (
+            "The required 'attentions' package is not installed."
+            "The package can be installed with sgl_kernel_npu"
+        )
+    ) from e
+
+logger = init_logger(__name__)
+
+
+class LaserAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.LASER_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["LaserAttentionImpl"]:
+        return LaserAttentionImpl
+
+
+class LaserAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.softmax_scale = softmax_scale
+
+        # After preprocess input layout should be BNSD.
+        self.seqlen_base = 256
+        self.seqlen_index = 2
+        self.dim_index = 3
+        self.dim_base = 128
+        self.max_token = 2**31 - 1
+        self.seq_len_pad_base = 256
+
+        # the laser attention operator has issues with small seq_len
+        self.min_seqlen = 2048
+        self.sdpa_impl = SDPABackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _pad(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Pad the input tensor along the sequence length and head dimension.
+        to multiples of base values. self.seqlen_index and self.dim_index should be positive integers.
+        """
+
+        seq_len = input_tensor.size(self.seqlen_index)
+        head_dim = input_tensor.size(self.dim_index)
+
+        pad_seq = 0
+        if seq_len % self.seqlen_base != 0:
+            pad_seq = ((seq_len // self.seqlen_base) + 1) * self.seqlen_base - seq_len
+
+        pad_dim = 0
+        if head_dim % self.dim_base != 0:
+            pad_dim = ((head_dim // self.dim_base) + 1) * self.dim_base - head_dim
+
+        if pad_seq == 0 and pad_dim == 0:
+            return input_tensor
+
+        pad_list = [0] * (2 * input_tensor.ndim)
+
+        pad_list[len(pad_list) - 2 * self.seqlen_index - 1] = pad_seq
+        pad_list[len(pad_list) - 2 * self.dim_index - 1] = pad_dim
+
+        return torch.nn.functional.pad(input_tensor, pad_list)
+
+    def _la_preprocess_input(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Currently BSND input layout is not supported
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+
+        if q.dtype != torch.float16:
+            q = q.to(torch.float16)
+            k = k.to(torch.float16)
+            v = v.to(torch.float16)
+
+        q = self._pad(q)
+        k = self._pad(k)
+        v = self._pad(v)
+
+        return q, k, v
+
+    def _la_postprocess_output(
+        self,
+        attention_out: torch.Tensor,
+        dtype: torch.dtype,
+        qseqlen: int,
+        head_dim: int,
+    ) -> torch.Tensor:
+        if dtype != attention_out.dtype:
+            attention_out = attention_out.to(dtype)
+
+        attention_out = attention_out[:, :, :qseqlen, :head_dim]
+        attention_out = attention_out.transpose(1, 2).contiguous()
+        return attention_out
+
+    def _laser_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        head_num: int,
+        pre_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.la(
+            query=query,
+            key=key,
+            value=value,
+            atten_mask=None,
+            alibi_mask=None,
+            drop_mask=None,
+            scale_value=self.softmax_scale,
+            head_num=head_num,
+            input_layout="BNSD",
+            keep_prob=1.0,
+            pre_tokens=pre_tokens,
+            next_tokens=1,
+            is_highPrecision=True,
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        q_seqlen, head_dim = query.shape[1], query.shape[3]
+        kv_seqlen = key.shape[1]
+
+        if q_seqlen < self.min_seqlen or kv_seqlen != q_seqlen:
+            output = self.sdpa_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            pre_tokens = self.max_token
+            if kv_seqlen % self.seq_len_pad_base != 0:
+                pre_tokens = (
+                    kv_seqlen // self.seq_len_pad_base + 1
+                ) * self.seq_len_pad_base - kv_seqlen
+
+            q, k, v = self._la_preprocess_input(query, key, value)
+            _, la_output = self._laser_attention(q, k, v, q.shape[1], pre_tokens)
+            output = self._la_postprocess_output(
+                la_output, query.dtype, q_seqlen, head_dim
+            )
+
+        return output

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

@@ -0,0 +1,414 @@
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import attentions  # noqa: F401
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class RainFusionAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.RAIN_FUSION_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["RainFusionAttentionImpl"]:
+        return RainFusionAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["RainFusionAttentionMetadata"]:
+        return RainFusionAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["RainFusionAttentionMetadataBuilder"]:
+        return RainFusionAttentionMetadataBuilder
+
+
+@dataclass
+class RainFusionAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    latent_shape: list[int]
+
+
+class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> RainFusionAttentionMetadata:
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        latent_shape = raw_latent_shape[-3:]
+        latent_shape = [latent_shape[i] // patch_size[i] for i in range(3)]
+
+        return RainFusionAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            latent_shape=latent_shape,
+        )
+
+
+class RainFusionAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.block_size = 128
+        self.inner_precise = 0
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _avgpool(
+        self, input_tensor: torch.Tensor, pool_size: int = 128
+    ) -> torch.Tensor:
+        batch, seqlen, heads, dim = input_tensor.shape
+        x = input_tensor.permute(0, 2, 3, 1).reshape(batch * heads, dim, seqlen)
+
+        pooled = torch.nn.functional.avg_pool1d(
+            x, kernel_size=pool_size, stride=pool_size, ceil_mode=True
+        )
+        out = pooled.reshape(batch, heads, dim, -1).permute(0, 3, 1, 2).contiguous()
+
+        return out
+
+    def _get_mask_index(self, mask: torch.Tensor) -> torch.Tensor:
+        batch_size, num_heads, seq_len, _ = mask.shape
+
+        mask_reshaped = mask.reshape(-1, seq_len)
+        row_indices = torch.arange(
+            seq_len, device=mask.device, dtype=torch.float32
+        ).unsqueeze(0)
+
+        sorted_vals = torch.where(mask_reshaped, row_indices, seq_len)
+        sorted_vals, _ = torch.sort(sorted_vals, dim=-1)
+        valid_count = mask_reshaped.sum(dim=-1, keepdim=True)
+        keep_mask = row_indices < valid_count
+        result = torch.where(keep_mask, sorted_vals, -1)
+
+        pos_matrix = result.reshape(batch_size, num_heads, seq_len, seq_len).to(
+            torch.int64
+        )
+        return pos_matrix
+
+    def _get_blockwise_mask(
+        self,
+        qkv_pool: torch.Tensor,
+        sparsity: float,
+        scale: float,
+        pool_size: int,
+        latent_shape: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        first_frame_len = latent_shape[1] * latent_shape[2]
+
+        query_pool, key_pool, value_pool = torch.chunk(qkv_pool, 3, dim=0)
+        attn_scores = (
+            query_pool.permute(0, 2, 1, 3) @ key_pool.permute(0, 2, 3, 1) * scale
+        )
+
+        keep_len = math.ceil(attn_scores.shape[-1] * (1 - sparsity))
+
+        topk_values, _ = torch.topk(attn_scores, k=keep_len, dim=-1)
+        mask = attn_scores >= topk_values[..., -1:]
+
+        firstframe_block_num = (first_frame_len + pool_size - 1) // pool_size
+        if firstframe_block_num > 0:
+            mask[:, :, :firstframe_block_num, :] = True
+            mask[:, :, :, :firstframe_block_num] = True
+
+        select_idx = self._get_mask_index(mask)
+        select_idx = select_idx[0].transpose(0, 1)
+        select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
+        return select_idx, select_num_idx
+
+    def _rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        """
+        b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d
+        or
+        b n (f hn hb wn wb) d -> b n (f hn wn hb wb) d
+        """
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+                tensor_h_r = tensor_h_r.reshape(b, frame_num - 1, -1, n, d)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+                tensor_w_r = tensor_w_r.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn hb) (wn wb) n d -> b f (hn wn hb wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _inv_rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+            tensor_hwt = tensor_hwt.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn wn hb wb) n d -> b f (hn hb) (wn wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=3)
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn wn hb wb) n h -> b (f hn hb wn wb) n h",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _do_tensor_rearrange_pooling(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pool_size: int,
+        latent_shape: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Tensor block rearrangement + pooling operation
+        """
+        tensor = torch.cat((query, key, value), dim=0)
+
+        tensor = self._rearrange_with_remaining(tensor, latent_shape)
+        tensor_pool = self._avgpool(tensor, pool_size)
+
+        query_, key_, value_ = torch.chunk(tensor, 3, dim=0)
+        return query_, key_, value_, tensor_pool
+
+    def _rain_fusion_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        select_idx: torch.Tensor,
+        select_num_idx: torch.Tensor,
+        blockshape: List[int],
+        scale: float = 1.0,
+        head_num: int = 1,
+        input_layout: str = "TND",
+        actual_seq_lengths=Optional[torch.Tensor],
+        actual_seq_lengths_kv=Optional[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.rainfusionattention(
+            query=query,
+            key=key,
+            value=value,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=blockshape,
+            attn_mask=None,
+            actual_seq_qlen=actual_seq_lengths,
+            actual_seq_kvlen=actual_seq_lengths_kv,
+            block_table=None,
+            q_input_layout=input_layout,
+            kv_input_layout=input_layout,
+            head_num=head_num,
+            mask_type=0,
+            scale=scale,
+            inner_precise=self.inner_precise,
+            block_size=0,
+        )
+
+    def _rain_fusion_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        latent_shape: tuple[int, int, int],
+        sparsity: float,
+    ):
+        q, k, v, qkv_pool = self._do_tensor_rearrange_pooling(
+            query, key, value, self.block_size, latent_shape
+        )
+
+        select_idx, select_num_idx = self._get_blockwise_mask(
+            qkv_pool,
+            sparsity,
+            self.softmax_scale,
+            self.block_size,
+            latent_shape,
+        )
+
+        batch_size, seqlen_q, head_num, head_dim = q.shape
+        seqlen_kv = k.shape[1]
+
+        layout = "TND"
+        q = q.reshape(-1, head_num, head_dim)
+        k = k.reshape(-1, head_num, head_dim)
+        v = v.reshape(-1, head_num, head_dim)
+
+        actual_seq_lengths = [seqlen_q] * batch_size
+        actual_seq_lengths_kv = [seqlen_kv] * batch_size
+
+        out, _ = self._rain_fusion_attention(
+            q,
+            k,
+            v,
+            scale=self.softmax_scale,
+            head_num=head_num,
+            input_layout=layout,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=[self.block_size, self.block_size],
+            actual_seq_lengths=actual_seq_lengths,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+        )
+
+        out = out.reshape(batch_size, seqlen_q, head_num, head_dim)
+        out = self._inv_rearrange_with_remaining(out, latent_shape)
+        return out
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._rain_fusion_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.latent_shape,
+                attn_metadata.sparsity,
+            )
+        return output

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

@@ -0,0 +1,74 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+from sageattention import sageattn
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+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_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+
+class SageAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        *,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        output = sageattn(
+            query,
+            key,
+            value,
+            # since input is (batch_size, seq_len, head_num, head_dim)
+            tensor_layout="NHD",
+            is_causal=self.causal,
+            sm_scale=self.softmax_scale,
+            return_lse=return_softmax_lse,
+        )
+        if return_softmax_lse:
+            output, softmax_lse = output
+            return output, softmax_lse
+        return output

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

@@ -0,0 +1,92 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn.functional as F
+from sageattn3 import sageattn3_blackwell
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN_3
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA
+        # (Hq != Hkv), fall back to torch SDPA which supports GQA.
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning(
+                    "SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA."
+                )
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
+        output = output.transpose(1, 2)
+        return output

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

@@ -0,0 +1,95 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import nullcontext
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+
+class SDPABackend(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_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.TORCH_SDPA
+
+    @staticmethod
+    def get_impl_cls() -> type["SDPAImpl"]:
+        return SDPAImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SDPAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+        self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_kwargs = {
+            "attn_mask": None,
+            "dropout_p": self.dropout,
+            "is_causal": self.causal,
+            "scale": self.softmax_scale,
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        sdpa_context = (
+            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            if self.allow_cudnn_sdp and query.device.type == "cuda"
+            else nullcontext()
+        )
+        with sdpa_context:
+            output = torch.nn.functional.scaled_dot_product_attention(
+                query, key, value, **attn_kwargs
+            )
+        output = output.transpose(1, 2)
+        return output

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

@@ -0,0 +1,316 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.distributed import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    ForwardContext,
+    get_forward_context,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import dict_to_3d_list
+
+try:
+    from st_attn import sliding_tile_attention
+
+    st_attn_backend_available = True
+except Exception:
+    st_attn_backend_available = False
+
+logger = init_logger(__name__)
+
+
+class RangeDict(dict):
+
+    def __getitem__(self, item: int) -> str:
+        for key in self.keys():
+            if isinstance(key, tuple):
+                low, high = key
+                if low <= item <= high:
+                    return str(super().__getitem__(key))
+            elif key == item:
+                return str(super().__getitem__(key))
+        raise KeyError(f"seq_len {item} not supported for STA")
+
+
+class SlidingTileAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # TODO(will-refactor): check this
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLIDING_TILE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
+        return SlidingTileAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
+        return SlidingTileAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
+        return SlidingTileAttentionMetadataBuilder
+
+
+@dataclass
+class SlidingTileAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    STA_param: list[
+        list[Any]
+    ]  # each timestep with one metadata, shape [num_layers, num_heads]
+
+
+class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        STA_param: list[list[Any]],
+        current_timestep: int,
+        **kwargs: dict[str, Any],
+    ) -> SlidingTileAttentionMetadata:
+        param = STA_param
+        if param is None:
+            return SlidingTileAttentionMetadata(
+                current_timestep=current_timestep, STA_param=[]
+            )
+        return SlidingTileAttentionMetadata(
+            current_timestep=current_timestep, STA_param=param[current_timestep]
+        )
+
+
+class SlidingTileAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if not st_attn_backend_available:
+            raise ValueError("st attn not supported")
+        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
+        # have a more general config for STA?
+        mask_strategy_file_path = (
+            get_global_server_args().attention_backend_config.mask_strategy_file_path
+        )
+        if mask_strategy_file_path is None:
+            raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set")
+
+        # TODO(kevin): get mask strategy for different STA modes
+        with open(mask_strategy_file_path) as f:
+            mask_strategy = json.load(f)
+        self.mask_strategy = dict_to_3d_list(mask_strategy)
+
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+        # STA config
+        self.STA_base_tile_size = [6, 8, 8]
+        self.dit_seq_shape_mapping = RangeDict(
+            {
+                (115200, 115456): "30x48x80",
+                82944: "36x48x48",
+                69120: "18x48x80",
+            }
+        )
+        self.full_window_mapping = {
+            "30x48x80": [5, 6, 10],
+            "36x48x48": [6, 6, 6],
+            "18x48x80": [3, 6, 10],
+        }
+
+    def tile(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+
+    def untile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(
+            x,
+            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+        return x
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        img_sequence_length = qkv.shape[1]
+        self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length]
+        self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str]
+        self.dit_seq_shape_int = list(map(int, self.dit_seq_shape_str.split("x")))
+        self.img_seq_length = (
+            self.dit_seq_shape_int[0]
+            * self.dit_seq_shape_int[1]
+            * self.dit_seq_shape_int[2]
+        )
+        return self.tile(qkv)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        if self.mask_strategy is None:
+            raise ValueError("mask_strategy cannot be None for SlidingTileAttention")
+        if self.mask_strategy[0] is None:
+            raise ValueError("mask_strategy[0] cannot be None for SlidingTileAttention")
+
+        timestep = attn_metadata.current_timestep
+        forward_context: ForwardContext = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        if forward_batch is None:
+            raise ValueError("forward_batch cannot be None")
+        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
+        layer_idx = int(self.prefix.split(".")[-3])
+        if attn_metadata.STA_param is None or len(attn_metadata.STA_param) <= layer_idx:
+            raise ValueError("Invalid STA_param")
+        STA_param = attn_metadata.STA_param[layer_idx]
+
+        text_length = q.shape[1] - self.img_seq_length
+        has_text = text_length > 0
+
+        query = q.transpose(1, 2).contiguous()
+        key = k.transpose(1, 2).contiguous()
+        value = v.transpose(1, 2).contiguous()
+
+        head_num = query.size(1)
+        sp_group = get_sp_group()
+        current_rank = sp_group.rank_in_group
+        start_head = current_rank * head_num
+
+        # searching or tuning mode
+        if len(STA_param) < head_num * sp_group.world_size:
+            sparse_attn_hidden_states_all = []
+            full_mask_window = STA_param[-1]
+            for window_size in STA_param[:-1]:
+                sparse_hidden_states = sliding_tile_attention(
+                    query,
+                    key,
+                    value,
+                    [window_size] * head_num,
+                    text_length,
+                    has_text,
+                    self.dit_seq_shape_str,
+                ).transpose(1, 2)
+                sparse_attn_hidden_states_all.append(sparse_hidden_states)
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                [full_mask_window] * head_num,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+            attn_L2_loss = []
+            attn_L1_loss = []
+            # average loss across all heads
+            for sparse_attn_hidden_states in sparse_attn_hidden_states_all:
+                # L2 loss
+                attn_L2_loss_ = (
+                    torch.mean(
+                        (sparse_attn_hidden_states.float() - hidden_states.float())
+                        ** 2,
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_]
+                attn_L2_loss.append(attn_L2_loss_)
+                # L1 loss
+                attn_L1_loss_ = (
+                    torch.mean(
+                        torch.abs(
+                            sparse_attn_hidden_states.float() - hidden_states.float()
+                        ),
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_]
+                attn_L1_loss.append(attn_L1_loss_)
+
+            layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss}
+
+            if forward_batch.is_cfg_negative:
+                if forward_batch.mask_search_final_result_neg is not None:
+                    forward_batch.mask_search_final_result_neg[timestep].append(
+                        layer_loss_save
+                    )
+            else:
+                if forward_batch.mask_search_final_result_pos is not None:
+                    forward_batch.mask_search_final_result_pos[timestep].append(
+                        layer_loss_save
+                    )
+        else:
+            windows = [STA_param[head_idx + start_head] for head_idx in range(head_num)]
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                windows,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+        return hidden_states

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

@@ -0,0 +1,695 @@
+"""
+Copyright (c) 2025 by SLA team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+
+This implementation is adapted from: from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py and https://github.com/thu-ml/SLA/blob/main/SageSLA/core.py
+Citation (please cite if you use this code):
+
+@article{zhang2025sla,
+  title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention},
+  author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and Joseph E. Gonzalez and Jun Zhu and Jianfei Chen},
+  journal={arXiv preprint arXiv:2509.24006},
+  year={2025}
+}
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# ==================================SLA Functions===================================
+def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
+    arg_k = k - torch.mean(
+        k, dim=-2, keepdim=True
+    )  # smooth-k technique in SageAttention
+    pooled_qblocks = mean_pool(q, BLKQ)
+    pooled_kblocks = mean_pool(arg_k, BLKK)
+    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
+
+    K = pooled_score.shape[-1]
+    topk = min(K, int(topk_ratio * K))
+    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
+
+    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
+    sparse_map.scatter_(-1, lut, 1)
+    return sparse_map, lut, topk
+
+
+def mean_pool(x, BLK):
+    assert x.is_contiguous()
+
+    B, H, L, D = x.shape
+    L_BLOCKS = (L + BLK - 1) // BLK
+    x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
+
+    grid = (L_BLOCKS, B * H)
+    compress_kernel[grid](x, x_mean, L, D, BLK)
+    return x_mean
+
+
+@triton.jit
+def compress_kernel(
+    X,
+    XM,
+    L: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+):
+    idx_l = tl.program_id(0)
+    idx_bh = tl.program_id(1)
+
+    offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    offs_d = tl.arange(0, D)
+
+    x_offset = idx_bh * L * D
+    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
+    x = tl.load(
+        X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
+    )
+
+    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
+    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
+    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
+
+
+@triton.jit
+def _attn_fwd(
+    Q,
+    K,
+    V,
+    qk_scale: tl.constexpr,
+    topk: tl.constexpr,
+    LUT,
+    LSE,
+    OS,
+    L: tl.constexpr,
+    M_BLOCKS: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_m = tl.program_id(0).to(tl.int64)
+    idx_bh = tl.program_id(1).to(tl.int64)
+
+    qkv_offset = idx_bh * L * D
+    lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
+    lse_offset = idx_bh * L
+    offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, D)
+
+    Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
+    V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
+    OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    LUT_ptr = LUT + lut_offset
+    LSE_ptrs = LSE + lse_offset + offs_m
+
+    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
+
+    q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
+    for block_idx in tl.range(topk):
+        idx_n = tl.load(LUT_ptr + block_idx)
+        n_mask = offs_n < L - idx_n * BLOCK_N
+
+        k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
+        qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634)  # = 1 / ln(2)
+        if L - idx_n * BLOCK_N < BLOCK_N:
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+        v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
+        local_m = tl.max(qk, 1)
+        new_m = tl.maximum(m_i, local_m)
+        qk = qk - new_m[:, None]
+
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        alpha = tl.math.exp2(m_i - new_m)
+        o_s = o_s * alpha[:, None]
+        o_s += tl.dot(p.to(v.dtype), v)
+
+        l_i = l_i * alpha + l_ij
+        m_i = new_m
+
+    o_s = o_s / l_i[:, None]
+    tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
+
+    m_i += tl.math.log2(l_i)
+    tl.store(LSE_ptrs, m_i, mask=offs_m < L)
+
+
+def _get_cuda_arch(device_index: int) -> str:
+    """Get CUDA architecture string for the given device."""
+    major, minor = torch.cuda.get_device_capability(device_index)
+    return f"sm{major}{minor}"
+
+
+# ==================================SLA Class===================================
+class SparseLinearAttentionBackend(AttentionBackend):
+    """Sparse Linear Attention Backend for efficient attention computation."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseLinearAttentionImpl"]:
+        return SparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseLinearAttentionMetadata"]:
+        return SparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseLinearAttentionMetadataBuilder"]:
+        return SparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SparseLinearAttentionMetadata:
+        return SparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    """Implementation of sparse linear attention for the backend."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        # SLA-specific parameters - matched to TurboDiffusion defaults
+        topk_ratio: float = 0.1,  # TurboDiffusion uses topk=0.1
+        feature_map: str = "softmax",
+        BLKQ: int = 128,  # TurboDiffusion uses BLKQ=128
+        BLKK: int = 64,  # TurboDiffusion uses BLKK=64
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        # SLA-specific config
+        self.topk_ratio = topk_ratio
+        self.BLKQ = BLKQ
+        self.BLKK = BLKK
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(self, q, k, v):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseLinearAttentionMetadata = None,
+    ) -> torch.Tensor:
+        """Forward pass for sparse linear attention.
+
+        Args:
+            query: query tensor of shape (B, H, L, D)
+            key: key tensor of shape (B, H, L, D)
+            value: value tensor of shape (B, H, L, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, H, L, D)
+        """
+        dtype = query.dtype
+
+        # Transpose for computation
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        # Get sparse attention map
+        sparse_map, lut, real_topk = get_block_map(
+            query, key, topk_ratio=self.topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK
+        )
+
+        # Convert to computation dtype
+        query = query.to(self.dtype)
+        key = key.to(self.dtype)
+        value = value.to(self.dtype)
+
+        # Sparse attention computation
+        o_s = _attention.apply(
+            query, key, value, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
+        )
+
+        # Apply feature maps
+        query = self.feature_map_q(query).to(self.dtype)  # c_q
+        key = self.feature_map_k(key).to(self.dtype)  # c_k
+        # Linear attention computation
+        o_l = self._calc_linear_attention_with_torch(query, key, value)
+
+        # Apply projection and combine results
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear attention
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output
+
+
+class _attention(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert k_block_id.is_contiguous() and lut.is_contiguous()
+
+        # We recommend the following two settings
+        assert BLOCK_M == 64 or BLOCK_M == 128
+        assert BLOCK_N == 64
+
+        B, H, L, D = q.shape
+        if qk_scale is None:
+            qk_scale = D**-0.5
+
+        M_BLOCKS = triton.cdiv(L, BLOCK_M)
+
+        o_s = torch.empty_like(v)
+        lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
+
+        grid = (M_BLOCKS, B * H)
+        _attn_fwd[grid](
+            q,
+            k,
+            v,
+            qk_scale,
+            topk,
+            lut,
+            lse,
+            o_s,
+            L,
+            M_BLOCKS,
+            D,
+            BLOCK_M,
+            BLOCK_N,
+            num_warps=4 if q.shape[-1] == 64 else 8,
+            num_stages=3,
+        )
+
+        ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
+        ctx.qk_scale = qk_scale
+        ctx.topk = topk
+        ctx.BLOCK_M = BLOCK_M
+        ctx.BLOCK_N = BLOCK_N
+        return o_s
+
+
+# ==================================SageSLA Class===================================
+SAGESLA_ENABLED = True
+try:
+    import spas_sage_attn._fused as fused
+    import spas_sage_attn._qattn as qattn
+    from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant
+except ImportError:
+    SAGESLA_ENABLED = False
+
+SAGE2PP_ENABLED = True
+try:
+    from spas_sage_attn._qattn import (
+        qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold,
+    )
+except ImportError:
+    SAGE2PP_ENABLED = False
+
+
+class SageSparseLinearAttentionBackend(AttentionBackend):
+    """Quantized Sparse-Linear Attention backend using SageAttention kernels."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageSparseLinearAttentionImpl"]:
+        return SageSparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SageSparseLinearAttentionMetadata"]:
+        return SageSparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SageSparseLinearAttentionMetadataBuilder"]:
+        return SageSparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SageSparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sage Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SageSparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SageSparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SageSparseLinearAttentionMetadata:
+        return SageSparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        topk_ratio: float = 0.5,
+        feature_map: str = "softmax",
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        assert (
+            SAGESLA_ENABLED
+        ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA."
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5
+        self.causal = causal
+        self.prefix = prefix
+
+        self.topk_ratio = topk_ratio
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Forward pass for Sage Sparse Linear attention with quantized kernels.
+        Args:
+            query: query tensor of shape (B, L, H, D)
+            key: key tensor of shape (B, L, H, D)
+            value: value tensor of shape (B, L, H, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, L, H, D)
+        """
+        dtype = query.dtype
+
+        # Transpose from (B, L, H, D) to SLA format (B, H, L, D)
+        q = query.transpose(1, 2).contiguous()
+        k = key.transpose(1, 2).contiguous()
+        v = value.transpose(1, 2).contiguous()
+
+        # Determine block sizes based on GPU architecture
+        arch = _get_cuda_arch(q.device.index)
+
+        if arch == "sm90":
+            BLKQ = 64
+            BLKK = 128
+        else:
+            BLKQ = 128
+            BLKK = 64
+        # Compute block-sparse attention pattern
+        sparse_map, lut, real_topk = get_block_map(
+            q, k, topk_ratio=self.topk_ratio, BLKQ=BLKQ, BLKK=BLKK
+        )
+
+        # Convert to compute dtype
+        q = q.to(self.dtype)
+        k = k.to(self.dtype)
+        v = v.to(self.dtype)
+
+        ########## SPARGE BEGIN ##########
+        km = k.mean(dim=-2, keepdim=True)
+        headdim = q.size(-1)
+        assert headdim in [
+            64,
+            128,
+        ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale."
+
+        # Quantize Q, K to INT8
+        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK)
+        lut, valid_block_num = block_map_lut_triton(sparse_map)
+        scale = 1.0 / (headdim**0.5)
+
+        o_s = torch.empty_like(q)
+
+        if arch in ("sm80", "sm86", "sm87"):
+            pvthreshold = torch.full(
+                (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+            )
+            v_fp16 = v.to(torch.float16)
+            qattn.qk_int8_sv_f16_accum_f16_block_sparse_attn_inst_buf_with_pv_threshold(
+                q_int8,
+                k_int8,
+                v_fp16,
+                o_s,
+                lut,
+                valid_block_num,
+                pvthreshold,
+                q_scale,
+                k_scale,
+                1,
+                False,
+                1,
+                scale,
+                0,
+            )
+        else:
+            b, h_kv, kv_len, head_dim = v.shape
+            padded_len = (kv_len + 127) // 128 * 128
+            v_transposed_permutted = torch.empty(
+                (b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device
+            )
+            fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1)
+            v_fp8 = torch.empty(
+                v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device
+            )
+            v_scale = torch.empty(
+                (b, h_kv, head_dim), dtype=torch.float32, device=v.device
+            )
+            fused.scale_fuse_quant_cuda(
+                v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1
+            )
+
+            if arch == "sm90":
+                qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90(
+                    q_int8,
+                    k_int8,
+                    v_fp8,
+                    o_s,
+                    lut,
+                    valid_block_num,
+                    q_scale,
+                    k_scale,
+                    v_scale,
+                    1,
+                    False,
+                    1,
+                    scale,
+                )
+            else:
+                pvthreshold = torch.full(
+                    (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+                )
+                if SAGE2PP_ENABLED:
+                    qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+                else:
+                    qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+
+        ########## SPARGE END ##########
+
+        # Linear attention with feature maps
+        q_linear = self.feature_map_q(q).to(self.dtype)
+        k_linear = self.feature_map_k(k).to(self.dtype)
+        o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v)
+
+        # Project linear attention output and combine
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear outputs
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output

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

@@ -0,0 +1,562 @@
+"""
+Sparse Video Gen 2 (SAP) attention backend.
+
+This is a baseline integration that wires the backend into the
+attention framework.
+
+Adapted from https://github.com/svg-project/Sparse-VideoGen/blob/main/svg/models/wan/attention.py
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+try:
+    from svg.kernels.triton.permute import (
+        apply_inverse_permutation_triton,
+        permute_tensor_by_labels_triton,
+    )
+    from svg.kmeans_utils import (
+        batch_kmeans_Euclid,
+        dynamic_block_sparse_fwd_flashinfer,
+        identify_dynamic_map,
+    )
+
+    svg2_available = True
+except ImportError:
+    svg2_available = False
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SparseVideoGen2AttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseVideoGen2AttentionImpl"]:
+        return SparseVideoGen2AttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseVideoGen2AttentionMetadata"]:
+        return SparseVideoGen2AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseVideoGen2AttentionMetadataBuilder"]:
+        return SparseVideoGen2AttentionMetadataBuilder
+
+
+@dataclass
+class Svg2LayerCache:
+    # centroids for kmeans clustering
+    q_centroids: torch.Tensor | None = None
+    k_centroids: torch.Tensor | None = None
+    centroids_initialized: bool = False
+
+
+@dataclass
+class Svg2Cache:
+    layers: dict[int, Svg2LayerCache] = field(default_factory=dict)
+
+    def get_layer(self, layer_idx: int) -> Svg2LayerCache:
+        layer_cache = self.layers.get(layer_idx)
+        if layer_cache is None:
+            layer_cache = Svg2LayerCache()
+            self.layers[layer_idx] = layer_cache
+        return layer_cache
+
+
+@dataclass
+class SparseVideoGen2AttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    num_q_centroids: int
+    num_k_centroids: int
+    top_p_kmeans: float
+    min_kc_ratio: float
+    kmeans_iter_init: int
+    kmeans_iter_step: int
+    zero_step_kmeans_init: bool
+    first_layers_fp: float
+    first_times_fp: float
+    context_length: int
+    num_frame: int
+    frame_size: int
+    cache: Svg2Cache
+    prompt_length: int | None = None
+    max_seqlen_q: int | None = None
+    max_seqlen_k: int | None = None
+
+
+def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any:
+    if name not in kwargs:
+        raise ValueError(
+            f"Missing required argument for SparseVideoGen2Attention: {name}"
+        )
+    return kwargs[name]
+
+
+class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore[override]
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, ...],
+        patch_size: tuple[int, int, int],
+        cache: Svg2Cache,
+        num_q_centroids: int,
+        num_k_centroids: int,
+        top_p_kmeans: float,
+        min_kc_ratio: float,
+        kmeans_iter_init: int,
+        kmeans_iter_step: int,
+        zero_step_kmeans_init: bool,
+        first_layers_fp: float,
+        first_times_fp: float,
+        context_length: int = 0,
+        prompt_length: int | None = None,
+        **kwargs: dict[str, Any],
+    ) -> SparseVideoGen2AttentionMetadata:
+        raw_shape = tuple(raw_latent_shape)
+        if len(raw_shape) == 5:
+            t, h, w = raw_shape[2:5]
+        elif len(raw_shape) == 3:
+            t, h, w = raw_shape
+        else:
+            raise ValueError(
+                "raw_latent_shape must be (T, H, W) or (B, C, T, H, W) for SAP attention"
+            )
+        pt, ph, pw = patch_size
+        if t % pt != 0 or h % ph != 0 or w % pw != 0:
+            raise ValueError(
+                "raw_latent_shape must be divisible by patch_size for SAP attention"
+            )
+
+        num_frame = t // pt
+        frame_size = (h // ph) * (w // pw)
+
+        return SparseVideoGen2AttentionMetadata(
+            current_timestep=current_timestep,
+            num_q_centroids=num_q_centroids,
+            num_k_centroids=num_k_centroids,
+            top_p_kmeans=top_p_kmeans,
+            min_kc_ratio=min_kc_ratio,
+            kmeans_iter_init=kmeans_iter_init,
+            kmeans_iter_step=kmeans_iter_step,
+            zero_step_kmeans_init=zero_step_kmeans_init,
+            first_layers_fp=first_layers_fp,
+            first_times_fp=first_times_fp,
+            context_length=context_length,
+            prompt_length=prompt_length,
+            num_frame=num_frame,
+            frame_size=frame_size,
+            cache=cache,
+        )
+
+
+class SparseVideoGen2AttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if causal:
+            raise ValueError(
+                "Sparse Video Gen 2 attention does not support causal attention"
+            )
+        if not svg2_available:
+            raise ImportError(
+                "Sparse Video Gen 2 attention backend requires svg package to be installed"
+                "Please install it by following the instructions at "
+                "https://github.com/svg-project/Sparse-VideoGen"
+            )
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+    def _get_layer_idx(self, prefix: str) -> int:
+        parts = prefix.split(".")
+        if len(parts) < 3:
+            raise ValueError(
+                f"Invalid prefix for SparseVideoGen2AttentionImpl: {prefix}"
+            )
+        return int(parts[-3])
+
+    def kmeans_init(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_step(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.q_centroids,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.k_centroids,
+        )
+
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_clustering(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        if not layer_cache.centroids_initialized:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_init(query, key, attn_metadata)
+            layer_cache.centroids_initialized = True
+            logger.debug(
+                "Centroids initialized at layer %s (init iters: %s).",
+                self.layer_idx,
+                attn_metadata.kmeans_iter_init,
+            )
+        else:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_step(query, key, attn_metadata)
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def semantic_aware_permutation(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+
+        # 1. Kmeans clustering
+        (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        ) = self.kmeans_clustering(query, key, attn_metadata)
+
+        # 2. Identify dynamic map
+        q_cluster_sizes = qcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_q_centroids
+        )
+        k_cluster_sizes = kcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_k_centroids
+        )
+
+        dynamic_map = identify_dynamic_map(
+            qcentroids.view(cfg, num_heads, attn_metadata.num_q_centroids, dim),
+            kcentroids.view(cfg, num_heads, attn_metadata.num_k_centroids, dim),
+            q_cluster_sizes,
+            k_cluster_sizes,
+            attn_metadata.top_p_kmeans,
+            attn_metadata.min_kc_ratio,
+        )
+
+        # 3. Permute the query, key, value
+        q_permuted, q_sorted_indices = permute_tensor_by_labels_triton(
+            query, qlabels, dim=2
+        )
+        k_permuted, k_sorted_indices = permute_tensor_by_labels_triton(
+            key, klabels, dim=2
+        )
+        v_permuted, v_sorted_indices = permute_tensor_by_labels_triton(
+            value, klabels, dim=2, sorted_indices=k_sorted_indices
+        )
+
+        return (
+            q_permuted,
+            k_permuted,
+            v_permuted,
+            dynamic_map,
+            q_cluster_sizes,
+            k_cluster_sizes,
+            q_sorted_indices,
+        )
+
+    def _hunyuan_dynamic_map_post_processing(
+        self,
+        q_perm: torch.Tensor,
+        k_perm: torch.Tensor,
+        v_perm: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        dyn_map: torch.Tensor,
+        qc_sz_s: torch.Tensor,
+        kc_sz_s: torch.Tensor,
+        q_sorted_indices: torch.Tensor,
+        video_length: int,
+        context_length: int,
+        prompt_length: int,
+        unprompt_length: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        # Place the permuted video tokens back and keep text tokens at the tail.
+        query[:, :, :-context_length, :] = q_perm
+        key[:, :, :-context_length, :] = k_perm
+        value[:, :, :-context_length, :] = v_perm
+
+        # Add prompt/unprompt clusters to the dynamic map.
+        dyn_map = F.pad(dyn_map, (0, 2, 0, 2), value=0)
+        dyn_map[:, :, -2, :-1] = True
+        dyn_map[:, :, :-1, -2] = True
+        dyn_map[:, :, -1, -1] = True
+
+        qc_sz_s = F.pad(qc_sz_s, (0, 2), value=0)
+        qc_sz_s[:, :, -2] = prompt_length
+        qc_sz_s[:, :, -1] = unprompt_length
+        kc_sz_s = F.pad(kc_sz_s, (0, 2), value=0)
+        kc_sz_s[:, :, -2] = prompt_length
+        kc_sz_s[:, :, -1] = unprompt_length
+
+        q_sorted_indices = F.pad(q_sorted_indices, (0, context_length), value=0)
+        q_sorted_indices[:, video_length:] = torch.arange(
+            video_length,
+            video_length + context_length,
+            device=q_sorted_indices.device,
+        )
+        return query, key, value, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ) -> torch.Tensor:
+        torch.backends.cuda.preferred_linalg_library(backend="magma")
+        res = None
+        # bshd -> bhsd
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        batch_size, num_heads, seq_len, dim = query.size()
+
+        context_length, num_frame, frame_size = (
+            attn_metadata.context_length,
+            attn_metadata.num_frame,
+            attn_metadata.frame_size,
+        )
+        prompt_length = attn_metadata.prompt_length
+        if prompt_length is None:
+            prompt_length = context_length
+
+        assert (
+            seq_len == context_length + num_frame * frame_size
+        ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}"
+
+        # Determine if we use Full Attention to calculate
+        full_attention_flag = False
+
+        if self.layer_idx < attn_metadata.first_layers_fp:
+            full_attention_flag = True
+        if attn_metadata.current_timestep > attn_metadata.first_times_fp:
+            full_attention_flag = True
+
+        if full_attention_flag:
+            if attn_metadata.zero_step_kmeans_init:
+                video_length = attn_metadata.num_frame * attn_metadata.frame_size
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                self.kmeans_clustering(query_video, key_video, attn_metadata)
+
+            with sdpa_kernel(
+                SDPBackend.CUDNN_ATTENTION
+            ):  # not sure why we need to force cudnn here, but it's faster than flash attention
+                output_hidden_states = torch.nn.functional.scaled_dot_product_attention(
+                    query, key, value, dropout_p=0.0, is_causal=False
+                )
+
+            res = output_hidden_states.reshape(
+                batch_size, num_heads, seq_len, dim
+            ).transpose(1, 2)
+        else:
+            if context_length > 0:
+                video_length = num_frame * frame_size
+                unprompt_length = max(context_length - prompt_length, 0)
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                value_video = value[:, :, :video_length, :].contiguous()
+
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(
+                    query_video, key_video, value_video, attn_metadata
+                )
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self._hunyuan_dynamic_map_post_processing(
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    query,
+                    key,
+                    value,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                    video_length,
+                    context_length,
+                    prompt_length,
+                    unprompt_length,
+                )
+            else:
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(query, key, value, attn_metadata)
+
+            output_permuted = dynamic_block_sparse_fwd_flashinfer(
+                q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, is_cpu=False
+            )
+
+            attn_output = apply_inverse_permutation_triton(
+                output_permuted, q_sorted_indices, dim=2
+            )
+
+            res = attn_output.reshape(batch_size, num_heads, seq_len, dim).transpose(
+                1, 2
+            )
+
+        torch.backends.cuda.preferred_linalg_library(
+            backend="default"
+        )  # reset to default
+        return res.contiguous()

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

@@ -0,0 +1,330 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import functools
+import math
+from dataclasses import dataclass
+
+import torch
+
+try:
+    from vsa import video_sparse_attn
+except ImportError:
+    video_sparse_attn = None
+
+from typing import Any
+
+from sglang.multimodal_gen.runtime.distributed import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+VSA_TILE_SIZE = (4, 4, 4)
+
+
+@functools.lru_cache(maxsize=10)
+def get_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    T, H, W = dit_seq_shape
+    ts, hs, ws = tile_size
+    indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
+    ls = []
+    for t in range(math.ceil(T / ts)):
+        for h in range(math.ceil(H / hs)):
+            for w in range(math.ceil(W / ws)):
+                ls.append(
+                    indices[
+                        t * ts : min(t * ts + ts, T),
+                        h * hs : min(h * hs + hs, H),
+                        w * ws : min(w * ws + ws, W),
+                    ].flatten()
+                )
+    index = torch.cat(ls, dim=0)
+    return index
+
+
+@functools.lru_cache(maxsize=10)
+def get_reverse_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))
+
+
+@functools.lru_cache(maxsize=10)
+def construct_variable_block_sizes(
+    dit_seq_shape: tuple[int, int, int],
+    num_tiles: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    """
+    Compute the number of valid (non‑padded) tokens inside every
+    (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order
+    (t‑tile, h‑tile, w‑tile) that `rearrange` uses.
+
+    Returns
+    -------
+    torch.LongTensor  # shape: [∏ full_window_size]
+    """
+    # unpack
+    t, h, w = dit_seq_shape
+    ts_t, ts_h, ts_w = VSA_TILE_SIZE
+    n_t, n_h, n_w = num_tiles
+
+    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
+        """Vector with the size of each tile along one dimension."""
+        sizes = torch.full((n_tiles,), tile, dtype=torch.int, device=device)
+        # size of last (possibly partial) tile
+        remainder = dim_len - (n_tiles - 1) * tile
+        sizes[-1] = remainder if remainder > 0 else tile
+        return sizes
+
+    t_sizes = _sizes(t, ts_t, n_t)  # [n_t]
+    h_sizes = _sizes(h, ts_h, n_h)  # [n_h]
+    w_sizes = _sizes(w, ts_w, n_w)  # [n_w]
+
+    # broadcast‑multiply to get voxels per tile, then flatten
+    block_sizes = (
+        t_sizes[:, None, None]  # [n_t, 1,   1]
+        * h_sizes[None, :, None]  # [1,   n_h, 1]
+        * w_sizes[None, None, :]  # [1,   1,   n_w]
+    ).reshape(
+        -1
+    )  # [n_t * n_h * n_w]
+
+    return block_sizes
+
+
+@functools.lru_cache(maxsize=10)
+def get_non_pad_index(
+    variable_block_sizes: torch.LongTensor,
+    max_block_size: int,
+):
+    n_win = variable_block_sizes.shape[0]
+    device = variable_block_sizes.device
+    starts_pad = torch.arange(n_win, device=device) * max_block_size
+    index_pad = (
+        starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :]
+    )
+    index_mask = (
+        torch.arange(max_block_size, device=device)[None, :]
+        < variable_block_sizes[:, None]
+    )
+    return index_pad[index_mask]
+
+
+class VideoSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.VIDEO_SPARSE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
+        return VideoSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
+        return VideoSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
+        return VideoSparseAttentionMetadataBuilder
+
+
+@dataclass
+class VideoSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    dit_seq_shape: list[int]
+    VSA_sparsity: float
+    num_tiles: list[int]
+    total_seq_length: int
+    tile_partition_indices: torch.LongTensor
+    reverse_tile_partition_indices: torch.LongTensor
+    variable_block_sizes: torch.LongTensor
+    non_pad_index: torch.LongTensor
+    untile_combined_index: torch.LongTensor
+    tile_buf: torch.Tensor | None = None
+
+    # adaption for FastWan2.1-T2V-1.3B-Diffusers
+    # Sequence lengths for the forward batch
+    # Maximum sequence length for query
+    max_seqlen_q: int = 1
+    # Maximum sequence length for key
+    max_seqlen_k: int = 0
+
+
+class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        VSA_sparsity: float,
+        device: torch.device,
+        **kwargs: dict[str, Any],
+    ) -> VideoSparseAttentionMetadata:
+        patch_size = patch_size
+        dit_seq_shape = (
+            raw_latent_shape[0] // patch_size[0],
+            raw_latent_shape[1] // patch_size[1],
+            raw_latent_shape[2] // patch_size[2],
+        )
+
+        num_tiles = (
+            math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]),
+            math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
+            math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]),
+        )
+        total_seq_length = math.prod(dit_seq_shape)
+
+        tile_partition_indices = get_tile_partition_indices(
+            dit_seq_shape, VSA_TILE_SIZE, device
+        )
+        reverse_tile_partition_indices = get_reverse_tile_partition_indices(
+            dit_seq_shape, VSA_TILE_SIZE, device
+        )
+        variable_block_sizes = construct_variable_block_sizes(
+            dit_seq_shape, num_tiles, device
+        )
+        non_pad_index = get_non_pad_index(
+            variable_block_sizes, math.prod(VSA_TILE_SIZE)
+        )
+        untile_combined_index = non_pad_index[reverse_tile_partition_indices]
+
+        return VideoSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            dit_seq_shape=dit_seq_shape,  # type: ignore
+            VSA_sparsity=VSA_sparsity,  # type: ignore
+            num_tiles=num_tiles,  # type: ignore
+            total_seq_length=total_seq_length,  # type: ignore
+            tile_partition_indices=tile_partition_indices,  # type: ignore
+            reverse_tile_partition_indices=reverse_tile_partition_indices,
+            variable_block_sizes=variable_block_sizes,
+            non_pad_index=non_pad_index,
+            untile_combined_index=untile_combined_index,
+        )
+
+
+class VideoSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+
+    def tile(
+        self,
+        x: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        num_tiles = attn_metadata.num_tiles
+        t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
+        h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
+        w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
+        target_shape = (
+            x.shape[0],
+            t_padded_size * h_padded_size * w_padded_size,
+            x.shape[-2],
+            x.shape[-1],
+        )
+
+        buf = attn_metadata.tile_buf
+        if (
+            buf is None
+            or buf.shape != target_shape
+            or buf.dtype != x.dtype
+            or buf.device != x.device
+        ):
+            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
+            attn_metadata.tile_buf = buf
+
+        buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
+        return buf
+
+    def untile(
+        self,
+        x: torch.Tensor,
+        untile_combined_index: torch.LongTensor,
+    ) -> torch.Tensor:
+        return x[:, untile_combined_index]
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.tile(qkv, attn_metadata)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output, attn_metadata.untile_combined_index)
+
+    def forward(  # type: ignore[override]
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        gate_compress: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        gate_compress = gate_compress.transpose(1, 2).contiguous()
+
+        VSA_sparsity = attn_metadata.VSA_sparsity
+
+        cur_topk = math.ceil(
+            (1 - VSA_sparsity)
+            * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))
+        )
+
+        if video_sparse_attn is None:
+            raise NotImplementedError("video_sparse_attn is not installed")
+        hidden_states = video_sparse_attn(
+            query,
+            key,
+            value,
+            variable_block_sizes=attn_metadata.variable_block_sizes,
+            topk=cur_topk,
+            block_size=VSA_TILE_SIZE,
+            compress_attn_weight=gate_compress,
+        ).transpose(1, 2)
+
+        return hidden_states

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

@@ -0,0 +1,259 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from kernel.attn.vmoba_attn.vmoba import (
+    moba_attn_varlen,
+    process_moba_input,
+    process_moba_output,
+)
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class VMOBAAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.VMOBA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["VMOBAAttentionImpl"]:
+        return VMOBAAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoMobaAttentionMetadata"]:
+        return VideoMobaAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoMobaAttentionMetadataBuilder"]:
+        return VideoMobaAttentionMetadataBuilder
+
+
+@dataclass
+class VideoMobaAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+
+    temporal_chunk_size: int
+    temporal_topk: int
+    spatial_chunk_size: tuple[int, int]
+    spatial_topk: int
+    st_chunk_size: tuple[int, int, int]
+    st_topk: int
+
+    moba_select_mode: str
+    moba_threshold: float
+    moba_threshold_type: str
+    patch_resolution: list[int]
+
+    first_full_step: int = 12
+    first_full_layer: int = 0
+    # temporal_layer -> spatial_layer -> st_layer
+    temporal_layer: int = 1
+    spatial_layer: int = 1
+    st_layer: int = 1
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    """
+    Arguments:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.
+        batch: int, batch size for the padded sequence.
+        seqlen: int, maximum sequence length for the padded sequence.
+    Return:
+        hidden_states: (batch, seqlen, ...)
+    """
+    dim = hidden_states.shape[1:]
+    output = torch.zeros(
+        (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype
+    )
+    output[indices] = hidden_states
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        temporal_chunk_size: int,
+        temporal_topk: int,
+        spatial_chunk_size: tuple[int, int],
+        spatial_topk: int,
+        st_chunk_size: tuple[int, int, int],
+        st_topk: int,
+        moba_select_mode: str = "threshold",
+        moba_threshold: float = 0.25,
+        moba_threshold_type: str = "query_head",
+        device: torch.device = None,
+        first_full_layer: int = 0,
+        first_full_step: int = 12,
+        temporal_layer: int = 1,
+        spatial_layer: int = 1,
+        st_layer: int = 1,
+        **kwargs,
+    ) -> VideoMobaAttentionMetadata:
+        if device is None:
+            device = torch.device("cpu")
+        assert (
+            raw_latent_shape[0] % patch_size[0] == 0
+            and raw_latent_shape[1] % patch_size[1] == 0
+            and raw_latent_shape[2] % patch_size[2] == 0
+        ), f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}"
+        patch_resolution = [
+            t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False)
+        ]
+
+        return VideoMobaAttentionMetadata(
+            current_timestep=current_timestep,
+            temporal_chunk_size=temporal_chunk_size,
+            temporal_topk=temporal_topk,
+            spatial_chunk_size=spatial_chunk_size,
+            spatial_topk=spatial_topk,
+            st_chunk_size=st_chunk_size,
+            st_topk=st_topk,
+            moba_select_mode=moba_select_mode,
+            moba_threshold=moba_threshold,
+            moba_threshold_type=moba_threshold_type,
+            patch_resolution=patch_resolution,
+            first_full_layer=first_full_layer,
+            first_full_step=first_full_step,
+            temporal_layer=temporal_layer,
+            spatial_layer=spatial_layer,
+            st_layer=st_layer,
+        )
+
+
+class VMOBAAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads,
+        head_size,
+        softmax_scale,
+        causal=False,
+        num_kv_heads=None,
+        prefix="",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+        self.pad_input = pad_input
+
+    def _get_layer_idx(self, prefix: str) -> int | None:
+        match = re.search(r"blocks\.(\d+)", prefix)
+        if not match:
+            raise ValueError(f"Invalid prefix: {prefix}")
+        return int(match.group(1))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """
+        query: [B, L, H, D]
+        key:   [B, L, H, D]
+        value: [B, L, H, D]
+        attn_metadata: AttentionMetadata
+        """
+        batch_size, sequence_length, num_heads, head_dim = query.shape
+
+        # select chunk type according to layer idx:
+        loop_layer_num = (
+            attn_metadata.temporal_layer
+            + attn_metadata.spatial_layer
+            + attn_metadata.st_layer
+        )
+        moba_layer = self.layer_idx - attn_metadata.first_full_layer
+        if moba_layer % loop_layer_num < attn_metadata.temporal_layer:
+            moba_chunk_size = attn_metadata.temporal_chunk_size
+            moba_topk = attn_metadata.temporal_topk
+        elif (
+            moba_layer % loop_layer_num
+            < attn_metadata.temporal_layer + attn_metadata.spatial_layer
+        ):
+            moba_chunk_size = attn_metadata.spatial_chunk_size
+            moba_topk = attn_metadata.spatial_topk
+        elif (
+            moba_layer % loop_layer_num
+            < attn_metadata.temporal_layer
+            + attn_metadata.spatial_layer
+            + attn_metadata.st_layer
+        ):
+            moba_chunk_size = attn_metadata.st_chunk_size
+            moba_topk = attn_metadata.st_topk
+
+        query, chunk_size = process_moba_input(
+            query, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        key, chunk_size = process_moba_input(
+            key, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        value, chunk_size = process_moba_input(
+            value, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        max_seqlen = query.shape[1]
+        indices_q = torch.arange(
+            0, query.shape[0] * query.shape[1], device=query.device
+        )
+        cu_seqlens = torch.arange(
+            0,
+            query.shape[0] * query.shape[1] + 1,
+            query.shape[1],
+            dtype=torch.int32,
+            device=query.device,
+        )
+        query = rearrange(query, "b s ... -> (b s) ...")
+        key = rearrange(key, "b s ... -> (b s) ...")
+        value = rearrange(value, "b s ... -> (b s) ...")
+
+        # current_timestep=attn_metadata.current_timestep
+        hidden_states = moba_attn_varlen(
+            query,
+            key,
+            value,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=chunk_size,
+            moba_topk=moba_topk,
+            select_mode=attn_metadata.moba_select_mode,
+            simsum_threshold=attn_metadata.moba_threshold,
+            threshold_type=attn_metadata.moba_threshold_type,
+        )
+        hidden_states = self.pad_input(
+            hidden_states, indices_q, batch_size, sequence_length
+        )
+        hidden_states = process_moba_output(
+            hidden_states, attn_metadata.patch_resolution, moba_chunk_size
+        )
+
+        return hidden_states

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

@@ -0,0 +1,122 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from functools import lru_cache
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+try:
+    from sgl_kernel.flash_attn import flash_attn_varlen_func
+
+    flash_attn_func = flash_attn_varlen_func
+except ImportError as e:
+    raise e
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    FlashAttentionMetadataBuilder,
+)
+
+
+class XPUAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["XPUAttentionImpl"]:
+        return XPUAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        """XPU backend does not require special metadata."""
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+@lru_cache(maxsize=128)
+def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor:
+    return torch.arange(
+        0,
+        (bsz + 1) * seqlen,
+        step=seqlen,
+        device=torch.device("xpu", device_index),
+        dtype=torch.int32,
+    )
+
+
+class XPUAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        bsz, seqlen_q, nheads_q, d = tuple(query.shape)
+        _, seqlen_k, nheads_k, _ = tuple(key.shape)
+
+        max_seqlen_q = seqlen_q
+        max_seqlen_k = seqlen_k
+
+        q_ = query.contiguous().reshape(bsz * seqlen_q, nheads_q, d)
+        k_ = key.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        v_ = value.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        cu_q = _get_cu_seqlens(q_.device.index, bsz, seqlen_q)
+        cu_k = _get_cu_seqlens(q_.device.index, bsz, seqlen_k)
+
+        out = flash_attn_func(
+            q=q_,
+            k=k_,
+            v=v_,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+            return_softmax_lse=return_softmax_lse,
+        )
+
+        if return_softmax_lse:
+            out_tensor, softmax_lse = out[:2]
+            result = out_tensor.reshape(bsz, seqlen_q, nheads_q, d)
+            return result, softmax_lse
+
+        result = out.reshape(bsz, seqlen_q, nheads_q, d)
+        return result