diff --git a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py index 3a9939c4d..3430c001f 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py +++ b/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__() diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index 3a999683b..b9f5ea6a2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/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.""" diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index b81391170..e29261cfa 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py b/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py index 1b40782be..39852c384 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py new file mode 100644 index 000000000..4d28601b0 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index cee93f3bf..80f19d035 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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( - num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1", - ) + if attention_type == "sla": + self.attn1 = MinimalA2AAttnOp( + SparseLinearAttention( + dim // num_heads, topk=sla_topk, BLKQ=128, BLKK=64 + ) + ) + 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) ]