From ee1f2e8dfdbb264727134aeac4b8dec155ec6efe Mon Sep 17 00:00:00 2001 From: HuangJi <32611516+IPostYellow@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:31:45 +0800 Subject: [PATCH] [diffusion][Minimax H3]support subblock sparse attention on SM90 (#34680) --- .../flash_attn/cute/block_sparse_utils.py | 41 ++++- .../attention/flash_attn/cute/interface.py | 37 ++++- .../backends/subblock_sparse/README.md | 20 +-- .../backends/subblock_sparse/__init__.py | 7 +- .../backends/subblock_sparse_attn.py | 147 ++++++++++++++++-- .../multimodal_gen/runtime/platforms/cuda.py | 50 +++--- .../unit/test_subblock_sparse_attention.py | 120 +++++++++++--- .../cpu/test_subblock_sparse_attention.py | 54 +++++++ .../attention/test_subblock_sparse_sm90.py | 58 +++++++ 9 files changed, 459 insertions(+), 75 deletions(-) create mode 100644 test/registered/cpu/test_subblock_sparse_attention.py create mode 100644 test/registered/kernels/ops/attention/test_subblock_sparse_sm90.py diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py b/python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py index 5072f9d60..55e7c1e2c 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py +++ b/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( diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py b/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py index aefa4861f..437bca8ca 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py +++ b/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( 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 3a5dcc4a4..dfe6e551a 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 @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/__init__.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/__init__.py index 312eddeff..d1c757595 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/__init__.py +++ b/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 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 f79328fa3..b1c10fd5c 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 @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 1c4bc869f..2dc303448 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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): 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 ffa1feb9b..6eba04fc6 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 @@ -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): diff --git a/test/registered/cpu/test_subblock_sparse_attention.py b/test/registered/cpu/test_subblock_sparse_attention.py new file mode 100644 index 000000000..d86de3831 --- /dev/null +++ b/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) diff --git a/test/registered/kernels/ops/attention/test_subblock_sparse_sm90.py b/test/registered/kernels/ops/attention/test_subblock_sparse_sm90.py new file mode 100644 index 000000000..e8fe1242e --- /dev/null +++ b/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)