From e776239afdd16c82ce340e64a6a50e8d48f79ff3 Mon Sep 17 00:00:00 2001 From: HuangJi <32611516+IPostYellow@users.noreply.github.com> Date: Wed, 21 Jan 2026 18:13:51 +0800 Subject: [PATCH] [diffusion] feat: support SageSparseLinearAttention attention backend (#17399) --- .../multimodal_gen/docs/support_matrix.md | 38 +- .../attention/backends/sparse_linear_attn.py | 548 ++++++++++++++---- .../runtime/layers/attention/turbo_layer.py | 13 +- .../runtime/models/dits/wanvideo.py | 7 +- .../multimodal_gen/runtime/platforms/cuda.py | 3 + .../runtime/platforms/interface.py | 1 + 6 files changed, 479 insertions(+), 131 deletions(-) diff --git a/python/sglang/multimodal_gen/docs/support_matrix.md b/python/sglang/multimodal_gen/docs/support_matrix.md index 2f9baec13..eb06afc4a 100644 --- a/python/sglang/multimodal_gen/docs/support_matrix.md +++ b/python/sglang/multimodal_gen/docs/support_matrix.md @@ -16,25 +16,27 @@ default parameters when initializing and generating videos. ### Video Generation Models -| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention（SLA） -|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:| -| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ | -| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | -| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | -| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | -| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | -| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | -| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | -| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | -| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | -| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | -| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | -| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | -| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | +| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention（SLA）| Sage Sparse Linear Attention（SageSLA）| +|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|:-----------------------------------------------:| +| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | +| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | +| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | +| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | +| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | +| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | +| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | +| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | -**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. +**Note**:
+1.Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
+2.SageSLA Based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` ### Image Generation Models 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 index 9781a9c93..c6c25ebf2 100644 --- 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 @@ -1,5 +1,18 @@ -# SPDX-License-Identifier: Apache-2.0 -# Adapted from turbo_layer.py for Attention Backend integration +""" +Copyright (c) 2025 by SLA team. + +Licensed under the Apache License, Version 2.0 (the "License"); + +This implementation is adapted from: from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py and https://github.com/thu-ml/SLA/blob/main/SageSLA/core.py +Citation (please cite if you use this code): + +@article{zhang2025sla, + title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention}, + author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and Joseph E. Gonzalez and Jun Zhu and Jianfei Chen}, + journal={arXiv preprint arXiv:2509.24006}, + year={2025} +} +""" from collections.abc import Callable from dataclasses import dataclass @@ -23,6 +36,136 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +# ==================================SLA Functions=================================== +def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): + arg_k = k - torch.mean( + k, dim=-2, keepdim=True + ) # smooth-k technique in SageAttention + pooled_qblocks = mean_pool(q, BLKQ) + pooled_kblocks = mean_pool(arg_k, BLKK) + pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + + K = pooled_score.shape[-1] + topk = min(K, int(topk_ratio * K)) + lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices + + sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) + sparse_map.scatter_(-1, lut, 1) + return sparse_map, lut, topk + + +def mean_pool(x, BLK): + assert x.is_contiguous() + + B, H, L, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + + grid = (L_BLOCKS, B * H) + compress_kernel[grid](x, x_mean, L, D, BLK) + return x_mean + + +@triton.jit +def compress_kernel( + X, + XM, + L: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_bh = tl.program_id(1) + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + + x_offset = idx_bh * L * D + xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D + x = tl.load( + X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L + ) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx + tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) + + +@triton.jit +def _attn_fwd( + Q, + K, + V, + qk_scale: tl.constexpr, + topk: tl.constexpr, + LUT, + LSE, + OS, + L: tl.constexpr, + M_BLOCKS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + qkv_offset = idx_bh * L * D + lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk + lse_offset = idx_bh * L + offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None] + V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + LUT_ptr = LUT + lut_offset + LSE_ptrs = LSE + lse_offset + offs_m + + m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32) + + q = tl.load(Q_ptrs, mask=offs_m[:, None] < L) + for block_idx in tl.range(topk): + idx_n = tl.load(LUT_ptr + block_idx) + n_mask = offs_n < L - idx_n * BLOCK_N + + k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :]) + qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2) + if L - idx_n * BLOCK_N < BLOCK_N: + qk = tl.where(n_mask[None, :], qk, float("-inf")) + + v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) + local_m = tl.max(qk, 1) + new_m = tl.maximum(m_i, local_m) + qk = qk - new_m[:, None] + + p = tl.math.exp2(qk) + l_ij = tl.sum(p, 1) + alpha = tl.math.exp2(m_i - new_m) + o_s = o_s * alpha[:, None] + o_s += tl.dot(p.to(v.dtype), v) + + l_i = l_i * alpha + l_ij + m_i = new_m + + o_s = o_s / l_i[:, None] + tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L) + + m_i += tl.math.log2(l_i) + tl.store(LSE_ptrs, m_i, mask=offs_m < L) + + +def _get_cuda_arch(device_index: int) -> str: + """Get CUDA architecture string for the given device.""" + major, minor = torch.cuda.get_device_capability(device_index) + return f"sm{major}{minor}" + + +# ==================================SLA Class=================================== class SparseLinearAttentionBackend(AttentionBackend): """Sparse Linear Attention Backend for efficient attention computation.""" @@ -241,123 +384,312 @@ class _attention(torch.autograd.Function): 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) +# ==================================SageSLA Class=================================== +SAGESLA_ENABLED = True +try: + import spas_sage_attn._fused as fused + import spas_sage_attn._qattn as qattn + from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant +except ImportError: + SAGESLA_ENABLED = False - 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 +SAGE2PP_ENABLED = True +try: + from spas_sage_attn._qattn import ( + qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold, ) - - 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)) +except ImportError: + SAGE2PP_ENABLED = False -@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) +class SageSparseLinearAttentionBackend(AttentionBackend): + """Quantized Sparse-Linear Attention backend using SageAttention kernels.""" - 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) + accept_output_buffer: bool = True - 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 + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] - 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) + @staticmethod + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.SAGE_SLA_ATTN - 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 + @staticmethod + def get_impl_cls() -> type["SageSparseLinearAttentionImpl"]: + return SageSparseLinearAttentionImpl - 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")) + @staticmethod + def get_metadata_cls() -> type["SageSparseLinearAttentionMetadata"]: + return SageSparseLinearAttentionMetadata - 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] + @staticmethod + def get_builder_cls() -> type["SageSparseLinearAttentionMetadataBuilder"]: + return SageSparseLinearAttentionMetadataBuilder - 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 +@dataclass +class SageSparseLinearAttentionMetadata(AttentionMetadata): + """Metadata for Sage Sparse Linear Attention computation.""" - o_s = o_s / l_i[:, None] - tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L) + # Basic attention parameters + current_timestep: int - m_i += tl.math.log2(l_i) - tl.store(LSE_ptrs, m_i, mask=offs_m < L) + # Sparse attention configuration + topk_ratio: float = 0.1 + + +class SageSparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder): + """Builder for SageSparseLinearAttentionMetadata.""" + + def __init__(self) -> None: + pass + + def prepare(self) -> None: + pass + + def build( + self, + current_timestep: int, + topk_ratio: float = 0.1, + **kwargs: dict[str, Any], + ) -> SageSparseLinearAttentionMetadata: + return SageSparseLinearAttentionMetadata( + current_timestep=current_timestep, + topk_ratio=topk_ratio, + ) + + +class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool = False, + softmax_scale: float | None = None, + num_kv_heads: int | None = None, + prefix: str = "", + topk_ratio: float = 0.5, + feature_map: str = "softmax", + use_bf16: bool = True, + **extra_impl_args, + ) -> None: + nn.Module.__init__(self) + + assert ( + SAGESLA_ENABLED + ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." + + self.num_heads = num_heads + self.head_size = head_size + self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5 + self.causal = causal + self.prefix = prefix + + self.topk_ratio = topk_ratio + self.dtype = torch.bfloat16 if use_bf16 else torch.float16 + + # Learnable linear projection for combining sparse + linear attention + self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32) + + # Feature map for linear attention + # Type annotation for callables + self.feature_map_q: Callable[[torch.Tensor], torch.Tensor] + self.feature_map_k: Callable[[torch.Tensor], torch.Tensor] + if feature_map == "elu": + self.feature_map_q = lambda x: F.elu(x) + 1 + self.feature_map_k = lambda x: F.elu(x) + 1 + elif feature_map == "relu": + self.feature_map_q = F.relu + self.feature_map_k = F.relu + elif feature_map == "softmax": + self.feature_map_q = lambda x: F.softmax(x, dim=-1) + self.feature_map_k = lambda x: F.softmax(x, dim=-1) + else: + raise ValueError(f"Unknown feature map: {feature_map}") + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize projection weights to zero for residual-like behavior.""" + with torch.no_grad(): + nn.init.zeros_(self.proj_l.weight) + nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type] + + def _calc_linear_attention_with_torch( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ): + kv = torch.matmul(k.transpose(-1, -2), v) + k_sum = torch.sum(k, dim=-2, keepdim=True) + return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2))) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """Forward pass for Sage Sparse Linear attention with quantized kernels. + Args: + query: query tensor of shape (B, L, H, D) + key: key tensor of shape (B, L, H, D) + value: value tensor of shape (B, L, H, D) + attn_metadata: attention metadata containing configuration + Returns: + output tensor of shape (B, L, H, D) + """ + dtype = query.dtype + + # Transpose from (B, L, H, D) to SLA format (B, H, L, D) + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + + # Determine block sizes based on GPU architecture + arch = _get_cuda_arch(q.device.index) + + if arch == "sm90": + BLKQ = 64 + BLKK = 128 + else: + BLKQ = 128 + BLKK = 64 + # Compute block-sparse attention pattern + sparse_map, lut, real_topk = get_block_map( + q, k, topk_ratio=self.topk_ratio, BLKQ=BLKQ, BLKK=BLKK + ) + + # Convert to compute dtype + q = q.to(self.dtype) + k = k.to(self.dtype) + v = v.to(self.dtype) + + ########## SPARGE BEGIN ########## + km = k.mean(dim=-2, keepdim=True) + headdim = q.size(-1) + assert headdim in [ + 64, + 128, + ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." + + # Quantize Q, K to INT8 + q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK) + lut, valid_block_num = block_map_lut_triton(sparse_map) + scale = 1.0 / (headdim**0.5) + + o_s = torch.empty_like(q) + + if arch in ("sm80", "sm86", "sm87"): + pvthreshold = torch.full( + (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device + ) + v_fp16 = v.to(torch.float16) + qattn.qk_int8_sv_f16_accum_f16_block_sparse_attn_inst_buf_with_pv_threshold( + q_int8, + k_int8, + v_fp16, + o_s, + lut, + valid_block_num, + pvthreshold, + q_scale, + k_scale, + 1, + False, + 1, + scale, + 0, + ) + else: + b, h_kv, kv_len, head_dim = v.shape + padded_len = (kv_len + 127) // 128 * 128 + v_transposed_permutted = torch.empty( + (b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device + ) + fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1) + v_fp8 = torch.empty( + v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device + ) + v_scale = torch.empty( + (b, h_kv, head_dim), dtype=torch.float32, device=v.device + ) + fused.scale_fuse_quant_cuda( + v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1 + ) + + if arch == "sm90": + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90( + q_int8, + k_int8, + v_fp8, + o_s, + lut, + valid_block_num, + q_scale, + k_scale, + v_scale, + 1, + False, + 1, + scale, + ) + else: + pvthreshold = torch.full( + (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device + ) + if SAGE2PP_ENABLED: + qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( + q_int8, + k_int8, + v_fp8, + o_s, + lut, + valid_block_num, + pvthreshold, + q_scale, + k_scale, + v_scale, + 1, + False, + 1, + scale, + 0, + ) + else: + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( + q_int8, + k_int8, + v_fp8, + o_s, + lut, + valid_block_num, + pvthreshold, + q_scale, + k_scale, + v_scale, + 1, + False, + 1, + scale, + 0, + ) + + ########## SPARGE END ########## + + # Linear attention with feature maps + q_linear = self.feature_map_q(q).contiguous().to(self.dtype) + k_linear = self.feature_map_k(k).contiguous().to(self.dtype) + o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v) + + # Project linear attention output and combine + with torch.amp.autocast("cuda", dtype=self.dtype): + o_l = self.proj_l(o_l) + + # Combine sparse and linear outputs + output = (o_s + o_l).to(dtype).transpose(1, 2) + + return output 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 4cd680eb8..89a8b5ceb 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py @@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionImpl, ) from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn import ( + SageSparseLinearAttentionBackend, SparseLinearAttentionBackend, ) from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend @@ -239,11 +240,17 @@ class MinimalA2AAttnOp(DistributedAttention): 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: + if attn_backend not in ( + SparseLinearAttentionBackend, + SageSparseLinearAttentionBackend, + ): logger.warning( - "TurboWan now only supports `sla_attn` and has been automatically set to `sla_attn`. Please set --attention-backend to `sla_attn`." + "TurboWan now only supports `sla_attn` or `sage_sla_attn` and has been automatically set to attention_type. Please set --attention-backend to `sla_attn` or `sage_sla_attn`." ) - attn_backend = SparseLinearAttentionBackend + if attention_type == "sagesla": + attn_backend = SageSparseLinearAttentionBackend + else: + attn_backend = SparseLinearAttentionBackend impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls() local_attn = impl_cls( num_heads=num_heads, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 31ed3c945..99e88a849 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -307,13 +307,16 @@ class WanTransformerBlock(nn.Module): self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True) - if attention_type == "sla": + if attention_type in ("sla", "sagesla"): self.attn1 = MinimalA2AAttnOp( 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}, + supported_attention_backends={ + AttentionBackendEnum.SLA_ATTN, + AttentionBackendEnum.SAGE_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 c2ec28a46..cebd7dd53 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -243,6 +243,9 @@ class CudaPlatformBase(Platform): 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 == AttentionBackendEnum.SAGE_SLA_ATTN: + logger.info("Using Sage Sparse Linear Attention backend") + return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend" 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 6f3a2ebe6..5f5c37ddb 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -34,6 +34,7 @@ class AttentionBackendEnum(enum.Enum): VMOBA_ATTN = enum.auto() AITER = enum.auto() SLA_ATTN = enum.auto() + SAGE_SLA_ATTN = enum.auto() NO_ATTENTION = enum.auto() def __str__(self):