[diffusion] refactor: move SLA to attention_backend folder (#17020)

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

@@ -14,10 +14,7 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
     USPAttention,
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
-from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import (
+from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp
-    MinimalA2AAttnOp,
-    SparseLinearAttention,
-)
 
 __all__ = [
     "USPAttention",
@@ -25,7 +22,6 @@ __all__ = [
     "UlyssesAttention",
     "UlyssesAttention_VSA",
     "MinimalA2AAttnOp",
-    "SparseLinearAttention",
     "AttentionBackend",
     "AttentionMetadata",
     "AttentionMetadataBuilder",

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

@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from turbo_layer.py for Attention Backend integration
+
+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__)
+
+
+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).contiguous().to(self.dtype)  # c_q
+        key = self.feature_map_k(key).contiguous().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
+
+
+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)

python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py

@@ -1,18 +1,31 @@
 # copy and modify from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/rcm/utils/a2a_cp.py and https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py
 
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, List, Tuple, Type, Union
 
 import torch
 import torch.distributed as dist
-import torch.nn as nn
-import torch.nn.functional as F
-import triton
-import triton.language as tl
 from einops import rearrange
 from torch import Tensor
 from torch.distributed import ProcessGroup
 from torch.nn import Module
 
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionImpl,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn import (
+    SparseLinearAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    ForwardContext,
+    get_forward_context,
+)
+from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import get_compute_dtype
+
+logger = init_logger(__name__)
+
 
 def post_all2all(local_seq_2_local_head, seq_world_size):
     def post_func(input):
@@ -110,128 +123,6 @@ def async_a2a_communicate(
     return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs
 
 
-@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_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))
-
-
 class _SeqAllToAll(torch.autograd.Function):
     @staticmethod
     def forward(
@@ -301,7 +192,7 @@ class DistributedAttention(torch.nn.Module):
         self.stream = None
 
     def forward(
-        self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
+        self, query: Tensor, key: Tensor, value: Tensor, ctx_attn_metadata
     ) -> Tensor:
         """forward
 
@@ -309,22 +200,21 @@ class DistributedAttention(torch.nn.Module):
             query (Tensor): query input to the layer
             key (Tensor): key input to the layer
             value (Tensor): value input to the layer
-            args: other args
 
         Returns:
             * output (Tensor): context output
         """
         if self.pg is None:
-            return self.local_attn(query, key, value, *args, **kwargs)
+            return self.local_attn(query, key, value, ctx_attn_metadata)
         pg_size = dist.get_world_size(self.pg)
         if pg_size < 2:
-            return self.local_attn(query, key, value, *args, **kwargs)
+            return self.local_attn(query, key, value, ctx_attn_metadata)
 
         query_layer, key_layer, value_layer = _SeqAllToAllQKV.apply(
             self.pg, query, key, value, pg_size, self.stream, True
         )
         context_layer = self.local_attn(
-            query_layer, key_layer, value_layer, *args, **kwargs
+            query_layer, key_layer, value_layer, ctx_attn_metadata
         )
 
         output = _SeqAllToAll.apply(self.pg, context_layer, False)
@@ -336,8 +226,30 @@ class DistributedAttention(torch.nn.Module):
 
 
 class MinimalA2AAttnOp(DistributedAttention):
-    def __init__(self, local_attn=None, *args, **kwargs):
+    def __init__(
-        del args, kwargs
+        self,
+        num_heads: int,
+        head_size: int,
+        attention_type: str,
+        topk: float,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+    ):
+        dtype = get_compute_dtype()
+        attn_backend = get_attn_backend(
+            head_size, dtype, supported_attention_backends=supported_attention_backends
+        )
+        # Maintained for compatibility purposes; can be removed when CI allows setting Attention_backend or when TurboWan supports FA.
+        if attn_backend is not SparseLinearAttentionBackend:
+            logger.warning(
+                "TurboWan now only supports `sla_attn` and has been automatically set to `sla_attn`. Please set --attention-backend to `sla_attn`."
+            )
+            attn_backend = SparseLinearAttentionBackend
+        impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls()
+        local_attn = impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            topk_ratio=topk,
+        )
         super(MinimalA2AAttnOp, self).__init__(local_attn)
 
     def set_context_parallel_group(self, process_group, ranks, stream):
@@ -347,157 +259,7 @@ class MinimalA2AAttnOp(DistributedAttention):
     def forward(
         self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
     ) -> Tensor:
-        results = super().forward(query, key, value, *args, **kwargs)
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+        results = super().forward(query, key, value, ctx_attn_metadata)
         return rearrange(results, "b ... h l -> b ... (h l)")
-
-
-class SparseLinearAttention(nn.Module):
-    def __init__(
-        self,
-        head_dim,
-        topk,
-        feature_map="softmax",
-        BLKQ=64,
-        BLKK=64,
-        use_bf16=True,
-        tie_feature_map_qk=True,
-    ):
-        R"""
-        Args:
-            head_dim: dimension of each head.
-            topk: ratio of keys selected for sparse attention, shared across all queries.
-            feature_map: feature map for linear attention, one of ['hedgehog', 'elu', 'relu', 'softmax'].
-            BLKQ: block size for query.
-            BLKK: block size for key.
-            use_bf16: whether to use bfloat16 (default) or float16 for computation. The conversion to bf16/fp16 is done inside the module.
-            tie_feature_map_qk: whether to use the same feature map for query and key.
-        """
-        super().__init__()
-        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
-        self.topk = topk
-        self.BLKQ = BLKQ
-        self.BLKK = BLKK
-        self.proj_l = nn.Linear(head_dim, head_dim, dtype=torch.float32)
-
-        if feature_map == "elu":
-            self.feature_map_q = self._elu_feature_map
-            self.feature_map_k = self.elu_feature_map
-        elif feature_map == "relu":
-            self.feature_map_q = nn.ReLU()
-            self.feature_map_k = nn.ReLU()
-        elif feature_map == "softmax":
-            self.feature_map_q = self._softmax_feature_map
-            self.feature_map_k = self._softmax_feature_map
-        else:
-            raise NotImplementedError(f"Not supported feature map {feature_map}.")
-
-        if tie_feature_map_qk:
-            self.feature_map_k = self.feature_map_q
-
-        self.init_weights_()
-
-    def init_weights_(self):
-        with torch.no_grad():
-            nn.init.zeros_(self.proj_l.weight)
-            nn.init.zeros_(self.proj_l.bias)
-
-    def forward(self, q, k, v, return_sparsity=False):
-        R"""
-        Args:
-            q: queries of shape (B, H, L, D).
-            k: keys of shape (B, H, L, D).
-            v: values of shape (B, H, L, D).
-            return_sparsity: whether to return the actual sparsity.
-        """
-        dtype = q.dtype
-
-        q = q.transpose(1, 2).contiguous()
-        k = k.transpose(1, 2).contiguous()
-        v = v.transpose(1, 2).contiguous()
-
-        sparse_map, lut, real_topk = get_block_map(
-            q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK
-        )
-
-        q = q.to(self.dtype)
-        k = k.to(self.dtype)
-        v = v.to(self.dtype)
-        o_s = _attention.apply(
-            q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
-        )
-
-        q = self.feature_map_q(q).contiguous().to(self.dtype)  # c_q
-        k = self.feature_map_k(k).contiguous().to(self.dtype)  # c_k
-
-        o_l = self._torch_calc_linear(q, k, v)
-
-        with torch.amp.autocast("cuda", dtype=self.dtype):
-            o_l = self.proj_l(o_l)
-        o = (o_s + o_l).to(dtype).transpose(1, 2)
-
-        if return_sparsity:
-            return o, real_topk / sparse_map.shape[-1]
-        else:
-            return o
-
-    def _torch_calc_linear(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 _calc_linear(self, q, k, v):
-        kvsum = k.transpose(-1, -2) @ v
-        ksum = torch.sum(k, dim=-2, keepdim=True)
-        return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
-
-    def _softmax_feature_map(self, x):
-        return F.softmax(x, dim=-1)
-
-    def _elu_feature_map(self, x):
-        return F.elu(x) + 1
-
-
-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

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

@@ -18,7 +18,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
 )
 from sglang.multimodal_gen.runtime.layers.attention import (
     MinimalA2AAttnOp,
-    SparseLinearAttention,
     UlyssesAttention_VSA,
     USPAttention,
 )
@@ -311,9 +310,11 @@ class WanTransformerBlock(nn.Module):
         self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True)
         if attention_type == "sla":
             self.attn1 = MinimalA2AAttnOp(
-                SparseLinearAttention(
+                num_heads=divide(num_heads, get_tensor_model_parallel_world_size()),
-                    dim // num_heads, topk=sla_topk, BLKQ=128, BLKK=64
+                head_size=dim // num_heads,
-                )
+                attention_type=attention_type,
+                topk=sla_topk,
+                supported_attention_backends={AttentionBackendEnum.SLA_ATTN},
             )
         else:
             self.attn1 = USPAttention(

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

@@ -245,6 +245,9 @@ class CudaPlatformBase(Platform):
         elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
             logger.info("Using Torch SDPA backend")
             return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        elif selected_backend == AttentionBackendEnum.SLA_ATTN:
+            logger.info("Using Sparse Linear Attention backend")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend"
         elif selected_backend in [
             AttentionBackendEnum.FA,
         ]:

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

@@ -33,6 +33,7 @@ class AttentionBackendEnum(enum.Enum):
     VIDEO_SPARSE_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()
     AITER = enum.auto()
+    SLA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()
 
     def __str__(self):