[diffusion][Minimax H3]support subblock sparse attention on SM90 (#34680)

test/registered/cpu/test_subblock_sparse_attention.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
+    _get_subblock_sparse_attention_runner,
+    _sm90_sparse_attention,
+    _sm100_sparse_attention,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+
+
+class TestSubBlockSparseAttentionDispatch(CustomTestCase):
+    def setUp(self):
+        _get_subblock_sparse_attention_runner.cache_clear()
+        self.addCleanup(_get_subblock_sparse_attention_runner.cache_clear)
+
+    def test_dispatch_is_resolved_once_per_device(self):
+        device = torch.device("cuda:0")
+        with patch(
+            "torch.cuda.get_device_capability", return_value=(9, 0)
+        ) as get_capability:
+            first = _get_subblock_sparse_attention_runner(device)
+            second = _get_subblock_sparse_attention_runner(device)
+
+        self.assertIs(first, _sm90_sparse_attention)
+        self.assertIs(second, first)
+        get_capability.assert_called_once_with(device)
+
+    def test_dispatches_sm100(self):
+        device = torch.device("cuda:0")
+        with patch("torch.cuda.get_device_capability", return_value=(10, 0)):
+            runner = _get_subblock_sparse_attention_runner(device)
+
+        self.assertIs(runner, _sm100_sparse_attention)
+
+    def test_rejects_unsupported_compute_capability(self):
+        device = torch.device("cuda:0")
+        with patch("torch.cuda.get_device_capability", return_value=(10, 3)):
+            with self.assertRaisesRegex(
+                RuntimeError,
+                "supports compute capability 9.0 or 10.0;.*10.3 device",
+            ):
+                _get_subblock_sparse_attention_runner(device)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)

test/registered/kernels/ops/attention/test_subblock_sparse_sm90.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SM90-specific invariants for SubBlock sparse attention."""
+
+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=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+requires_sm90 = unittest.skipUnless(
+    torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0),
+    "requires SM90 (Hopper)",
+)
+
+
+@requires_sm90
+class TestSubBlockSparseSM90(CustomTestCase):
+    def test_64x64_routing_mask_uses_matching_compute_tile(self):
+        """A tile spanning routing rows would apply one row's mask to another row."""
+        from sglang.kernels.ops.attention.flash_attn.cute.interface import (
+            _tile_size_fwd_sm90,
+        )
+
+        config = _tile_size_fwd_sm90(
+            head_dim=128,
+            head_dim_v=128,
+            is_causal=False,
+            is_local=False,
+            sparse_block_size_q=64,
+            sparse_block_size_kv=64,
+        )
+
+        self.assertEqual(config.m_block_size, 64)
+        self.assertEqual(config.n_block_size, 64)
+
+    def test_64x64_special_case_is_limited_to_head_dim_128(self):
+        from sglang.kernels.ops.attention.flash_attn.cute.interface import (
+            _tile_size_fwd_sm90,
+        )
+
+        config = _tile_size_fwd_sm90(
+            head_dim=96,
+            head_dim_v=96,
+            is_causal=False,
+            is_local=False,
+            sparse_block_size_q=64,
+            sparse_block_size_kv=64,
+        )
+
+        self.assertEqual(config.m_block_size, 128)
+        self.assertEqual(config.n_block_size, 128)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)