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

2026-09-20 16:40:28 +08:00
parent 414adef060
commit 2a0cb2f04e
6 changed files with 277 additions and 44 deletions
@@ -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()
@@ -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 (