diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py b/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py index 39852c384..aacb7412b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py +++ b/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 ( - MinimalA2AAttnOp, - SparseLinearAttention, -) +from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp __all__ = [ "USPAttention", @@ -25,7 +22,6 @@ __all__ = [ "UlyssesAttention", "UlyssesAttention_VSA", "MinimalA2AAttnOp", - "SparseLinearAttention", "AttentionBackend", "AttentionMetadata", "AttentionMetadataBuilder", diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py new file mode 100644 index 000000000..9781a9c93 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py index 3074dc802..4cd680eb8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/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): - del args, kwargs + def __init__( + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 5e8e7e714..444e42472 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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( - dim // num_heads, topk=sla_topk, BLKQ=128, BLKK=64 - ) + num_heads=divide(num_heads, get_tensor_model_parallel_world_size()), + head_size=dim // num_heads, + attention_type=attention_type, + topk=sla_topk, + supported_attention_backends={AttentionBackendEnum.SLA_ATTN}, ) else: self.attn1 = USPAttention( diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 3deff525d..49358322a 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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, ]: diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index abdb55dc0..283381ab9 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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):