From dde0ecdb902345fffad4e84bc0336ac6ec92a2a1 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 2 Sep 2026 09:42:33 +0800 Subject: [PATCH] [diffusion] feat: support spargeattention (#37437) --- .../sglang-diffusion/attention_backends.mdx | 49 +++++ .../layers/attention/backends/sparge_attn.py | 156 ++++++++++++++++ .../runtime/models/dits/base.py | 1 + .../runtime/models/dits/qwen_image.py | 1 + .../multimodal_gen/runtime/platforms/cuda.py | 35 ++++ .../runtime/platforms/interface.py | 2 + .../test/unit/test_cuda_attention_backend.py | 44 ++++- .../unit/test_sparge_attention_backend.py | 176 ++++++++++++++++++ 8 files changed, 463 insertions(+), 1 deletion(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/attention/backends/sparge_attn.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_sparge_attention_backend.py diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index 5cc57575e..23e0752b8 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/docs/docs/sglang-diffusion/attention_backends.mdx @@ -64,6 +64,11 @@ For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBa `SAGE_ATTN_3` Requires SageAttention3 installed per upstream instructions. + + `sparge_attn` + `SPARGE_ATTN` + Training-free sparse SageAttention2. CUDA SM80/86/87/89/90, FP16/BF16, head dim 64/128, and square self-attention with sequence length at least 128. Other attention shapes use dense SDPA. Install pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation; configure retained block ratio with topk. + `sol_attn` `SOL_ATTN` @@ -213,6 +218,16 @@ Some backends require additional configuration. You can pass these parameters vi +**SpargeAttention (`sparge_attn`)** + +| Parameter | Type | Description | Default | +|---|---|---|---| +| `topk` | `float` | Fraction of predicted attention blocks retained per head. Higher values preserve more attention work and generally improve quality. Must be in `(0, 1]`. | `0.5` | + +SpargeAttention is approximate even when `topk=1`: the recommended upstream +kernel quantizes attention through SageAttention2. Validate output quality and +end-to-end latency on the target model and resolution before deployment. + **V-MoBA (`vmoba_attn`)** @@ -494,6 +509,16 @@ Some backends require additional configuration. You can pass these parameters vi + + + + + + + + + + @@ -678,6 +703,30 @@ dense switching; under ring parallelism the target must be ring-capable. Note `sage_attn` / `sage_attn_3` are lossy (quantized attention) — validate quality on your workload. +### Using SpargeAttention + +Install the optional CUDA extension, then select the backend explicitly: + +```bash +pip install ninja +pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation + +sglang generate \ + --model-path \ + --prompt "..." \ + --attention-backend sparge_attn \ + --attention-backend-config topk=0.5 +``` + +`sparge_attn` applies to compatible self-attention layers. Cross-attention uses +the existing dense fallback; short self-attention and asymmetric Q/KV also use +dense SDPA. LTX-2, LTX-2.3, and LTX-2.5 have compatible video (128) and audio +(64) head dimensions, so their sufficiently long, unmasked self-attention uses +SpargeAttention while prompt and audio-video cross-attention remain dense. +Ulysses sequence parallelism is supported, but ring attention is not because the +upstream kernel does not expose the softmax LSE needed for ring merging. K/V-gather +SP therefore follows the normal sparse-backend rule and uses Ulysses instead. + ### Sage then Sol hybrid `sol_attn` keeps the first `dense_steps` steps dense. Set diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparge_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparge_attn.py new file mode 100644 index 000000000..f402ad11d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparge_attn.py @@ -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() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index e5d545ca9..4d6a11af1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index be711a340..40ab1ab4f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -899,6 +899,7 @@ class QwenImageCrossAttention(nn.Module): AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.SAGE_ATTN_3, + AttentionBackendEnum.SPARGE_ATTN, }, ) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 590cb6cf8..1c1e54294 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 030dea6d7..a7bd1c8a8 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): 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, diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py index 3109175b9..603e95cb5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py +++ b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py @@ -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" diff --git a/python/sglang/multimodal_gen/test/unit/test_sparge_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_sparge_attention_backend.py new file mode 100644 index 000000000..e1422acf9 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_sparge_attention_backend.py @@ -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()
❌ CUDA-only (optional dependency).
`sparge_attn`YesNoNoNo❌❌CUDA SM80/86/87/89/90 only. Requires SpargeAttn; head dim 64/128 and square self-attention with sequence length at least 128.
`sol_attn` Yes