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

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(

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
+```

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
                 ),

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

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()

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 (