[diffusion] model: support TurboWan2.1-T2V-1.3B/14B SLA (#15888)

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

@@ -30,6 +30,7 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2",
             r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
             r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn1\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.local_attn.proj_l.\2",
             r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
             r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
             r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
@@ -87,6 +88,8 @@ class WanVideoArchConfig(DiTArchConfig):
     )
     num_frames_per_block: int = 3
     sliding_window_num_frames: int = 21
+    attention_type: str = "original"
+    sla_topk: float = 0.1
 
     def __post_init__(self):
         super().__post_init__()

python/sglang/multimodal_gen/configs/pipeline_configs/wan.py

@@ -92,6 +92,16 @@ class WanT2V480PConfig(PipelineConfig):
         self.vae_config.load_decoder = True
 
 
+@dataclass
+class TurboWanT2V480PConfig(WanT2V480PConfig):
+    """Base configuration for Wan T2V 1.3B pipeline architecture."""
+
+    flow_shift: float | None = 8.0
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [988, 932, 852, 608]
+    )
+
+
 @dataclass
 class WanT2V720PConfig(WanT2V480PConfig):
     """Base configuration for Wan T2V 14B 720P pipeline architecture."""

python/sglang/multimodal_gen/registry.py

@@ -36,6 +36,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     FastWan2_1_T2V_480P_Config,
     FastWan2_2_TI2V_5B_Config,
+    TurboWanT2V480PConfig,
     Wan2_2_I2V_A14B_Config,
     Wan2_2_T2V_A14B_Config,
     Wan2_2_TI2V_5B_Config,
@@ -320,6 +321,13 @@ def _register_configs():
         ],
         model_detectors=[lambda hf_id: "wanpipeline" in hf_id.lower()],
     )
+    register_configs(
+        sampling_param_cls=WanT2V_1_3B_SamplingParams,
+        pipeline_config_cls=TurboWanT2V480PConfig,
+        hf_model_paths=[
+            "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
+        ],
+    )
     register_configs(
         sampling_param_cls=WanT2V_14B_SamplingParams,
         pipeline_config_cls=WanT2V720PConfig,
@@ -327,6 +335,13 @@ def _register_configs():
             "Wan-AI/Wan2.1-T2V-14B-Diffusers",
         ],
     )
+    register_configs(
+        sampling_param_cls=WanT2V_14B_SamplingParams,
+        pipeline_config_cls=TurboWanT2V480PConfig,
+        hf_model_paths=[
+            "IPostYellow/TurboWan2.1-T2V-14B-Diffusers",
+        ],
+    )
     register_configs(
         sampling_param_cls=WanI2V_14B_480P_SamplingParam,
         pipeline_config_cls=WanI2V480PConfig,

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

@@ -14,12 +14,18 @@ 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 (
+    MinimalA2AAttnOp,
+    SparseLinearAttention,
+)
 
 __all__ = [
     "USPAttention",
     "LocalAttention",
     "UlyssesAttention",
     "UlyssesAttention_VSA",
+    "MinimalA2AAttnOp",
+    "SparseLinearAttention",
     "AttentionBackend",
     "AttentionMetadata",
     "AttentionMetadataBuilder",

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

@@ -0,0 +1,500 @@
+# 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
+
+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
+
+
+def post_all2all(local_seq_2_local_head, seq_world_size):
+    def post_func(input):
+        # b, s, n, h
+        if local_seq_2_local_head:
+            output = rearrange(input, "w bs seq h d -> bs (w seq) h d")
+        else:
+            output = rearrange(input, "w bs s h d -> bs s (w h) d", w=seq_world_size)
+
+        return output
+
+    return post_func
+
+
+def single_all_to_all(input, local_seq_2_local_head, group, async_op=False):
+    seq_world_size = dist.get_world_size(group)
+
+    # b, s, n, h
+    if local_seq_2_local_head:
+        bs, local_seq_len, num_total_head, head_dim = input.shape
+        assert (
+            num_total_head % seq_world_size == 0
+        ), f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!"
+        input_t = rearrange(
+            input,
+            "bs seq_len (w h) d -> w bs seq_len h d",
+            w=seq_world_size,
+            h=num_total_head // seq_world_size,
+        ).contiguous()
+        post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size)
+    else:
+        bs, global_seq_len, num_local_head, head_dim = input.shape
+        input_t = rearrange(
+            input,
+            "bs (w s) h d -> w bs s h d",
+            w=seq_world_size,
+            s=global_seq_len // seq_world_size,
+        ).contiguous()
+        post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size)
+
+    output = torch.empty_like(input_t)
+    dist.all_to_all_single(output, input_t, group=group, async_op=async_op)
+
+    res = post_all2all_fun(output)
+    return res
+
+
+def async_a2a_communicate(
+    a2a_inputs: Union[torch.Tensor, List[torch.Tensor]],
+    cp_size: int,
+    cp_group: ProcessGroup,
+    cp_stream: torch.cuda.Stream,
+    local_seq_2_local_head: bool,
+) -> Union[torch.Tensor, List[torch.Tensor]]:
+    """
+    A2A communication for context parallelism. best used in communicate qkv
+    Modified from Nvidia Transformer Engine.
+    """
+    a2a_inputs = [a2a_inputs] if not isinstance(a2a_inputs, list) else a2a_inputs
+    a2a_outputs, a2a_reqs = [None] * len(a2a_inputs), [None] * len(a2a_inputs)
+    a2a_post_fns = [None] * len(a2a_inputs)
+    if local_seq_2_local_head:
+        for i in range(len(a2a_inputs) + 2):
+            if 0 < i < len(a2a_inputs) + 1:
+                a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1])
+                a2a_reqs[i - 1] = torch.distributed.all_to_all_single(
+                    a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True
+                )
+                a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size)
+            if i > 1:
+                with torch.cuda.stream(cp_stream):
+                    a2a_reqs[i - 2].wait()
+                    a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
+            if i < len(a2a_inputs):
+                a2a_inputs[i] = rearrange(
+                    a2a_inputs[i], "bs seq_len (w h) d -> w bs seq_len h d", w=cp_size
+                ).contiguous()
+    else:
+        for i in range(len(a2a_inputs) + 2):
+            if 0 < i < len(a2a_inputs) + 1:
+                a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1])
+                a2a_reqs[i - 1] = torch.distributed.all_to_all_single(
+                    a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True
+                )
+                a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size)
+            if i < len(a2a_inputs):
+                a2a_inputs[i] = rearrange(
+                    a2a_inputs[i], "bs (w s) h d -> w bs s h d", w=cp_size
+                ).contiguous()
+            if i > 1:
+                with torch.cuda.stream(cp_stream):
+                    a2a_reqs[i - 2].wait()
+                    a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
+    torch.cuda.current_stream().wait_stream(cp_stream)
+    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(
+        ctx: Any, group: dist.ProcessGroup, input: Tensor, local_seq_2_local_head: bool
+    ) -> Tensor:
+        ctx.group = group
+        res = single_all_to_all(input, local_seq_2_local_head, group, False)
+        ctx.local_seq_2_local_head = local_seq_2_local_head
+        return res
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None]:
+        return (
+            None,
+            _SeqAllToAll.apply(ctx.group, *grad_output, not ctx.local_seq_2_local_head),
+            None,
+        )
+
+
+class _SeqAllToAllQKV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        cp_size: int,
+        cp_stream: torch.cuda.Stream,
+        local_seq_2_local_head: bool,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        ctx.group = group
+        ctx.cp_size = cp_size
+        ctx.cp_stream = cp_stream
+        ctx.local_seq_2_local_head = local_seq_2_local_head
+        q, k, v = async_a2a_communicate(
+            [q, k, v], cp_size, group, cp_stream, local_seq_2_local_head
+        )
+        return q, k, v
+
+    @staticmethod
+    def backward(
+        ctx: Any, *grad_output: Tensor
+    ) -> Tuple[None, Tensor, Tensor, Tensor, None, None, None]:
+        q_grad, k_grad, v_grad = _SeqAllToAllQKV.apply(
+            ctx.group,
+            *grad_output,
+            ctx.cp_size,
+            ctx.cp_stream,
+            not ctx.local_seq_2_local_head,
+        )
+        return (None, q_grad, k_grad, v_grad, None, None, None)
+
+
+class DistributedAttention(torch.nn.Module):
+    """Initialization.
+
+    Arguments:
+        local_attention (Module): local attention with q,k,v
+        sequence_process_group (ProcessGroup): sequence parallel process group
+    """
+
+    def __init__(self, local_attention: Union[Module, Callable]) -> None:
+        super(DistributedAttention, self).__init__()
+        self.local_attn = local_attention
+        self.pg = None
+        self.stream = None
+
+    def forward(
+        self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
+    ) -> Tensor:
+        """forward
+
+        Arguments:
+            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)
+        pg_size = dist.get_world_size(self.pg)
+        if pg_size < 2:
+            return self.local_attn(query, key, value, *args, **kwargs)
+
+        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
+        )
+
+        output = _SeqAllToAll.apply(self.pg, context_layer, False)
+        return output
+
+    def set_context_parallel_group(self, group, stream):
+        self.pg = group
+        self.stream = stream
+
+
+class MinimalA2AAttnOp(DistributedAttention):
+    def __init__(self, local_attn=None, *args, **kwargs):
+        del args, kwargs
+        super(MinimalA2AAttnOp, self).__init__(local_attn)
+
+    def set_context_parallel_group(self, process_group, ranks, stream):
+        del ranks
+        super().set_context_parallel_group(process_group, stream)
+
+    def forward(
+        self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
+    ) -> Tensor:
+        results = super().forward(query, key, value, *args, **kwargs)
+        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":
+
+            def elu_feature_map(x):
+                return F.elu(x) + 1
+
+            self.feature_map_q = elu_feature_map
+            self.feature_map_k = elu_feature_map
+        elif feature_map == "relu":
+            self.feature_map_q = nn.ReLU()
+            self.feature_map_k = nn.ReLU()
+        elif feature_map == "softmax":
+
+            def softmax_feature_map(x):
+                return F.softmax(x, dim=-1)
+
+            self.feature_map_q = softmax_feature_map
+            self.feature_map_k = 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
+
+        def calc_linear(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))
+
+        o_l = 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
+
+
+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

@@ -13,6 +13,8 @@ from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
 from sglang.multimodal_gen.configs.sample.wan import WanTeaCacheParams
 from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
 from sglang.multimodal_gen.runtime.layers.attention import (
+    MinimalA2AAttnOp,
+    SparseLinearAttention,
     UlyssesAttention_VSA,
     USPAttention,
 )
@@ -262,6 +264,8 @@ class WanTransformerBlock(nn.Module):
         added_kv_proj_dim: int | None = None,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
         prefix: str = "",
+        attention_type: str = "original",
+        sla_topk: float = 0.1,
     ):
         super().__init__()
 
@@ -272,13 +276,20 @@ class WanTransformerBlock(nn.Module):
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
 
         self.to_out = ReplicatedLinear(dim, dim, bias=True)
-        self.attn1 = USPAttention(
+        if attention_type == "sla":
-            num_heads=num_heads,
+            self.attn1 = MinimalA2AAttnOp(
-            head_size=dim // num_heads,
+                SparseLinearAttention(
-            causal=False,
+                    dim // num_heads, topk=sla_topk, BLKQ=128, BLKK=64
-            supported_attention_backends=supported_attention_backends,
+                )
-            prefix=f"{prefix}.attn1",
+            )
-        )
+        else:
+            self.attn1 = USPAttention(
+                num_heads=num_heads,
+                head_size=dim // num_heads,
+                causal=False,
+                supported_attention_backends=supported_attention_backends,
+                prefix=f"{prefix}.attn1",
+            )
 
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
@@ -648,6 +659,8 @@ class WanTransformer3DModel(CachableDiT):
                     self._supported_attention_backends
                     | {AttentionBackendEnum.VIDEO_SPARSE_ATTN},
                     prefix=f"{config.prefix}.blocks.{i}",
+                    attention_type=config.attention_type,
+                    sla_topk=config.sla_topk,
                 )
                 for i in range(config.num_layers)
             ]