chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] feat: support SageSparseLinearAttention attention backend (#17399)

Browse Source
This commit is contained in:
HuangJi
2026-01-21 18:13:51 +08:00
committed by GitHub
parent 1b97fa769b
commit e776239afd
6 changed files with 479 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
python/sglang/multimodal_gen/docs/support_matrix.md
View File

@@ -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<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280<br>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<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280<br>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**: <br>
1.Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.<br>
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

548
python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py
View File

@@ -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

13
python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py
View File

@@ -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,

7
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -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(

3
python/sglang/multimodal_gen/runtime/platforms/cuda.py
View File

@@ -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,
]:

1
python/sglang/multimodal_gen/runtime/platforms/interface.py
View File

@@ -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):