minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support spargeattention (#37437)

Browse Source
This commit is contained in:
Mick
2026-09-02 09:42:33 +08:00
committed by GitHub
parent dc276264cb
commit dde0ecdb90
8 changed files with 463 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/sparge_attn.py
+156
View File
@@ -0,0 +1,156 @@
# SPDX-License-Identifier: Apache-2.0
from numbers import Real
from typing import Any
import torch
from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
AttentionBackend,
AttentionImpl,
AttentionMetadata,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
_SUPPORTED_HEAD_SIZES = (64, 128)
_MIN_SEQUENCE_LENGTH = 128
_DEFAULT_TOPK = 0.5
def _validate_topk(value: Any) -> float:
if isinstance(value, bool) or not isinstance(value, Real):
raise ValueError("SpargeAttention 'topk' must be a number in (0, 1].")
topk = float(value)
if not 0.0 < topk <= 1.0:
raise ValueError("SpargeAttention 'topk' must be in (0, 1].")
return topk
def _dense_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
*,
softmax_scale: float,
causal: bool,
) -> torch.Tensor:
return (
torch.nn.functional.scaled_dot_product_attention(
query.transpose(1, 2),
key.transpose(1, 2),
value.transpose(1, 2),
dropout_p=0.0,
is_causal=causal,
scale=softmax_scale,
)
.transpose(1, 2)
.contiguous()
)
class SpargeAttentionBackend(AttentionBackend):
"""Training-free sparse, quantized attention backed by SpargeAttention."""
@staticmethod
def get_supported_head_sizes() -> tuple[int, ...]:
return _SUPPORTED_HEAD_SIZES
@staticmethod
def get_enum() -> AttentionBackendEnum:
return AttentionBackendEnum.SPARGE_ATTN
@staticmethod
def get_impl_cls() -> type["SpargeAttentionImpl"]:
return SpargeAttentionImpl
class SpargeAttentionImpl(AttentionImpl):
def __init__(
self,
num_heads: int,
head_size: int,
softmax_scale: float,
causal: bool = False,
num_kv_heads: int | None = None,
prefix: str = "",
**extra_impl_args,
) -> None:
if head_size not in _SUPPORTED_HEAD_SIZES:
raise ValueError(
"SpargeAttention supports head sizes 64 and 128, "
f"but received {head_size}."
)
if num_kv_heads is not None and num_kv_heads != num_heads:
raise ValueError(
"SpargeAttention does not support grouped-query attention."
)
self.num_heads = num_heads
self.head_size = head_size
self.softmax_scale = softmax_scale
self.causal = causal
config = get_global_server_args().attention_backend_config or {}
self.topk = _validate_topk(config.get("topk", _DEFAULT_TOPK))
self._attention_op = spas_sage2_attn_meansim_topk_cuda
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: AttentionMetadata,
) -> torch.Tensor:
if query.device.type != "cuda":
raise ValueError("SpargeAttention requires CUDA tensors.")
if query.dtype not in (torch.float16, torch.bfloat16):
raise ValueError("SpargeAttention requires FP16 or BF16 inputs.")
if key.device != query.device or value.device != query.device:
raise ValueError("SpargeAttention requires Q, K, and V on one device.")
if key.dtype != query.dtype or value.dtype != query.dtype:
raise ValueError("SpargeAttention requires Q, K, and V with one dtype.")
if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
raise ValueError(
"SpargeAttention expects Q, K, and V in "
"[batch, sequence, heads, head_dim] layout."
)
if (
query.shape[0] != key.shape[0]
or query.shape[0] != value.shape[0]
or query.shape[2:] != key.shape[2:]
or query.shape[2:] != value.shape[2:]
or key.shape[1] != value.shape[1]
or query.shape[-1] != self.head_size
):
raise ValueError(
"SpargeAttention requires compatible Q, K, and V batch, head, "
f"and head-dim shapes with head_dim={self.head_size}."
)
# The upstream sparse kernel is square self-attention only and requires
# at least one 128-token block. LTX can produce shorter audio sequences
# and asymmetric Q/KV in its replicated-audio SP path.
if (
query.shape != key.shape
or query.shape != value.shape
or query.shape[1] < _MIN_SEQUENCE_LENGTH
):
return _dense_attention(
query,
key,
value,
softmax_scale=self.softmax_scale,
causal=self.causal,
)
output = self._attention_op(
query,
key,
value,
is_causal=self.causal,
scale=self.softmax_scale,
tensor_layout="NHD",
topk=self.topk,
)
return output.contiguous()
python/sglang/multimodal_gen/runtime/models/dits/base.py
+1
View File
@@ -40,6 +40,7 @@ class BaseDiT(nn.Module, ABC):
_supported_attention_backends: set[AttentionBackendEnum] = {
AttentionBackendEnum.SLIDING_TILE_ATTN,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.SPARGE_ATTN,
AttentionBackendEnum.FA,
AttentionBackendEnum.AITER,
AttentionBackendEnum.AITER_SAGE,
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+1
View File
@@ -899,6 +899,7 @@ class QwenImageCrossAttention(nn.Module):
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.SAGE_ATTN_3,
AttentionBackendEnum.SPARGE_ATTN,
},
)
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+35
View File
@@ -200,6 +200,40 @@ class _SageAttention3BackendResolver(_CudaAttentionBackendResolver):
return AttentionBackendEnum.TORCH_SDPA
class _SpargeAttentionBackendResolver(_CudaAttentionBackendResolver):
backend = AttentionBackendEnum.SPARGE_ATTN
supported_capabilities = {(8, 0), (8, 6), (8, 7), (8, 9), (9, 0)}
@classmethod
def resolve(cls, platform) -> str:
capability = platform.get_device_capability()
capability_tuple = (
(capability.major, capability.minor) if capability is not None else None
)
if capability_tuple not in cls.supported_capabilities:
found = capability.as_version_str() if capability else "unknown"
raise ValueError(
"SpargeAttention supports CUDA compute capabilities "
f"8.0, 8.6, 8.7, 8.9, and 9.0; found {found}."
)
try:
from spas_sage_attn import ( # noqa: F401
spas_sage2_attn_meansim_topk_cuda,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.sparge_attn import ( # noqa: F401
SpargeAttentionBackend,
)
return "sglang.multimodal_gen.runtime.layers.attention.backends.sparge_attn.SpargeAttentionBackend"
except ImportError as e:
raise ImportError(
"SpargeAttention is not installed. Install it with "
"`pip install git+https://github.com/thu-ml/SpargeAttn.git "
"--no-build-isolation`."
) from e
class _VideoSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN
@@ -383,6 +417,7 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = {
_SlidingTileAttentionBackendResolver,
_SageAttentionBackendResolver,
_SageAttention3BackendResolver,
_SpargeAttentionBackendResolver,
_VideoSparseAttentionBackendResolver,
_SparseVideoGen2AttentionBackendResolver,
_SolAttnBackendResolver,
python/sglang/multimodal_gen/runtime/platforms/interface.py
+2
View File
@@ -33,6 +33,7 @@ class AttentionBackendEnum(enum.Enum):
DYNAMIC_CUDNN_SDPA = enum.auto()
SAGE_ATTN = enum.auto()
SAGE_ATTN_3 = enum.auto()
SPARGE_ATTN = enum.auto()
VIDEO_SPARSE_ATTN = enum.auto()
SPARSE_VIDEO_GEN_2_ATTN = enum.auto()
VMOBA_ATTN = enum.auto()
@@ -59,6 +60,7 @@ class AttentionBackendEnum(enum.Enum):
AttentionBackendEnum.VMOBA_ATTN,
AttentionBackendEnum.SLA_ATTN,
AttentionBackendEnum.SAGE_SLA_ATTN,
AttentionBackendEnum.SPARGE_ATTN,
AttentionBackendEnum.LASER_ATTN,
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
AttentionBackendEnum.RAIN_FUSION_ATTN,
python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py
+43 -1
View File
@@ -11,8 +11,12 @@ from sglang.multimodal_gen.runtime.layers.attention.selector import (
from sglang.multimodal_gen.runtime.platforms.cuda import (
CudaPlatformBase,
_SageAttentionBackendResolver,
_SpargeAttentionBackendResolver,
)
from sglang.multimodal_gen.runtime.platforms.interface import (
AttentionBackendEnum,
DeviceCapability,
)
from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
SDPA_BACKEND_CLS_STR = (
"sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
@@ -24,6 +28,7 @@ class FakeCudaPlatform(CudaPlatformBase):
is_blackwell_device = False
is_hopper_device = False
supports_flash_attention = True
device_capability = DeviceCapability(8, 0)
@classmethod
def is_sm120(cls):
@@ -45,6 +50,10 @@ class FakeCudaPlatform(CudaPlatformBase):
) -> bool:
return cls.supports_flash_attention
@classmethod
def get_device_capability(cls, device_id: int = 0):
return cls.device_capability
class TestCudaAttentionBackendSelection(unittest.TestCase):
def setUp(self):
@@ -52,6 +61,7 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
FakeCudaPlatform.is_blackwell_device = False
FakeCudaPlatform.is_hopper_device = False
FakeCudaPlatform.supports_flash_attention = True
FakeCudaPlatform.device_capability = DeviceCapability(8, 0)
_cached_get_attn_backend.cache_clear()
def resolve(
@@ -142,6 +152,38 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
AttentionBackendEnum.FA,
)
def test_sparge_attention_resolver(self):
module = types.ModuleType("spas_sage_attn")
module.spas_sage2_attn_meansim_topk_cuda = object()
backend_module = (
"sglang.multimodal_gen.runtime.layers.attention.backends.sparge_attn"
)
try:
with patch.dict(sys.modules, {"spas_sage_attn": module}):
self.assertEqual(
self.resolve(AttentionBackendEnum.SPARGE_ATTN),
f"{backend_module}.SpargeAttentionBackend",
)
finally:
sys.modules.pop(backend_module, None)
def test_sparge_attention_rejects_pre_ampere_cuda(self):
FakeCudaPlatform.device_capability = DeviceCapability(7, 5)
with self.assertRaisesRegex(ValueError, "found 7.5"):
_SpargeAttentionBackendResolver.resolve(FakeCudaPlatform)
def test_sparge_attention_rejects_unsupported_blackwell_cuda(self):
FakeCudaPlatform.device_capability = DeviceCapability(10, 0)
with self.assertRaisesRegex(ValueError, "found 10.0"):
_SpargeAttentionBackendResolver.resolve(FakeCudaPlatform)
def test_sparge_attention_missing_dependency_fails_closed(self):
with patch.dict(sys.modules, {"spas_sage_attn": None}):
with self.assertRaisesRegex(
ImportError, "SpargeAttention is not installed"
):
_SpargeAttentionBackendResolver.resolve(FakeCudaPlatform)
def test_explicit_backend_rejected_by_a_model_fails_closed(self):
with self.assertRaisesRegex(
ValueError, "not supported by this attention layer"
python/sglang/multimodal_gen/test/unit/test_sparge_attention_backend.py
+176
View File
@@ -0,0 +1,176 @@
import importlib
import sys
import types
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch
import torch
from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig
from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
_sparge = types.ModuleType("spas_sage_attn")
_sparge.spas_sage2_attn_meansim_topk_cuda = Mock()
_BACKEND_MODULE_NAME = (
"sglang.multimodal_gen.runtime.layers.attention.backends.sparge_attn"
)
with patch.dict(sys.modules, {"spas_sage_attn": _sparge}):
_backend_module = importlib.import_module(_BACKEND_MODULE_NAME)
sys.modules.pop(_BACKEND_MODULE_NAME, None)
SpargeAttentionBackend = _backend_module.SpargeAttentionBackend
SpargeAttentionImpl = _backend_module.SpargeAttentionImpl
_validate_topk = _backend_module._validate_topk
class TestSpargeAttentionBackend(unittest.TestCase):
def _make_impl(self, *, head_size: int = 128, topk: float = 0.5):
op = Mock()
with (
patch.object(
_backend_module,
"get_global_server_args",
return_value=SimpleNamespace(attention_backend_config={"topk": topk}),
),
patch.object(
_backend_module,
"spas_sage2_attn_meansim_topk_cuda",
op,
),
):
impl = SpargeAttentionImpl(
num_heads=8,
head_size=head_size,
softmax_scale=head_size**-0.5,
)
return impl, op
def test_enum_and_generic_dit_capability(self):
self.assertEqual(str(AttentionBackendEnum.SPARGE_ATTN), "sparge_attn")
self.assertTrue(AttentionBackendEnum.SPARGE_ATTN.is_sparse)
self.assertIn(
AttentionBackendEnum.SPARGE_ATTN,
BaseDiT._supported_attention_backends,
)
self.assertIs(
SpargeAttentionBackend.get_enum(), AttentionBackendEnum.SPARGE_ATTN
)
def test_ltx2_self_attention_head_dims_are_supported(self):
arch = LTX2ArchConfig()
self.assertIn(
arch.attention_head_dim,
SpargeAttentionBackend.get_supported_head_sizes(),
)
self.assertIn(
arch.audio_attention_head_dim,
SpargeAttentionBackend.get_supported_head_sizes(),
)
def test_topk_validation(self):
self.assertEqual(_validate_topk(0.5), 0.5)
self.assertEqual(_validate_topk(0.75), 0.75)
for invalid in (True, "0.5", 0.0, -0.1, 1.1):
with self.subTest(invalid=invalid):
with self.assertRaisesRegex(ValueError, "topk"):
_validate_topk(invalid)
def test_head_size_validation(self):
with self.assertRaisesRegex(ValueError, "head sizes 64 and 128"):
self._make_impl(head_size=96)
def test_grouped_query_attention_is_rejected(self):
with patch.object(
_backend_module,
"get_global_server_args",
return_value=SimpleNamespace(attention_backend_config={}),
):
with self.assertRaisesRegex(ValueError, "grouped-query"):
SpargeAttentionImpl(
num_heads=8,
num_kv_heads=4,
head_size=128,
softmax_scale=128**-0.5,
)
def test_forward_uses_nhd_layout_and_configured_topk(self):
impl, op = self._make_impl(topk=0.75)
qkv = SimpleNamespace(
device=torch.device("cuda"),
dtype=torch.bfloat16,
shape=(1, 128, 8, 128),
ndim=4,
)
expected = op.return_value.contiguous.return_value
self.assertIs(impl.forward(qkv, qkv, qkv, None), expected)
op.assert_called_once_with(
qkv,
qkv,
qkv,
is_causal=False,
scale=128**-0.5,
tensor_layout="NHD",
topk=0.75,
)
op.return_value.contiguous.assert_called_once_with()
def test_short_sequence_falls_back_to_dense_attention(self):
impl, op = self._make_impl()
qkv = SimpleNamespace(
device=torch.device("cuda"),
dtype=torch.float16,
shape=(1, 127, 8, 128),
ndim=4,
)
dense_output = object()
with patch.object(
_backend_module, "_dense_attention", return_value=dense_output
) as dense_attention:
self.assertIs(impl.forward(qkv, qkv, qkv, None), dense_output)
dense_attention.assert_called_once_with(
qkv,
qkv,
qkv,
softmax_scale=128**-0.5,
causal=False,
)
op.assert_not_called()
def test_asymmetric_sequence_falls_back_to_dense_attention(self):
impl, op = self._make_impl(head_size=64)
query = SimpleNamespace(
device=torch.device("cuda"),
dtype=torch.bfloat16,
shape=(1, 256, 8, 64),
ndim=4,
)
key_value = SimpleNamespace(
device=torch.device("cuda"),
dtype=torch.bfloat16,
shape=(1, 512, 8, 64),
ndim=4,
)
dense_output = object()
with patch.object(
_backend_module, "_dense_attention", return_value=dense_output
) as dense_attention:
self.assertIs(impl.forward(query, key_value, key_value, None), dense_output)
dense_attention.assert_called_once_with(
query,
key_value,
key_value,
softmax_scale=64**-0.5,
causal=False,
)
op.assert_not_called()
if __name__ == "__main__":
unittest.main()