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

[Diffusion][MiniMax-H3] Add SM90 Sage compute for SubBlock sparse attention (#37982)

Browse Source
This commit is contained in:
HuangJi
2026-09-09 20:00:03 +08:00
committed by GitHub
parent dba34cc964
commit ffe98a4279
10 changed files with 900 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_subblock_sparse_attention.py
+231 -2
View File
@@ -1,17 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
import sys
import unittest
from types import SimpleNamespace
from types import ModuleType, SimpleNamespace
from unittest.mock import Mock, patch
import torch
from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
_load_sparge_attention_sm90_ops,
_routing_plan_to_block_map,
)
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
MiniMaxH3PipelineConfig,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
SubBlockSparseAttentionImpl,
_get_subblock_sparse_attention_runner,
_sm90_sage_fp8_sparse_attention,
_sm90_sparse_attention,
_sm100_sparse_attention,
_sm120_sparse_attention,
@@ -40,12 +46,87 @@ from sglang.multimodal_gen.runtime.platforms import (
from sglang.multimodal_gen.runtime.platforms.cuda import (
_SubBlockSparseAttentionBackendResolver,
)
from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=9, suite="stage-a-test-cpu-intel")
class TestSubBlockSageFp8PlanAdapter(CustomTestCase):
def test_variable_counts_ignore_each_row_suffix(self):
index = torch.tensor([[[[2, 0, 1, 3], [3, 1, 0, 2]]]], dtype=torch.int32)
counts = torch.tensor([[[3, 1]]], dtype=torch.int32)
block_map = _routing_plan_to_block_map(index, counts, 4)
torch.testing.assert_close(
block_map,
torch.tensor([[[[True, True, True, False], [False, False, False, True]]]]),
)
class TestSubBlockSageFp8DependencyLoader(CustomTestCase):
def setUp(self):
_load_sparge_attention_sm90_ops.cache_clear()
self.addCleanup(_load_sparge_attention_sm90_ops.cache_clear)
def test_missing_dependency_has_install_help(self):
missing_modules = {
name: None
for name in (
"spas_sage_attn",
"spas_sage_attn._fused",
"spas_sage_attn._qattn",
"spas_sage_attn.utils",
)
}
with (
patch.dict(sys.modules, missing_modules),
self.assertRaisesRegex(ImportError, "pip install.*SpargeAttn"),
):
_load_sparge_attention_sm90_ops()
def test_complete_dependency_exposes_required_ops(self):
package = ModuleType("spas_sage_attn")
package.__path__ = []
fused = ModuleType("spas_sage_attn._fused")
qattn = ModuleType("spas_sage_attn._qattn")
utils = ModuleType("spas_sage_attn.utils")
fused.transpose_pad_permute_cuda = Mock()
fused.scale_fuse_quant_cuda = Mock()
qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90 = (
Mock()
)
utils.block_map_lut_triton = Mock()
utils.get_vanilla_qk_quant = Mock()
package._fused = fused
package._qattn = qattn
package.utils = utils
with patch.dict(
sys.modules,
{
"spas_sage_attn": package,
"spas_sage_attn._fused": fused,
"spas_sage_attn._qattn": qattn,
"spas_sage_attn.utils": utils,
},
):
ops = _load_sparge_attention_sm90_ops()
self.assertEqual(
ops,
(
utils.get_vanilla_qk_quant,
utils.block_map_lut_triton,
fused.transpose_pad_permute_cuda,
fused.scale_fuse_quant_cuda,
qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90,
),
)
class TestSubBlockSparseAttentionDispatch(CustomTestCase):
def setUp(self):
_get_subblock_sparse_attention_runner.cache_clear()
@@ -70,6 +151,23 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase):
self.assertIs(runner, _sm100_sparse_attention)
def test_dispatches_sm90_sage_fp8_independently_from_bf16(self):
device = torch.device("cuda:0")
with patch("torch.cuda.get_device_capability", return_value=(9, 0)):
bf16_runner = _get_subblock_sparse_attention_runner(device, "bf16")
sage_runner = _get_subblock_sparse_attention_runner(device, "sage_fp8")
self.assertIs(bf16_runner, _sm90_sparse_attention)
self.assertIs(sage_runner, _sm90_sage_fp8_sparse_attention)
def test_rejects_sm90_sage_fp8_on_sm100_until_adapter_is_wired(self):
device = torch.device("cuda:0")
with (
patch("torch.cuda.get_device_capability", return_value=(10, 0)),
self.assertRaisesRegex(RuntimeError, "currently targets SM90"),
):
_get_subblock_sparse_attention_runner(device, "sage_fp8")
def test_dispatches_sm120(self):
device = torch.device("cuda:0")
with patch("torch.cuda.get_device_capability", return_value=(12, 0)):
@@ -134,6 +232,14 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase):
self.assertIs(kwargs["q2k_block_nums"], block_counts)
self.assertEqual(kwargs["softmax_scale"], 0.125)
def test_rejects_sage_fp8_on_sm120_until_adapter_is_wired(self):
device = torch.device("cuda:0")
with (
patch("torch.cuda.get_device_capability", return_value=(12, 0)),
self.assertRaisesRegex(RuntimeError, "currently targets SM90"),
):
_get_subblock_sparse_attention_runner(device, "sage_fp8")
def test_rejects_unsupported_compute_capability(self):
device = torch.device("cuda:0")
with patch("torch.cuda.get_device_capability", return_value=(10, 3)):
@@ -145,6 +251,128 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase):
class TestSubBlockSparseAttentionModalities(CustomTestCase):
@staticmethod
def _subblock_server_args(compute_mode: str):
return SimpleNamespace(
attention_backend="subblock_sparse_attn",
attention_backend_config={"compute_mode": compute_mode},
ring_degree=1,
resolve_component_attention_backend=lambda *_names: (
AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
"transformer",
),
)
def test_sage_fp8_dependency_is_checked_during_server_validation(self):
config = MiniMaxH3PipelineConfig()
server_args = self._subblock_server_args("sage_fp8")
loader = Mock()
with (
patch.object(current_platform, "is_mps", return_value=False),
patch.object(
current_platform,
"get_device_capability",
return_value=DeviceCapability(9, 0),
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_global_forced_attn_backend",
return_value=None,
),
patch(
"sglang.kernels.ops.attention.subblock_sage_fp8_sm90."
"_load_sparge_attention_sm90_ops",
loader,
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_attn_backend"
),
):
config.validate_server_args(server_args)
loader.assert_called_once_with()
def test_missing_sage_fp8_dependency_fails_server_validation(self):
config = MiniMaxH3PipelineConfig()
server_args = self._subblock_server_args("sage_fp8")
with (
patch.object(current_platform, "is_mps", return_value=False),
patch.object(
current_platform,
"get_device_capability",
return_value=DeviceCapability(9, 0),
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_global_forced_attn_backend",
return_value=None,
),
patch(
"sglang.kernels.ops.attention.subblock_sage_fp8_sm90."
"_load_sparge_attention_sm90_ops",
side_effect=ImportError("Install SpargeAttention"),
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_attn_backend"
) as get_backend,
self.assertRaisesRegex(ImportError, "Install SpargeAttention"),
):
config.validate_server_args(server_args)
get_backend.assert_not_called()
def test_bf16_does_not_require_sparge_attention(self):
config = MiniMaxH3PipelineConfig()
server_args = self._subblock_server_args("bf16")
loader = Mock()
with (
patch.object(current_platform, "is_mps", return_value=False),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_global_forced_attn_backend",
return_value=None,
),
patch(
"sglang.kernels.ops.attention.subblock_sage_fp8_sm90."
"_load_sparge_attention_sm90_ops",
loader,
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_attn_backend"
),
):
config.validate_server_args(server_args)
loader.assert_not_called()
def test_sage_fp8_rejects_non_sm90_during_server_validation(self):
config = MiniMaxH3PipelineConfig()
server_args = self._subblock_server_args("sage_fp8")
with (
patch.object(current_platform, "is_mps", return_value=False),
patch.object(
current_platform,
"get_device_capability",
return_value=DeviceCapability(10, 0),
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_global_forced_attn_backend",
return_value=None,
),
patch(
"sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
"get_attn_backend"
) as get_backend,
self.assertRaisesRegex(ValueError, "requires SM90.*found 10.0"),
):
config.validate_server_args(server_args)
get_backend.assert_not_called()
def test_transformer_subblock_with_ring_fails_admission(self):
config = MiniMaxH3PipelineConfig()
server_args = SimpleNamespace(
@@ -400,7 +628,7 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase):
impl = object.__new__(SubBlockSparseAttentionImpl)
impl.softmax_scale = 2**-0.5
impl.causal = False
impl.schedule = SimpleNamespace(sparsity=0.75)
impl.schedule = SimpleNamespace(sparsity=0.75, compute_mode="bf16")
plan = SimpleNamespace(
index=torch.tensor(
[[[[7, 1, 4], [6, 2, 0], [5, 0, 3]]]], dtype=torch.int32
@@ -419,6 +647,7 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase):
for runner, sparse_rows in (
(_sm90_sparse_attention, ([1, 4, 7], [0, 3, 5])),
(_sm90_sage_fp8_sparse_attention, ([7, 1, 4], [5, 0, 3])),
(_sm100_sparse_attention, ([7, 1, 4], [5, 0, 3])),
(_sm120_sparse_attention, ([7, 1, 4], [5, 0, 3])),
):
test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py
+140
View File
@@ -0,0 +1,140 @@
# SPDX-License-Identifier: Apache-2.0
"""Correctness tests for native SM90 SubBlock Sage FP8 attention."""
import math
import unittest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
def _native_sm90_sage_available() -> bool:
try:
import spas_sage_attn._qattn as qattn
except (ImportError, OSError):
return False
return hasattr(
qattn,
"qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90",
)
requires_native_sm90_sage = unittest.skipUnless(
torch.cuda.is_available()
and torch.cuda.get_device_capability() == (9, 0)
and _native_sm90_sage_available(),
"requires SM90 and a compiled SpargeAttention installation",
)
def _masked_reference(q, k, v, index, counts, scale, key_block_size=128):
logits = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * scale
mask = torch.zeros_like(logits, dtype=torch.bool)
for b in range(q.shape[0]):
for h in range(q.shape[2]):
for qb in range(index.shape[2]):
q_slice = slice(qb * 64, min((qb + 1) * 64, q.shape[1]))
for slot in range(int(counts[b, h, qb])):
kb = int(index[b, h, qb, slot])
k_slice = slice(
kb * key_block_size,
min((kb + 1) * key_block_size, k.shape[1]),
)
mask[b, h, q_slice, k_slice] = True
logits.masked_fill_(~mask, -float("inf"))
p = torch.softmax(logits, dim=-1)
return torch.einsum("bhqk,bkhd->bqhd", p, v.float()).to(torch.bfloat16)
def _cosine(a, b):
return float(
torch.nn.functional.cosine_similarity(
a.float().flatten(), b.float().flatten(), dim=0
)
)
@requires_native_sm90_sage
class TestSubBlockSageFp8NativeSm90(CustomTestCase):
def test_full_budget_ragged_tail_reproduces_dense(self):
from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
subblock_sage_fp8_sm90_attention,
)
torch.manual_seed(19)
seq_len = 1024 + 37
heads = 4
shape = (1, seq_len, heads, 128)
q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
scale = 1.0 / math.sqrt(128)
query_blocks = math.ceil(seq_len / 64)
key_blocks = math.ceil(seq_len / 128)
index = (
torch.arange(key_blocks, device="cuda", dtype=torch.int32)
.view(1, 1, 1, key_blocks)
.expand(1, heads, query_blocks, key_blocks)
.contiguous()
)
output = subblock_sage_fp8_sm90_attention(q, k, v, index, key_blocks, scale)
reference = torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
scale=scale,
).transpose(1, 2)
self.assertTrue(torch.isfinite(output.float()).all())
self.assertGreater(_cosine(output, reference), 0.998)
def test_sparse_k128_plan_and_variable_counts(self):
from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
subblock_sage_fp8_sm90_attention,
)
torch.manual_seed(23)
seq_len = 1024 + 37
heads = 4
shape = (1, seq_len, heads, 128)
q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
scale = 1.0 / math.sqrt(128)
query_blocks = math.ceil(seq_len / 64)
key_blocks = math.ceil(seq_len / 128)
width = 4
index = torch.stack(
[
torch.roll(torch.arange(key_blocks), shifts=query_block)[:width]
for query_block in range(query_blocks)
]
)
index = (
index.to(device="cuda", dtype=torch.int32)
.view(1, 1, query_blocks, width)
.expand(1, heads, query_blocks, width)
.contiguous()
)
counts = (
((torch.arange(query_blocks, device="cuda", dtype=torch.int32) % width) + 1)
.view(1, 1, query_blocks)
.expand(1, heads, query_blocks)
.contiguous()
)
output = subblock_sage_fp8_sm90_attention(q, k, v, index, width, scale, counts)
reference = _masked_reference(q, k, v, index, counts, scale)
self.assertTrue(torch.isfinite(output.float()).all())
self.assertGreater(_cosine(output, reference), 0.997)
if __name__ == "__main__":
unittest.main(verbosity=3)