From 2a0cb2f04edbd85778f1fb8c26272c7f668d9c34 Mon Sep 17 00:00:00 2001 From: HuangJi <32611516+IPostYellow@users.noreply.github.com> Date: Sun, 20 Sep 2026 16:40:28 +0800 Subject: [PATCH] [Diffusion][MiniMax-H3] Add SM120 Sage compute for SubBlock sparse attention (#40116) --- .../configs/pipeline_configs/minimax_h3.py | 19 ++- .../backends/subblock_sparse/README.md | 24 +++- .../backends/subblock_sparse_attn.py | 78 +++++++++--- .../unit/test_subblock_sparse_attention.py | 29 +++-- .../attention/test_subblock_sage_fp8_sm120.py | 51 ++++++++ .../cpu/test_subblock_sparse_attention.py | 120 ++++++++++++++++-- 6 files changed, 277 insertions(+), 44 deletions(-) create mode 100644 test/manual/attention/test_subblock_sage_fp8_sm120.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py index db7557bd6..9d9b47757 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py @@ -267,7 +267,7 @@ class MiniMaxH3PipelineConfig(PipelineConfig): ) if compute_mode == "sage_fp8": capability = current_platform.get_device_capability() - if capability is None or capability.to_int() != 90: + if capability is None or capability.to_int() not in (90, 120): found = ( capability.as_version_str() if capability is not None @@ -275,14 +275,21 @@ class MiniMaxH3PipelineConfig(PipelineConfig): ) raise ValueError( "MiniMax-H3 SubBlock compute_mode='sage_fp8' currently " - "requires SM90 (compute capability 9.0); " + "requires SM90 or SM120 (compute capability 9.0 or 12.0); " f"found {found}." ) - from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( - _load_sparge_attention_sm90_ops, - ) + if capability.to_int() == 90: + from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + _load_sparge_attention_sm90_ops, + ) - _load_sparge_attention_sm90_ops() + _load_sparge_attention_sm90_ops() + else: + from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( + _load_sm120_sage_ops, + ) + + _load_sm120_sage_ops() if selected_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3: if server_args.ring_degree > 1: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md index 68f2185d8..ac773b816 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md @@ -4,8 +4,8 @@ Routes a SubBlock plan to SGLang's CuTe-DSL block-sparse FlashAttention kernel on SM90 or FlashInfer's architecture-specific blk64 kernels on SM100 and SM120. The architecture-neutral `"compute_mode": "sage_fp8"` selects online Sage-style FP8 compute; its current SM90 implementation uses the native SageAttention2 -INT8-QK/FP8-PV Hopper kernel. Future SM100 and SM120 Sage FP8 implementations -can register under the same configuration value. +INT8-QK/FP8-PV Hopper kernel. SM120 uses FlashInfer's CuTe-DSL Sage kernel +with Q64 x K64 blocks. SM100 supports BF16 only. Nothing is trained and no weights change: a cheap estimator runs before attention and hands the selected kernel a `q2k_block_index`. @@ -46,7 +46,7 @@ are listed below. | | | | --- | --- | -| GPU | **compute capability 9.0, 10.0, or 12.0** — H100 / H200 use SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel or the current `sage_fp8` implementation; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` BF16 kernel; SM120 devices use FlashInfer's `bsa_attn_sm120_blk64_fwd` CuTe-DSL BF16 kernel. Other capabilities, including 10.3 (B300 / GB300), are rejected. | +| GPU | **compute capability 9.0, 10.0, or 12.0** — H100 / H200 use SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel or the current `sage_fp8` implementation; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` BF16 kernel; SM120 devices use FlashInfer's CuTe-DSL BF16 or Sage FP8 kernel. Other capabilities, including 10.3 (B300 / GB300), are rejected. | | dtype | bfloat16 | | head_dim | 128 | | attention | non-causal, one contiguous sequence per call | @@ -106,7 +106,17 @@ the normal router's 16-token key pooling cells. Install the optional dependency: pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation ``` -Enable it on SM90 with: +SM120 requires FlashInfer's CuTe-DSL SM120 Sage backend and +`quantize_sage_qkv_sm120`, added in +[FlashInfer #4691](https://github.com/flashinfer-ai/flashinfer/pull/4691). +The default FlashInfer `0.6.18` pin does not include these APIs. After installing +SGLang, install the FlashInfer source revision used for SM120 validation: + +```bash +pip install "git+https://github.com/flashinfer-ai/flashinfer.git@6a84331eb6013e5e61018dc2be532ae90520d30f" +``` + +Enable it on SM90 or SM120 with: ```bash --attention-backend-config '{"compute_mode":"sage_fp8"}' @@ -195,3 +205,9 @@ Tests: `test/unit/test_subblock_sparse_attention.py` and test covers the native production dispatch. Running at a full block budget must reproduce dense attention up to the expected quantization error, pinning routing indices, ragged tails, scale domains and the softmax scale in one check. + +Run the SM120 Sage numerical test with: + +```bash +python test/manual/attention/test_subblock_sage_fp8_sm120.py +``` diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py index 6433bb361..0fa30f918 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py @@ -21,10 +21,9 @@ individual keys of the defaults below:: Requirements inherited from the kernels: compute capability 9.0 (Hopper) or 10.0/12.0 (Blackwell), bf16, head_dim 128. Hopper uses SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel by default; ``compute_mode="sage_fp8"`` -selects its native SageAttention2 INT8-QK/FP8-PV kernel. The mode name is stable -across architectures, so future SM100/SM120 implementations can use their own -Sage FP8 arithmetic without changing server configuration. B200 and SM120 -devices currently use FlashInfer's architecture-specific blk64 BF16 kernels. +selects native SageAttention2 on SM90 and FlashInfer's CuTe-DSL Sage kernel +on SM120, both using INT8 QK and FP8 PV. SM100 uses the BF16 kernel because +its Sage path does not support SubBlock's variable counts and partial blocks. Inside the DiT, unsupported calls run dense instead; unsupported GPU architectures are rejected. @@ -110,6 +109,13 @@ DEFAULT_COMPUTE_MODE = "bf16" _DIT_LAYER_PREFIX = re.compile(r"^blocks\.(\d+)\.") +def _sage_key_block_size() -> int: + # Resolve on the worker's current device, after device initialization. + if torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0): + return SUBBLOCK_SPARSE_BLOCK_SIZE + return SAGE_FP8_SM90_KEY_BLOCK_SIZE + + def _dit_layer_index(prefix: str) -> int | None: match = _DIT_LAYER_PREFIX.match(prefix) return int(match.group(1)) if match else None @@ -233,6 +239,50 @@ def _sm100_sparse_attention( return out[0] if isinstance(out, tuple) else out +@functools.lru_cache(maxsize=1) +def _load_sm120_sage_ops(): + try: + from flashinfer.cute_dsl.sparse.bsa_attn_sm120 import ( + bsa_attn_sm120_blk64_sage_fwd, + ) + from flashinfer.cute_dsl.sparse.bsa_utils.sage_quant_sm120 import ( + quantize_sage_qkv_sm120, + ) + except (ImportError, OSError) as exc: + raise ImportError( + "SM120 SubBlock sage_fp8 requires FlashInfer with the CuTe-DSL " + "SM120 Sage backend and quantize_sage_qkv_sm120." + ) from exc + return quantize_sage_qkv_sm120, bsa_attn_sm120_blk64_sage_fwd + + +def _sm120_sage_fp8_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_block_index: torch.Tensor, + topk: int, + softmax_scale: float, + block_counts: torch.Tensor | None = None, +) -> torch.Tensor: + """Quantize BSHD activations online and execute a Q64 x K64 Sage plan.""" + quantize, attention = _load_sm120_sage_ops() + q_hnd, k_hnd, v_hnd = (x.transpose(1, 2).contiguous() for x in (q, k, v)) + quantized = quantize(q_hnd, k_hnd, v_hnd) + out = attention( + *quantized, + q2k_block_index.contiguous(), + topk, + block_sizes=_cached_block_sizes(k.shape[1], k.device), + q2k_block_nums=( + block_counts.contiguous() if block_counts is not None else None + ), + softmax_scale=softmax_scale, + backend="cute_dsl", + ) + return out.transpose(1, 2).contiguous() + + def _sm120_sparse_attention( q: torch.Tensor, k: torch.Tensor, @@ -273,16 +323,13 @@ def _get_subblock_sparse_attention_runner( if capability == (10, 0): if compute_mode != "bf16": raise RuntimeError( - f"SubBlock compute_mode={compute_mode!r} currently targets SM90; " - "the SM100 FlashInfer Sage adapter is not wired in SGLang yet." + "SM100 FlashInfer Sage does not support SubBlock variable block " + "counts and partial blocks; use compute_mode='bf16'." ) return _sm100_sparse_attention if capability == (12, 0): - if compute_mode != "bf16": - raise RuntimeError( - f"SubBlock compute_mode={compute_mode!r} currently targets SM90; " - "the SM120 FlashInfer Sage adapter is not wired in SGLang yet." - ) + if compute_mode == "sage_fp8": + return _sm120_sage_fp8_sparse_attention return _sm120_sparse_attention raise RuntimeError( "SubBlock sparse attention supports compute capability 9.0, 10.0, or 12.0; " @@ -378,9 +425,10 @@ class SubBlockSparseSchedule(msgspec.Struct, frozen=True): config = get_global_server_args().attention_backend_config or {} compute_mode = str(config.get("compute_mode", DEFAULT_COMPUTE_MODE)) - # Hopper's native kernel consumes 128-token K blocks. Eight sub-blocks - # preserve the default router's 16-token key pooling cells. - default_n_k = 8 if compute_mode == "sage_fp8" else DEFAULT_N_K + # Use 16-token pooling cells for both K128 (SM90) and K64 (SM120). + default_n_k = ( + _sage_key_block_size() // 16 if compute_mode == "sage_fp8" else DEFAULT_N_K + ) schedule = SubBlockSparseSchedule( sparsity=float(config.get("sparsity", DEFAULT_SPARSITY)), skip_first_steps=int( @@ -453,7 +501,7 @@ class SubBlockSparseAttentionImpl(AttentionImpl): n_k=self.schedule.n_k, n_q=self.schedule.n_q, block_size_k=( - SAGE_FP8_SM90_KEY_BLOCK_SIZE + _sage_key_block_size() if self.schedule.compute_mode == "sage_fp8" else SUBBLOCK_SPARSE_BLOCK_SIZE ), diff --git a/python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py b/python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py index 4dd2b0311..05dcb8f25 100644 --- a/python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py +++ b/python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py @@ -167,9 +167,15 @@ class TestSubBlockSparseSchedule(unittest.TestCase): self.assertEqual(schedule.compute_mode, "bf16") def test_sage_fp8_uses_16_token_key_subblocks_by_default(self): - with _patch_schedule({"compute_mode": "sage_fp8"}): - schedule = SubBlockSparseSchedule.from_server_args() - self.assertEqual(schedule.n_k, 8) + for capability, expected_n_k in (((9, 0), 8), ((12, 0), 4)): + with ( + self.subTest(capability=capability), + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=capability), + _patch_schedule({"compute_mode": "sage_fp8"}), + ): + schedule = SubBlockSparseSchedule.from_server_args() + self.assertEqual(schedule.n_k, expected_n_k) def test_explicit_sage_fp8_n_k_is_respected(self): with _patch_schedule({"compute_mode": "sage_fp8", "n_k": 4}): @@ -320,11 +326,18 @@ class TestSubBlockGating(unittest.TestCase): with _patch_step(20): self.assertFalse(impl._sparse_ready(q, q)) - def test_sage_fp8_builds_the_sm90_64x128_router(self): - impl = self._impl("blocks.9.attn", compute_mode="sage_fp8") - self.assertEqual(impl.router.block_size_k, 128) - self.assertEqual(impl.router.n_k, 8) - self.assertEqual(impl.router.budget_granularity, 1) + def test_sage_fp8_builds_architecture_specific_router(self): + for capability, key_block_size, n_k in (((9, 0), 128, 8), ((12, 0), 64, 4)): + with ( + self.subTest(capability=capability), + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=capability), + ): + impl = self._impl("blocks.9.attn", compute_mode="sage_fp8") + self.assertEqual(impl.router.block_size_k, key_block_size) + self.assertEqual(impl.router.n_k, n_k) + self.assertEqual(impl.router.block_size_k // impl.router.n_k, 16) + self.assertEqual(impl.router.budget_granularity, 1) @requires_subblock_kernel diff --git a/test/manual/attention/test_subblock_sage_fp8_sm120.py b/test/manual/attention/test_subblock_sage_fp8_sm120.py new file mode 100644 index 000000000..6efed2880 --- /dev/null +++ b/test/manual/attention/test_subblock_sage_fp8_sm120.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Manual SM120 Sage correctness check against selected-block FP32 attention. + +Requires an SM120 GPU and FlashInfer with the CuTe-DSL SM120 Sage backend. +Run: python test/manual/attention/test_subblock_sage_fp8_sm120.py +""" + +import unittest + +import torch + +from sglang.test.test_utils import CustomTestCase + + +@unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0), + "requires an SM120 GPU", +) +class TestSubBlockSageFp8Sm120(CustomTestCase): + def test_ragged_sparse_plan_with_empty_rows(self): + from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( + _sm120_sage_fp8_sparse_attention, + ) + + torch.manual_seed(19) + q = torch.randn(1, 65, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 129, 2, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + index = torch.tensor( + [[[[2, 0], [1, 0]], [[0, 2], [2, 1]]]], device="cuda", dtype=torch.int32 + ) + counts = torch.tensor([[[2, 1], [0, 2]]], device="cuda", dtype=torch.int32) + mask = torch.zeros(1, 2, 65, 129, device="cuda", dtype=torch.bool) + mask[0, 0, :64, :64] = True + mask[0, 0, :64, 128:] = True + mask[0, 0, 64:, 64:128] = True + mask[0, 1, 64:, 64:] = True + scale = 128**-0.5 + logits = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * scale + probs = logits.masked_fill(~mask, -float("inf")).softmax(-1).nan_to_num() + expected = torch.einsum("bhqk,bkhd->bqhd", probs, v.float()) + actual = _sm120_sage_fp8_sparse_attention(q, k, v, index, 2, scale, counts) + self.assertEqual(actual.dtype, torch.bfloat16) + self.assertTrue(actual.is_contiguous()) + self.assertTrue(torch.isfinite(actual).all()) + self.assertEqual(torch.count_nonzero(actual[0, :64, 1]).item(), 0) + torch.testing.assert_close(actual.float(), expected, atol=5e-2, rtol=5e-2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/cpu/test_subblock_sparse_attention.py b/test/registered/cpu/test_subblock_sparse_attention.py index 5a8daa8bc..a72504ba4 100644 --- a/test/registered/cpu/test_subblock_sparse_attention.py +++ b/test/registered/cpu/test_subblock_sparse_attention.py @@ -17,9 +17,11 @@ from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( SubBlockSparseAttentionImpl, _get_subblock_sparse_attention_runner, + _sage_key_block_size, _sm90_sage_fp8_sparse_attention, _sm90_sparse_attention, _sm100_sparse_attention, + _sm120_sage_fp8_sparse_attention, _sm120_sparse_attention, ) from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( @@ -160,11 +162,11 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase): 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): + def test_rejects_sage_fp8_on_sm100(self): device = torch.device("cuda:0") with ( patch("torch.cuda.get_device_capability", return_value=(10, 0)), - self.assertRaisesRegex(RuntimeError, "currently targets SM90"), + self.assertRaisesRegex(RuntimeError, "does not support SubBlock"), ): _get_subblock_sparse_attention_runner(device, "sage_fp8") @@ -175,6 +177,48 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase): self.assertIs(runner, _sm120_sparse_attention) + def test_sm120_sage_dispatch_and_block_geometry(self): + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=(12, 0)), + ): + self.assertEqual(_sage_key_block_size(), 64) + self.assertIs( + _get_subblock_sparse_attention_runner( + torch.device("cuda:0"), "sage_fp8" + ), + _sm120_sage_fp8_sparse_attention, + ) + + def test_sm120_sage_preserves_partial_blocks_and_variable_counts(self): + q = torch.randn(1, 65, 2, 128, dtype=torch.bfloat16) + k = torch.randn(1, 129, 2, 128, dtype=torch.bfloat16) + v = torch.randn_like(k) + index = torch.tensor([[[[2, 0], [1, 0]], [[0, 2], [2, 1]]]], dtype=torch.int32) + counts = torch.tensor([[[2, 1], [0, 2]]], dtype=torch.int32) + quantized = tuple(object() for _ in range(6)) + quantize = Mock(return_value=quantized) + attention = Mock(return_value=q.transpose(1, 2).contiguous()) + with patch( + "sglang.multimodal_gen.runtime.layers.attention.backends." + "subblock_sparse_attn._load_sm120_sage_ops", + return_value=(quantize, attention), + ): + result = _sm120_sage_fp8_sparse_attention(q, k, v, index, 2, 0.125, counts) + torch.testing.assert_close(result, q) + for actual, source in zip(quantize.call_args.args, (q, k, v)): + torch.testing.assert_close(actual, source.transpose(1, 2)) + self.assertTrue(actual.is_contiguous()) + self.assertEqual(attention.call_args.args[:6], quantized) + torch.testing.assert_close(attention.call_args.args[6], index) + kwargs = attention.call_args.kwargs + torch.testing.assert_close( + kwargs["block_sizes"], torch.tensor([64, 64, 1], dtype=torch.int32) + ) + torch.testing.assert_close(kwargs["q2k_block_nums"], counts) + self.assertEqual(kwargs["backend"], "cute_dsl") + self.assertEqual(kwargs["softmax_scale"], 0.125) + def test_platform_resolver_loads_sm120_dependency(self): capability = Mock(major=12, minor=0) capability.as_version_str.return_value = "12.0" @@ -232,14 +276,6 @@ 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)): @@ -323,6 +359,68 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase): get_backend.assert_not_called() + def test_sm120_sage_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(12, 0), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.layers.attention.backends." + "subblock_sparse_attn._load_sm120_sage_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_sm120_sage_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(12, 0), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.layers.attention.backends." + "subblock_sparse_attn._load_sm120_sage_ops", + side_effect=ImportError("FlashInfer SM120 Sage backend is unavailable"), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_attn_backend" + ) as get_backend, + self.assertRaisesRegex( + ImportError, "FlashInfer SM120 Sage backend is unavailable" + ), + ): + 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") @@ -348,7 +446,7 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase): loader.assert_not_called() - def test_sage_fp8_rejects_non_sm90_during_server_validation(self): + def test_sage_fp8_rejects_sm100_during_server_validation(self): config = MiniMaxH3PipelineConfig() server_args = self._subblock_server_args("sage_fp8") with (