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

python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py

@@ -294,8 +294,37 @@ def produce_block_sparse_loads(
     )
 
     mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits)
-    full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits)
     mask_empty = mask_begin == mask_end
+
+    # Normalization guarantees that the full count and index are both present
+    # or both absent. ``mask_empty`` is a runtime value, so CuTe still traces
+    # both sides of the dynamic branches below; without this specialization,
+    # the full-list side would subscript ``None`` and fail to compile.
+    if const_expr(blocksparse_tensors.full_block_cnt is None):
+        kv_producer_state = load_block_list(
+            curr_mask_block_idx,
+            mask_begin,
+            mask_end,
+            first_block_preloaded=False,
+            kv_producer_state=kv_producer_state,
+            load_K=load_K,
+            load_V=load_V,
+            pipeline_k=pipeline_k,
+            pipeline_v=pipeline_v,
+            intra_wg_overlap=intra_wg_overlap,
+        )
+        if const_expr(intra_wg_overlap) and not mask_empty:
+            kv_producer_state = finish_overlap_v_load(
+                curr_mask_block_idx,
+                mask_begin,
+                mask_end,
+                load_V,
+                pipeline_v,
+                kv_producer_state,
+            )
+        return kv_producer_state
+
+    full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits)
     full_empty = full_begin == full_end
 
     if mask_empty:
@@ -488,7 +517,10 @@ def consume_block_sparse_loads(
             if split_full_block_cnt == 0:
                 warp_scheduler_barrier_arrive()
 
-        if split_full_block_cnt > 0:
+        if (
+            const_expr(blocksparse_tensors.full_block_cnt is not None)
+            and split_full_block_cnt > 0
+        ):
             full_n_block = curr_full_block_idx[full_end - 1]
             if split_mask_block_cnt == 0:
                 warp_scheduler_barrier_sync()
@@ -561,7 +593,10 @@ def consume_block_sparse_loads(
                 )
                 O_should_accumulate = True
 
-        if split_full_block_cnt > 0:
+        if (
+            const_expr(blocksparse_tensors.full_block_cnt is not None)
+            and split_full_block_cnt > 0
+        ):
             full_n_block = curr_full_block_idx[full_end - 1]
             if split_mask_block_cnt == 0:
                 kv_consumer_state = process_first_half_block(

python/sglang/kernels/ops/attention/flash_attn/cute/interface.py

@@ -151,16 +151,34 @@ class FwdConfig:
 
 
 def _tile_size_fwd_sm90(
-    head_dim, head_dim_v, is_causal, is_local, sparse_block_size_q=None
+    head_dim,
+    head_dim_v,
+    is_causal,
+    is_local,
+    sparse_block_size_q=None,
+    sparse_block_size_kv=None,
 ):
     """Return FwdConfig for SM90 forward.
 
     Tile sizes and flags based on tile_size_fwd_sm90 in hopper/tile_size.h, adjusted
     for the Python kernel's different register/smem tradeoffs (benchmarked on H100 SXM).
 
-    When sparse_block_size_q is set, tile_m must divide it. For head_dim <= 96 the
-    optimal tile_m=192 is used when compatible, otherwise we fall back to 128.
+    When sparse block sizes are set, the compute tiles must respect both axes of
+    the sparse mask. The 64x64 case is used by SubBlock attention: every 64-row
+    query block has its own independently routed list of 64-row KV blocks, so it
+    cannot be coarsened to the usual 128x128 tile without changing the mask.
+
+    For other sparse masks, tile_m must divide sparse_block_size_q. For
+    head_dim <= 96 the optimal tile_m=192 is used when compatible, otherwise we
+    fall back to 128.
     """
+    if (
+        head_dim == 128
+        and sparse_block_size_q == 64
+        and sparse_block_size_kv == 64
+    ):
+        return FwdConfig(64, 64, True, True)
+
     if head_dim <= 64:
         # C++: 192×192 non-causal, 192×128 causal/local.
         # Python: 192×128 RS+OL is consistently best across seqlens.
@@ -718,8 +736,19 @@ def _flash_attn_fwd(
             fwd_cfg = FwdConfig(128, 64, True, True)  # SM80, should tune
         elif arch // 10 == 9:
             sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q)
+            sparse_block_size_kv = (
+                block_sparse_tensors.block_size[1]
+                if block_sparse_tensors is not None
+                and block_sparse_tensors.block_size is not None
+                else None
+            )
             fwd_cfg = _tile_size_fwd_sm90(
-                head_dim, head_dim_v, causal, local, sparse_block_size_q=sparse_q
+                head_dim,
+                head_dim_v,
+                causal,
+                local,
+                sparse_block_size_q=sparse_q,
+                sparse_block_size_kv=sparse_block_size_kv,
             )
     else:
         fwd_cfg = FwdConfig(

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md

@@ -1,8 +1,9 @@
 # SubBlock sparse attention — training-free block sparsity for the MiniMax-H3 DiT
 
-Routes FlashInfer's 64-token block-sparse kernel (`bsa_attn_blk64_fwd`) with a
-sub-block score. Nothing is trained and no weights change: a cheap estimator
-runs before attention and hands the kernel a `q2k_block_index`.
+Routes the same 64-token SubBlock plan to SGLang's CuTe-DSL block-sparse
+FlashAttention kernel on SM90 or FlashInfer's `bsa_attn_blk64_fwd` on SM100.
+Nothing is trained and no weights change: a cheap estimator runs before
+attention and hands the selected kernel a `q2k_block_index`.
 
 Spelled out in full, with every key at its default — which is the recommended
 configuration and what every number below was measured at:
@@ -32,11 +33,12 @@ quotes.
 
 ## What it runs on
 
-Everything below comes from `bsa_attn_blk64_fwd`, not from this backend.
+The backend selects an architecture-specific kernel; their shared constraints
+are listed below.
 
 | | |
 | --- | --- |
-| GPU | **compute capability 10.0 only** — B200 / GB200 class. The kernel is built `-gencode=arch=compute_100a,code=sm_100a`, which is arch-specific and does not forward-run on 10.3 (B300 / GB300) or 12.x (RTX PRO 6000, RTX 50xx). |
+| GPU | **compute capability 9.0 or 10.0** — H100 / H200 use SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` kernel. Other capabilities, including 10.3 (B300 / GB300) and 12.x (RTX PRO 6000, RTX 50xx), are rejected. |
 | dtype | bfloat16 |
 | head_dim | 128 |
 | attention | non-causal, one contiguous sequence per call |
@@ -46,10 +48,10 @@ refiner, sequences under `min_seq_len`, non-bf16 activations, head_dim != 128
 falls back to dense for that call, so no layer has to be excluded by hand.
 
 **On an unsupported GPU it is not a fallback, it is an error at startup.** The
-resolver checks the compute capability before anything loads and refuses
-anything but 10.0, so an H100 or a B300 fails at launch rather than after ten
-dense denoise steps. Do not rely on the kernel's own guard for this: it compares
-only the major version, so it would accept 10.3 and then fail with no cubin.
+resolver accepts exactly compute capability 9.0 or 10.0 before loading either
+kernel, so a B300 or an SM12x GPU fails at launch rather than after ten dense
+denoise steps. The exact 10.0 check is required because FlashInfer's kernel is
+built for `sm_100a` and has no forward-compatible 10.3 cubin.
 
 ## How the score works
 

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/__init__.py

@@ -5,9 +5,10 @@ Originally vendored from the standalone SubBlock repository; ``router.py`` and
 ``kernels.py`` have since diverged from it.
 
 ``router.py`` scores every (query block, key block) pair from sub-block-pooled
-Q/K and turns the scores into the ``q2k_block_index`` that FlashInfer's
-``bsa_attn_blk64_fwd`` consumes (SM100, bf16, head_dim 128). The estimator and
-the measurements behind its defaults are documented there.
+Q/K and turns the scores into a ``q2k_block_index`` consumed by SGLang's SM90
+CuTe-DSL block-sparse FlashAttention or FlashInfer's SM100
+``bsa_attn_blk64_fwd`` (bf16, head_dim 128). The estimator and the measurements
+behind its defaults are documented there.
 """
 
 from .router import SubBlockRouter, load_bsa_attn_blk64_fwd

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 """SubBlock block-sparse attention backend.
 
-Routes FlashInfer's 64-token block-sparse kernel with a K-side sub-block
-log-sum-exp score (see ``backends/subblock_sparse/``). Everything is training-free:
-the router runs before attention and produces the ``q2k_block_index`` the
-kernel consumes.
+Routes the same 64-token SubBlock plan to SGLang's CuTe-DSL block-sparse
+FlashAttention kernel on SM90 or FlashInfer's kernel on SM100. A log-sum-exp
+over query/key sub-block pairs selects the blocks (see ``backends/subblock_sparse/``).
+Everything is training-free: the router runs before attention and produces
+the ``q2k_block_index`` the selected kernel consumes.
 
 Sparsity is not applied everywhere. The early denoise steps settle the layout
 of the sample and tolerate approximation badly, so the backend falls back to
@@ -16,11 +17,12 @@ individual keys of the defaults below::
     --attention-backend subblock_sparse_attn \
     --attention-backend-config '{"sparsity": 0.85}'
 
-Requirements inherited from the kernel: compute capability 10.0 (B200 / GB200
-class -- it is built for ``sm_100a``, which does not forward-run on 10.3 or
-12.x), bf16, head_dim 128. Inside the DiT, any call the kernel cannot serve --
-cross/refiner attention, short sequences, non-bf16 -- runs dense instead. On any
-other GPU the resolver refuses the backend at startup rather than falling back.
+Requirements inherited from the kernels: compute capability 9.0 (Hopper) or
+10.0 (B200 / GB200), bf16, head_dim 128. Hopper uses SGLang's CuTe-DSL SM90
+block-sparse FlashAttention kernel; B200 uses FlashInfer's ``sm_100a`` blk64
+kernel. Inside the DiT, any call the kernels cannot serve -- cross/refiner
+attention, short sequences, non-bf16 -- runs dense instead. On any other GPU
+the resolver refuses the backend at startup rather than falling back.
 
 ``--attention-backend`` reaches every component, and the text encoder admits
 only fa / torch_sdpa / sage_attn_3, so pair it with
@@ -54,6 +56,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 # The kernel is fixed at 64-token blocks and 128-wide heads.
+SUBBLOCK_SPARSE_BLOCK_SIZE = 64
 SUBBLOCK_SPARSE_HEAD_DIM = 128
 
 # Defaults for the schedule; override through --attention-backend-config.
@@ -112,8 +115,124 @@ def _cached_block_sizes(seq_len: int, device: torch.device) -> torch.Tensor:
     return SubBlockRouter.block_sizes(seq_len, device)
 
 
-class SubBlockSparseAttentionBackend(AttentionBackend):
+@functools.lru_cache(maxsize=1)
+def _load_sm90_block_sparse_attention():
+    """Load the CuTe-DSL Hopper path only when an SM90 device selects it.
 
+    Keeping these imports lazy avoids pulling the sizeable CuTe dependency tree
+    into the existing SM100 path, whose FlashInfer blk64 kernel is plain CUDA.
+    """
+    from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (
+        BlockSparseTensorsTorch,
+    )
+    from sglang.kernels.ops.attention.flash_attn.cute.interface import flash_attn_func
+
+    return BlockSparseTensorsTorch, flash_attn_func
+
+
+def _sm90_sparse_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_block_index: torch.Tensor,
+    topk: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Run a SubBlock routing plan through the existing SM90 CuTe kernel."""
+    BlockSparseTensorsTorch, flash_attn_func = _load_sm90_block_sparse_attention()
+
+    # The router contract permits indices in any order, while the SM90 sparse
+    # pipeline consumes each list from high slot to low slot and applies
+    # sequence-tail masking to the first block. Sort explicitly so the largest
+    # block id -- the possible ragged tail -- occupies the highest slot without
+    # depending on the fused top-k kernel's current ascending output order.
+    ordered_index = q2k_block_index.sort(dim=-1).values
+    block_counts = torch.full(
+        ordered_index.shape[:-1],
+        topk,
+        dtype=torch.int32,
+        device=ordered_index.device,
+    )
+    sparse_tensors = BlockSparseTensorsTorch(
+        mask_block_cnt=block_counts,
+        mask_block_idx=ordered_index,
+        # There are no always-dense blocks in a SubBlock routing plan. The
+        # block-sparse broadcast pattern records both absent tensors as None
+        # and participates in the compile key, so mask-only and mask+full calls
+        # cannot share a compiled kernel.
+        full_block_cnt=None,
+        full_block_idx=None,
+        block_size=(SUBBLOCK_SPARSE_BLOCK_SIZE, SUBBLOCK_SPARSE_BLOCK_SIZE),
+    )
+    out, _ = flash_attn_func(
+        q,
+        k,
+        v,
+        softmax_scale=softmax_scale,
+        causal=False,
+        num_splits=1,
+        block_sparse_tensors=sparse_tensors,
+    )
+    return out
+
+
+def _sm100_sparse_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_block_index: torch.Tensor,
+    topk: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Run a SubBlock routing plan through FlashInfer's SM100 kernel."""
+    out = load_bsa_attn_blk64_fwd()(
+        q,
+        k,
+        v,
+        q2k_block_index,
+        topk,
+        block_sizes=_cached_block_sizes(k.shape[1], k.device),
+        q2k_block_nums=None,  # the budget is uniform across rows
+        softmax_scale=softmax_scale,
+    )
+    return out[0] if isinstance(out, tuple) else out
+
+
+@functools.lru_cache(maxsize=None)
+def _get_subblock_sparse_attention_runner(device: torch.device):
+    """Resolve the architecture-specific kernel once per CUDA device."""
+    capability = torch.cuda.get_device_capability(device)
+    if capability == (9, 0):
+        return _sm90_sparse_attention
+    if capability == (10, 0):
+        return _sm100_sparse_attention
+    raise RuntimeError(
+        "SubBlock sparse attention supports compute capability 9.0 or 10.0; "
+        f"this tensor is on a {capability[0]}.{capability[1]} device."
+    )
+
+
+def _run_subblock_sparse_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_block_index: torch.Tensor,
+    topk: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Dispatch the same 64x64 routing plan to Hopper or Blackwell."""
+    runner = _get_subblock_sparse_attention_runner(q.device)
+    return runner(
+        q,
+        k,
+        v,
+        q2k_block_index,
+        topk,
+        softmax_scale,
+    )
+
+
+class SubBlockSparseAttentionBackend(AttentionBackend):
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [SUBBLOCK_SPARSE_HEAD_DIM]
@@ -285,7 +404,6 @@ class SubBlockSparseAttentionImpl(AttentionImpl):
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
     ) -> torch.Tensor:
         """q, k, v: ``[1, S, H, 128]`` bf16 -> same shape."""
-        bsa_attn_blk64_fwd = load_bsa_attn_blk64_fwd()
         plan = self.router.route(
             q, k, sparsity=self.schedule.sparsity, softmax_scale=self.softmax_scale
         )
@@ -296,17 +414,14 @@ class SubBlockSparseAttentionImpl(AttentionImpl):
             f"keeping {plan.topk}/{plan.num_blocks} key blocks per query block "
             f"(sparsity {1 - plan.density:.4f})"
         )
-        out = bsa_attn_blk64_fwd(
+        return _run_subblock_sparse_attention(
             q,
             k,
             v,
             plan.index,
             plan.topk,
-            block_sizes=_cached_block_sizes(k.shape[1], k.device),
-            q2k_block_nums=None,  # the budget is uniform across rows
-            softmax_scale=self.softmax_scale,
+            self.softmax_scale,
         )
-        return out[0] if isinstance(out, tuple) else out
 
     def forward(
         self,

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -297,39 +297,53 @@ class _VMOBAAttentionBackendResolver(_CudaAttentionBackendResolver):
 class _SubBlockSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
     backend = AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
 
-    # The blk64 kernel is built `-gencode=arch=compute_100a,code=sm_100a`, which
-    # is arch-specific: 10.3 (B300 / GB300) and 12.x have no cubin. Its own guard
-    # only compares the major version, so it would accept 10.3 and fail later.
-    required_capability = (10, 0)
+    # Hopper uses SGLang's SM90 CuTe-DSL block-sparse kernel. Blackwell uses the
+    # FlashInfer blk64 kernel built specifically for sm_100a; 10.3 and 12.x do
+    # not have a compatible cubin and must still fail closed.
+    supported_capabilities = {(9, 0), (10, 0)}
 
     @classmethod
     def resolve(cls, platform) -> str:
         capability = platform.get_device_capability()
-        if capability is None or capability != cls.required_capability:
+        capability_tuple = (
+            (capability.major, capability.minor) if capability is not None else None
+        )
+        if capability_tuple not in cls.supported_capabilities:
             found = capability.as_version_str() if capability else "unknown"
             raise ValueError(
-                "SubBlock sparse attention needs compute capability "
-                f"{'.'.join(map(str, cls.required_capability))} (B200 / GB200); "
-                f"this device reports {found}."
+                "SubBlock sparse attention needs compute capability 9.0 "
+                f"(Hopper) or 10.0 (B200 / GB200); this device reports {found}."
             )
         try:
-            from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (  # noqa: F401
-                load_bsa_attn_blk64_fwd,
-            )
             from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (  # noqa: F401
                 SubBlockSparseAttentionBackend,
             )
 
-            # Importing the entry point catches a missing or broken FlashInfer;
-            # the CUDA extension itself is built lazily on the first call.
-            load_bsa_attn_blk64_fwd()
+            if capability_tuple == (9, 0):
+                # Importing catches missing/incompatible CuTe-DSL and Quack;
+                # the CUDA kernel itself is compiled lazily on the first call.
+                from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import (  # noqa: F401
+                    BlockSparseTensorsTorch,
+                )
+                from sglang.kernels.ops.attention.flash_attn.cute.interface import (  # noqa: F401
+                    flash_attn_func,
+                )
+            else:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (  # noqa: F401
+                    load_bsa_attn_blk64_fwd,
+                )
+
+                load_bsa_attn_blk64_fwd()
             return "sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn.SubBlockSparseAttentionBackend"
         except Exception as e:
             logger.error("Failed to import SubBlock sparse attention: %s", str(e))
-            raise ImportError(
-                "SubBlock sparse attention needs FlashInfer with the blk64 "
-                "block-sparse kernel (flashinfer.cute_dsl.sparse.bsa_attn_blk64_fwd)."
-            ) from e
+            dependency = (
+                "SGLang's SM90 CuTe-DSL FlashAttention dependencies"
+                if capability_tuple == (9, 0)
+                else "FlashInfer with the blk64 block-sparse kernel "
+                "(flashinfer.cute_dsl.sparse.bsa_attn_blk64_fwd)"
+            )
+            raise ImportError(f"SubBlock sparse attention needs {dependency}.") from e
 
 
 class _FlashAttention2BackendResolver(_CudaAttentionBackendResolver):

python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 """SubBlock block-sparse attention backend.
 
-The schedule tests are pure CPU. The numerical tests need an SM100 GPU with
+The schedule and adapter tests are pure CPU. The numerical tests need either
+an SM90 GPU with SGLang's CuTe-DSL dependencies or an SM100 GPU with
 FlashInfer's ``bsa_attn_blk64_fwd`` and are skipped otherwise.
 
 The trick that makes the sparse kernel checkable against dense attention: at
@@ -26,31 +27,40 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_att
     SubBlockSparseAttentionImpl,
     SubBlockSparseSchedule,
     _dit_layer_index,
+    _run_subblock_sparse_attention,
+    _sm90_sparse_attention,
 )
 
 HEAD_DIM = 128
 NUM_HEADS = 4
 
 
-def _sm100_available() -> bool:
+def _subblock_kernel_available() -> bool:
     if not torch.cuda.is_available():
         return False
-    # Exactly 10.0: the kernel is built for sm_100a, and 10.3 has no cubin.
-    if torch.cuda.get_device_capability(0) != (10, 0):
-        return False
+    capability = torch.cuda.get_device_capability(0)
     try:
-        from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
-            load_bsa_attn_blk64_fwd,
-        )
+        if capability == (9, 0):
+            from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
+                _load_sm90_block_sparse_attention,
+            )
 
-        load_bsa_attn_blk64_fwd()
+            _load_sm90_block_sparse_attention()
+        elif capability == (10, 0):
+            from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
+                load_bsa_attn_blk64_fwd,
+            )
+
+            load_bsa_attn_blk64_fwd()
+        else:
+            return False
     except Exception:
         return False
     return True
 
 
-requires_sm100 = unittest.skipUnless(
-    _sm100_available(), "needs SM100 and FlashInfer bsa_attn_blk64_fwd"
+requires_subblock_kernel = unittest.skipUnless(
+    _subblock_kernel_available(), "needs an SM90 or SM100 SubBlock attention kernel"
 )
 
 
@@ -180,6 +190,40 @@ class TestSubBlockSparseBackend(unittest.TestCase):
         )
         self.assertEqual(metadata.current_timestep, 7)
 
+    def test_sm90_adapter_sorts_indices_and_uses_64x64_blocks(self):
+        captured = {}
+
+        class _FakeBlockSparseTensors:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+                self.__dict__.update(kwargs)
+
+        def fake_flash_attn_func(q, k, v, **kwargs):
+            captured.update(kwargs)
+            return q, None
+
+        index = torch.tensor([[[[5, 1, 7, 3]]]], dtype=torch.int32)
+        q = torch.empty(1, 64, 1, HEAD_DIM, dtype=torch.bfloat16)
+        with patch(
+            "sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn._load_sm90_block_sparse_attention",
+            return_value=(_FakeBlockSparseTensors, fake_flash_attn_func),
+        ):
+            out = _sm90_sparse_attention(q, q, q, index, 4, HEAD_DIM**-0.5)
+
+        self.assertIs(out, q)
+        torch.testing.assert_close(
+            captured["mask_block_idx"],
+            torch.tensor([[[[1, 3, 5, 7]]]], dtype=torch.int32),
+        )
+        self.assertEqual(captured["mask_block_cnt"].item(), 4)
+        self.assertIsNone(captured["full_block_cnt"])
+        self.assertIsNone(captured["full_block_idx"])
+        self.assertEqual(captured["block_size"], (64, 64))
+        self.assertIs(
+            captured["block_sparse_tensors"].mask_block_idx,
+            captured["mask_block_idx"],
+        )
+
 
 class TestSubBlockGating(unittest.TestCase):
     """The schedule must decide sparsity from the layer and the step alone."""
@@ -241,7 +285,7 @@ class TestSubBlockGating(unittest.TestCase):
             self.assertFalse(impl._sparse_ready(q, q))
 
 
-@requires_sm100
+@requires_subblock_kernel
 class TestSubBlockNumerics(unittest.TestCase):
     seq_len = 8192
 
@@ -275,6 +319,46 @@ class TestSubBlockNumerics(unittest.TestCase):
         ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
         self.assertGreater(_cosine(out, ref), 0.999)
 
+    def test_unsorted_ragged_tail_oversubscribes_sms(self):
+        """Make tail-mask ordering observable across multiple SM waves."""
+        if torch.cuda.get_device_capability() != (9, 0):
+            self.skipTest("the reverse-consumption constraint is specific to SM90")
+
+        device = torch.device("cuda")
+        seq_len = self.seq_len + 37
+        shape = (1, seq_len, NUM_HEADS, HEAD_DIM)
+        q = torch.zeros(shape, device=device, dtype=torch.bfloat16)
+        k = torch.zeros_like(q)
+        v = torch.ones_like(q)
+
+        num_blocks = (seq_len + 63) // 64
+        num_tiles = NUM_HEADS * num_blocks
+        num_sms = torch.cuda.get_device_properties(device).multi_processor_count
+        self.assertGreater(num_tiles, 2 * num_sms)
+
+        # The SM90 consumer visits slots from high to low and applies the tail
+        # mask to the first block. Put the ragged block in the lowest slot, so
+        # removing the adapter's sort leaves its 27 padded rows unmasked. With
+        # zero Q/K and unit V that changes the output magnitude from 1 to
+        # (7 * 64 + 37) / (8 * 64), which an assert_close cannot overlook.
+        topk = 8
+        tail_block = num_blocks - 1
+        unsorted_blocks = torch.tensor(
+            [tail_block, 0, 1, 2, 3, 4, 5, 6],
+            device=device,
+            dtype=torch.int32,
+        )
+        unsorted_index = (
+            unsorted_blocks.view(1, 1, 1, topk)
+            .expand(1, NUM_HEADS, num_blocks, topk)
+            .clone()
+        )
+        out = _run_subblock_sparse_attention(
+            q, k, v, unsorted_index, topk, HEAD_DIM**-0.5
+        )
+
+        torch.testing.assert_close(out, torch.ones_like(out), rtol=0, atol=2e-3)
+
     def test_routing_finds_the_blocks_that_carry_the_mass(self):
         """At 0.75 sparsity the router must keep the blocks that matter.
 
@@ -282,11 +366,6 @@ class TestSubBlockNumerics(unittest.TestCase):
         number of blocks but chosen at random the output collapses, so a high
         cosine here measures the routing, not a forgiving fixture.
         """
-        from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
-            SubBlockRouter,
-            load_bsa_attn_blk64_fwd,
-        )
-
         device = torch.device("cuda")
         q, k, v = _structured_qkv(self.seq_len, device)
         ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
@@ -308,17 +387,14 @@ class TestSubBlockNumerics(unittest.TestCase):
             .argsort(dim=-1)[..., :topk]
             .to(torch.int32)
         )
-        random_out = load_bsa_attn_blk64_fwd()(
+        random_out = _run_subblock_sparse_attention(
             q,
             k,
             v,
             random_index,
             topk,
-            block_sizes=SubBlockRouter.block_sizes(self.seq_len, device),
-            q2k_block_nums=None,
-            softmax_scale=HEAD_DIM**-0.5,
+            HEAD_DIM**-0.5,
         )
-        random_out = random_out[0] if isinstance(random_out, tuple) else random_out
         self.assertLess(_cosine(random_out, ref), 0.9)
 
     def test_skipped_step_is_bitwise_dense(self):

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)