diff --git a/python/sglang/kernels/ops/attention/subblock_sage_fp8_sm90.py b/python/sglang/kernels/ops/attention/subblock_sage_fp8_sm90.py new file mode 100644 index 000000000..e4087163c --- /dev/null +++ b/python/sglang/kernels/ops/attention/subblock_sage_fp8_sm90.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SM90 adapter from a SubBlock plan to SpargeAttention's native Sage kernel. + +The public backend mode is ``sage_fp8`` across GPU generations. On Hopper, +the fastest available implementation is SageAttention2: it uses 64-token query +blocks and 128-token key blocks, quantizes Q/K to INT8, and quantizes V plus +softmax probabilities to E4M3 online. BF16 model activations and weights remain +unchanged. +""" + +from __future__ import annotations + +import functools + +import torch +import triton +import triton.language as tl + +SAGE_FP8_SM90_QUERY_BLOCK_SIZE = 64 +SAGE_FP8_SM90_KEY_BLOCK_SIZE = 128 + +_INSTALL_HELP = ( + "Install SpargeAttention with " + "`pip install git+https://github.com/thu-ml/SpargeAttn.git " + "--no-build-isolation` to use SubBlock compute_mode='sage_fp8' on SM90." +) + + +@functools.lru_cache(maxsize=1) +def _load_sparge_attention_sm90_ops(): + """Load every private SpargeAttention symbol used by the SM90 adapter.""" + try: + import spas_sage_attn._fused as fused + import spas_sage_attn._qattn as qattn + from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant + + transpose_pad_permute_cuda = fused.transpose_pad_permute_cuda + scale_fuse_quant_cuda = fused.scale_fuse_quant_cuda + kernel = ( + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90 + ) + except (ImportError, OSError, AttributeError) as exc: + raise ImportError(_INSTALL_HELP) from exc + return ( + get_vanilla_qk_quant, + block_map_lut_triton, + transpose_pad_permute_cuda, + scale_fuse_quant_cuda, + kernel, + ) + + +@triton.jit +def _routing_plan_to_block_map_kernel( + block_index, + block_counts, + block_map, + width: tl.constexpr, + num_key_blocks: tl.constexpr, + block: tl.constexpr, +): + row = tl.program_id(0) + slots = tl.arange(0, block) + count = tl.load(block_counts + row) + active = slots < count + key_block = tl.load( + block_index + row * width + slots, + mask=active, + other=0, + ) + tl.store( + block_map + row * num_key_blocks + key_block, + 1, + mask=active, + ) + + +def _routing_plan_to_block_map( + block_index: torch.Tensor, + block_counts: torch.Tensor, + num_key_blocks: int, +) -> torch.Tensor: + """Convert compact absolute block ids to SpargeAttention's dense bool map.""" + if block_index.ndim != 4: + raise ValueError( + "SubBlock Sage FP8 block_index must have shape [B, H, Gq, K], got " + f"{tuple(block_index.shape)}" + ) + if block_counts.shape != block_index.shape[:-1]: + raise ValueError( + "SubBlock Sage FP8 block_counts must match block_index[:-1], got " + f"{tuple(block_counts.shape)} and {tuple(block_index.shape)}" + ) + + width = block_index.shape[-1] + if block_index.device.type == "cuda": + flat_index = block_index.contiguous().view(-1, width) + flat_counts = block_counts.contiguous().view(-1) + block_map = torch.zeros( + (flat_index.shape[0], num_key_blocks), + dtype=torch.bool, + device=block_index.device, + ) + _routing_plan_to_block_map_kernel[(flat_index.shape[0],)]( + flat_index, + flat_counts, + block_map, + width=width, + num_key_blocks=num_key_blocks, + block=triton.next_power_of_2(width), + num_warps=4 if width >= 128 else 2, + ) + return block_map.view(*block_index.shape[:-1], num_key_blocks) + + flat_index = block_index.reshape(-1, width).long() + flat_counts = block_counts.reshape(-1) + slots = torch.arange(width, device=block_index.device) + active = slots[None, :] < flat_counts[:, None] + + # Only active entries are written. This matters for heterogeneous plans: + # ignored suffix values can duplicate an active id, and scatter(False) could + # otherwise race with scatter(True) for the same destination. + row = torch.arange(flat_index.shape[0], device=block_index.device) + row = row[:, None].expand_as(flat_index) + block_map = torch.zeros( + (flat_index.shape[0], num_key_blocks), + dtype=torch.bool, + device=block_index.device, + ) + block_map[row[active], flat_index[active]] = True + return block_map.view(*block_index.shape[:-1], num_key_blocks) + + +@torch.no_grad() +def subblock_sage_fp8_sm90_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_index: torch.Tensor, + topk: int, + softmax_scale: float, + block_counts: torch.Tensor | None = None, +) -> torch.Tensor: + """Run a 64x128 SubBlock plan with the native SM90 SageAttention2 kernel. + + Q/K/V use SGLang's normal ``[batch, sequence, heads, head_dim]`` layout. + Quantization is performed online and the returned tensor has the same BF16 + dtype and layout as Q. + """ + ( + get_vanilla_qk_quant, + block_map_lut_triton, + transpose_pad_permute_cuda, + scale_fuse_quant_cuda, + kernel, + ) = _load_sparge_attention_sm90_ops() + + if q.device.type != "cuda": + raise ValueError("SubBlock Sage FP8 requires CUDA tensors.") + if torch.cuda.get_device_capability(q.device) != (9, 0): + raise ValueError("This SubBlock Sage FP8 implementation requires SM90.") + if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("SubBlock Sage FP8 requires BF16 Q, K, and V.") + if k.device != q.device or v.device != q.device: + raise ValueError("SubBlock Sage FP8 requires Q, K, and V on one device.") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("SubBlock Sage FP8 expects Q, K, and V in [B, S, H, D].") + if ( + q.shape[0] != k.shape[0] + or q.shape[0] != v.shape[0] + or q.shape[2:] != k.shape[2:] + or q.shape[2:] != v.shape[2:] + or k.shape[1] != v.shape[1] + or q.shape[-1] != 128 + ): + raise ValueError( + "SubBlock Sage FP8 requires compatible Q, K, and V with head_dim=128." + ) + + expected_q_blocks = -(-q.shape[1] // SAGE_FP8_SM90_QUERY_BLOCK_SIZE) + num_key_blocks = -(-k.shape[1] // SAGE_FP8_SM90_KEY_BLOCK_SIZE) + if block_index.shape[:3] != (q.shape[0], q.shape[2], expected_q_blocks): + raise ValueError( + "SubBlock Sage FP8 routing shape does not match Q: expected " + f"{(q.shape[0], q.shape[2], expected_q_blocks)}, got " + f"{tuple(block_index.shape[:3])}" + ) + if block_counts is None: + block_counts = torch.full( + block_index.shape[:-1], + topk, + dtype=torch.int32, + device=block_index.device, + ) + block_map = _routing_plan_to_block_map(block_index, block_counts, num_key_blocks) + + # Keep the external package's online quantizers and native Hopper kernel, + # but not its second P/V-threshold sparsifier. SubBlock already chose the + # exact blocks to compute; another pruning rule would alter that plan and + # add a reduction to every iteration. + with torch.cuda.device(q.device): + q_hnd = q.transpose(1, 2).contiguous() + k_hnd = k.transpose(1, 2).contiguous() + v_hnd = v.transpose(1, 2).contiguous() + k_mean = k_hnd.mean(dim=-2, keepdim=True) + q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant( + q_hnd, + k_hnd, + k_mean, + SAGE_FP8_SM90_QUERY_BLOCK_SIZE, + SAGE_FP8_SM90_KEY_BLOCK_SIZE, + ) + lut, valid_block_counts = block_map_lut_triton(block_map) + + padded_kv_len = num_key_blocks * SAGE_FP8_SM90_KEY_BLOCK_SIZE + v_transposed = torch.empty( + (q.shape[0], q.shape[2], q.shape[3], padded_kv_len), + dtype=v.dtype, + device=v.device, + ) + transpose_pad_permute_cuda(v_hnd, v_transposed, 1) + v_fp8 = torch.empty_like(v_transposed, dtype=torch.float8_e4m3fn) + v_scale = torch.empty( + (q.shape[0], q.shape[2], q.shape[3]), + dtype=torch.float32, + device=v.device, + ) + scale_fuse_quant_cuda( + v_transposed, + v_fp8, + v_scale, + k.shape[1], + 2.25, + 1, + ) + + # The extension honors output strides. Let its HND kernel write through + # an HND view of BSHD storage, avoiding a full output transpose/copy. + output_bshd = torch.empty_like(q) + output_hnd = output_bshd.transpose(1, 2) + kernel( + q_int8, + k_int8, + v_fp8, + output_hnd, + lut, + valid_block_counts, + q_scale, + k_scale, + v_scale, + 1, # HND tensor layout + False, + 1, # per-block Q/K scales + softmax_scale, + ) + return output_bshd + + +__all__ = ["subblock_sage_fp8_sm90_attention"] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py index 20b8e829c..791bb0f4f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py @@ -260,6 +260,32 @@ class MiniMaxH3PipelineConfig(PipelineConfig): ) if selected_backend is None: return + if selected_backend is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN: + attention_config = server_args.attention_backend_config or {} + compute_mode = str(attention_config.get("compute_mode", "bf16")) + if compute_mode not in ("bf16", "sage_fp8"): + raise ValueError( + "SubBlock compute_mode must be 'bf16' or 'sage_fp8', got " + f"{compute_mode!r}." + ) + if compute_mode == "sage_fp8": + capability = current_platform.get_device_capability() + if capability is None or capability.to_int() != 90: + found = ( + capability.as_version_str() + if capability is not None + else "unknown" + ) + raise ValueError( + "MiniMax-H3 SubBlock compute_mode='sage_fp8' currently " + "requires SM90 (compute capability 9.0); " + f"found {found}." + ) + from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + _load_sparge_attention_sm90_ops, + ) + + _load_sparge_attention_sm90_ops() if selected_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3: if server_args.ring_degree > 1: raise ValueError( 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 31dbbb3c5..283fb8afb 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,11 @@ # SubBlock sparse attention — training-free block sparsity for the MiniMax-H3 DiT -Routes the same 64-token SubBlock plan to SGLang's CuTe-DSL block-sparse -FlashAttention kernel on SM90 or FlashInfer's architecture-specific blk64 -kernels on SM100 and SM120. +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. Nothing is trained and no weights change: a cheap estimator runs before attention and hands the selected kernel a `q2k_block_index`. @@ -16,7 +19,7 @@ sglang serve --model-path MiniMaxAI/MiniMax-H3 --model-variant fl2va \ --component-attention-backends text_encoder=fa \ --attention-backend-config '{"sparsity": 0.75, "n_k": 4, "n_q": 4, "skip_first_steps": 10, "skip_first_layers": 0, - "min_seq_len": 4096}' + "min_seq_len": 4096, "compute_mode": "bf16"}' ``` **The text-encoder override is not optional.** `--attention-backend` applies to @@ -43,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; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` kernel; SM120 devices use FlashInfer's `bsa_attn_sm120_blk64_fwd` CuTe-DSL 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 `bsa_attn_sm120_blk64_fwd` CuTe-DSL BF16 kernel. Other capabilities, including 10.3 (B300 / GB300), are rejected. | | dtype | bfloat16 | | head_dim | 128 | | attention | non-causal, one contiguous sequence per call | @@ -86,17 +89,33 @@ rejected. | key | default | meaning | | --- | ---: | --- | | `sparsity` | 0.75 | key blocks dropped per query block, as an upper bound | -| `n_k` | 4 | key sub-blocks per 64-token block (1, 2, 4, 8) | +| `n_k` | 4 (`sage_fp8` on SM90: 8) | key sub-blocks per key block (1, 2, 4, 8) | | `n_q` | 4 | query sub-blocks per 64-token block (1, 2, 4, 8) | | `skip_first_steps` | 10 | leading denoise forwards kept dense | | `skip_first_layers` | 0 | leading DiT blocks kept dense | | `min_seq_len` | 4096 | shorter sequences run dense | +| `compute_mode` | `"bf16"` | compute path: `"bf16"` or architecture-dispatched `"sage_fp8"` | -**`sparsity` is an upper bound, not an exact figure.** The kernel pads each query -row's block count up to a multiple of 8 with phantom slots it then masks out, so -148 blocks costs exactly what 152 costs; the router takes the 152. At 590 blocks, -0.75 requested delivers 0.7424, and the startup log reports what was kept. It is -the speed lever — see below — and the only knob most users should touch. +`compute_mode="sage_fp8"` is the stable public name, not a promise that every +architecture quantizes every operand identically. On SM90 it calls +SpargeAttention's native Hopper SageAttention2 kernel: Q/K are quantized online +to INT8 and V/P to E4M3. Its Q64 x K128 geometry defaults `n_k` to 8, preserving +the normal router's 16-token key pooling cells. Install the optional dependency: + +```bash +pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation +``` + +Enable it on SM90 with: + +```bash +--attention-backend-config '{"compute_mode":"sage_fp8"}' +``` + +**`sparsity` is an upper bound, not always an exact figure.** BF16 plans round +the budget up to groups of 8: 148 blocks costs what 152 costs. Native SM90 +`sage_fp8` uses 128-token key blocks without that rounding; at S=37,760 and +sparsity 0.75 it keeps 74 of 295 blocks. Startup logs report the actual budget. **`n_k` and `n_q` buy score accuracy, not speed.** They set how finely a block is cut before scoring: `n_k=4` means four 16-token key sub-blocks, and the block's @@ -168,10 +187,11 @@ are comparable. | --- | --- | | `router.py` | `SubBlockRouter` — pooling, scoring, selection, `RoutingPlan` | | `kernels.py` | Triton pooling / segmented-LSE / fused top-k | +| `../../../../../../kernels/ops/attention/subblock_sage_fp8_sm90.py` | production adapter to native SM90 SageAttention2 | | `../subblock_sparse_attn.py` | the `AttentionBackend`: schedule, gating, dense fallback | -Tests: `test/unit/test_subblock_sparse_attention.py`. The trick that makes the sparse -kernel checkable against dense is running it at a sparsity just above 0 — every -block is then inside the budget, so the result must reproduce dense attention up -to bf16 rounding, which pins the routing indices, the ragged tail block sizes -and the softmax scale in one assertion. +Tests: `test/unit/test_subblock_sparse_attention.py` and +`test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py`. The GPU +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. 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 f59c4b629..5a3838e3b 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 @@ -6,9 +6,9 @@ Originally vendored from the standalone SubBlock repository; ``router.py`` and ``router.py`` scores every (query block, key block) pair from sub-block-pooled Q/K and turns the scores into a ``q2k_block_index`` consumed by SGLang's SM90 -CuTe-DSL block-sparse FlashAttention or FlashInfer's architecture-specific -SM100/SM120 blk64 kernels (bf16, head_dim 128). The estimator and the -measurements behind its defaults are documented there. +CuTe-DSL block-sparse FlashAttention, the native SM90 Sage FP8 adapter, or +FlashInfer's architecture-specific SM100/SM120 blk64 kernels. The estimator and +the measurements behind its defaults are documented there. """ from .router import ( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py index b3d9c9ea6..cedc1e655 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py @@ -1,15 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -"""Sub-block block-sparse routing for FlashInfer's ``bsa_attn_blk64_fwd``. +"""Sub-block block-sparse routing for architecture-specific attention kernels. -Training-free. Runs *before* attention, produces the ``q2k_block_index`` tensor the -64-token block-sparse kernel consumes. +Training-free. Runs *before* attention and produces the ``q2k_block_index`` tensor +the selected block-sparse kernel consumes. The current BF16 kernels use Q64 x K64 +on SM90, SM100, and SM120; SM90 ``sage_fp8`` uses Q64 x K128. The router accepts +both block widths so each compute path receives indices in its native geometry. Why sub-blocks -------------- -The usual proxy score for a 64x64 block is ``mean(Q_block) . mean(K_block)``. Averaging -64 keys into one vector throws away exactly the variation that decides which keys a query -wants. Splitting each 64-token block into ``n`` sub-blocks of ``64/n`` tokens, scoring all -sub-block pairs and combining them with a log-sum-exp recovers most of that: +For the Q64 x K64 BF16 paths, the usual proxy score is +``mean(Q_block) . mean(K_block)``. Averaging a whole key block into one vector throws +away exactly the variation that decides which keys a query wants. Splitting the query +and key blocks into ``n_q`` and ``n_k`` sub-blocks, scoring all sub-block pairs, and +combining them with a log-sum-exp recovers most of that: score(i, j) = log sum_{a 64 / 32 / 16 / 8 tokens +BLOCK = 64 # default Q/K block size for the current BF16 consumers +BUDGET_GRANULARITY = 8 # Q64 x K64 default; SM90 sage_fp8 uses exact block counts +VALID_N = (1, 2, 4, 8) # sub-blocks per query/key block def _snap_up_to_8(topk: int, num_blocks: int) -> int: @@ -164,7 +170,14 @@ def _snap_up_to_8(topk: int, num_blocks: int) -> int: The consequence for the caller is that ``sparsity`` is an upper bound rather than an exact figure: 0.75 of 590 blocks becomes 152 kept, 0.7424 dropped. """ - return min(num_blocks, max(1, -(-topk // BUDGET_GRANULARITY)) * BUDGET_GRANULARITY) + return _snap_up(topk, num_blocks, BUDGET_GRANULARITY) + + +def _snap_up(topk: int, num_blocks: int, granularity: int) -> int: + """Clamp a requested budget and round it to a kernel's billing unit.""" + if granularity < 1: + raise ValueError(f"budget granularity must be positive, got {granularity}") + return min(num_blocks, max(1, -(-topk // granularity)) * granularity) class RoutingPlan(msgspec.Struct, frozen=True): @@ -183,11 +196,18 @@ class SubBlockRouter: """Builds ``q2k_block_index`` from sub-block-pooled Q/K. Args: - n_k: key sub-blocks per 64-token block (1, 2, 4 or 8). 1 reproduces plain avg - pooling; 4 is the quality/cost point the recall table above lands on. - n_q: query sub-blocks, same values. Splitting Q *alone* (n_q>1 with n_k=1) is - worse than not splitting; splitting both sides together is what the default - does. Costs n_q times the score matrix, 0.5% of the denoise time. + n_k: key sub-blocks per key block (1, 2, 4 or 8). 1 reproduces plain + average pooling. The BF16 Q64 x K64 paths default to 4; SM90 + ``sage_fp8`` defaults to 8 for its K128 blocks. + n_q: query sub-blocks per query block, with the same allowed values. + Splitting Q *alone* (n_q>1 with n_k=1) is worse than not splitting; + splitting both sides together is what the defaults do. Costs n_q + times the score matrix, 0.5% of denoise time. + block_size_k: native key-block width. It is 64 for the BF16 kernels and + 128 for SM90 ``sage_fp8``. + budget_granularity: rounding applied to the selected block count. The + Q64 x K64 paths retain the established default of 8 (required by + FlashInfer on SM100/SM120); SM90 ``sage_fp8`` uses 1. Structural block reservation (an attention sink, or forcing the diagonal j == i) was measured on 200 real H3 attention cells and is deliberately absent: at a fixed budget @@ -195,12 +215,30 @@ class SubBlockRouter: which did not survive to the pixels. """ - def __init__(self, n_k: int = 4, n_q: int = 4) -> None: + def __init__( + self, + n_k: int = 4, + n_q: int = 4, + *, + block_size_k: int = BLOCK, + budget_granularity: int = BUDGET_GRANULARITY, + ) -> None: if n_k not in VALID_N or n_q not in VALID_N: raise ValueError( f"n_q/n_k must be one of {VALID_N}, got n_q={n_q}, n_k={n_k}" ) + if block_size_k <= 0 or block_size_k % n_k: + raise ValueError( + f"key block size {block_size_k} must be positive and divisible " + f"by n_k={n_k}" + ) + if budget_granularity < 1: + raise ValueError( + f"budget granularity must be positive, got {budget_granularity}" + ) self.n_k, self.n_q = n_k, n_q + self.block_size_k = block_size_k + self.budget_granularity = budget_granularity @torch.no_grad() def scores( @@ -223,12 +261,14 @@ class SubBlockRouter: """ b, s, h, d = q.shape sk = k.shape[1] - gq, gk = -(-s // BLOCK), -(-sk // BLOCK) + gq = -(-s // BLOCK) + gk = -(-sk // self.block_size_k) nq, nk = self.n_q, self.n_k - sub_q, sub_k = BLOCK // nq, BLOCK // nk + sub_q = BLOCK // nq + sub_k = self.block_size_k // nk # Pooling handles the ragged tail on the *pooled* tensor: padding q/k up to - # G*BLOCK first would copy the whole 300+ MB activation to add a few rows. + # complete native blocks first would copy the whole 300+ MB activation. # Sub-cells past the last real token pool to zero, and `*_valid` tells the score # kernel to drop them -- left in, each would contribute an exp(0)=1 term that # both inflates the score and flattens the differences the ranking depends on. @@ -255,13 +295,14 @@ class SubBlockRouter: ) -> RoutingPlan: """Select the top ``(1 - sparsity)`` fraction of key blocks per query block.""" b, s, h, d = q.shape - gk = -(-k.shape[1] // BLOCK) + gk = -(-k.shape[1] // self.block_size_k) scores = self.scores(q, k, softmax_scale) # [B, H, Gq, Gk] gq = scores.shape[2] - topk = _snap_up_to_8(math.ceil((1.0 - sparsity) * gk), gk) + topk = _snap_up(math.ceil((1.0 - sparsity) * gk), gk, self.budget_granularity) # One pass over the score matrix instead of torch.topk's several. The - # output order is unspecified: SM100 consumes it directly, while the - # SM90 backend sorts compact active prefixes before heterogeneous expansion. + # output order is unspecified. SM100/SM120 BF16 consume it directly; + # SM90 sage_fp8 converts it to an order-independent block map. Only the + # SM90 CuTe BF16 consumer sorts compact active prefixes before expansion. index = fused_topk(scores.reshape(-1, gk), topk).view(b, h, gq, topk) return RoutingPlan(index=index, topk=topk, num_blocks=gk) 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 465e5179c..6433bb361 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,10 @@ # SPDX-License-Identifier: Apache-2.0 """SubBlock block-sparse attention backend. -Routes the same 64-token SubBlock plan to SGLang's CuTe-DSL block-sparse -FlashAttention kernel on SM90 or FlashInfer's architecture-specific kernels on -SM100 and SM120. A log-sum-exp over query/key sub-block pairs selects the blocks -(see ``backends/subblock_sparse/``). +Routes a SubBlock plan to SGLang's CuTe-DSL block-sparse FlashAttention kernel +on SM90 or FlashInfer's architecture-specific kernels on SM100 and SM120. 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. @@ -20,10 +20,13 @@ 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; B200 and SM120 devices use FlashInfer's -architecture-specific blk64 kernels. 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. +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. +Inside the DiT, unsupported calls run dense instead; unsupported GPU +architectures are rejected. ``--attention-backend`` reaches every component, and the text encoder admits only fa / torch_sdpa / sage_attn_3. Pair it with @@ -58,9 +61,10 @@ 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. +# Query blocks and heads are fixed; Hopper Sage FP8 widens K blocks to 128. SUBBLOCK_SPARSE_BLOCK_SIZE = 64 SUBBLOCK_SPARSE_HEAD_DIM = 128 +SAGE_FP8_SM90_KEY_BLOCK_SIZE = 128 # Defaults for the schedule; override through --attention-backend-config. # Sparsity is the speed lever, and it saturates: on MiniMax-H3 t2va at 37.7k @@ -97,6 +101,9 @@ DEFAULT_N_Q = 4 # Below this many keys the router costs more than the blocks it saves, and the # top-k budget collapses to a handful of blocks. DEFAULT_MIN_SEQ_LEN = 4096 +# Keep the established BF16 kernel as the default. ``sage_fp8`` is explicit +# until its model-level quality/performance envelope has been validated. +DEFAULT_COMPUTE_MODE = "bf16" # ``blocks..attn`` is a DiT layer; ``token_refiner.blocks..attn`` and # anything else is not and stays dense. @@ -178,6 +185,31 @@ def _sm90_sparse_attention( return out +def _sm90_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: + """Run a 64x128 plan with the native SM90 Sage INT8-QK/FP8-PV kernel.""" + from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + subblock_sage_fp8_sm90_attention, + ) + + return subblock_sage_fp8_sm90_attention( + q, + k, + v, + q2k_block_index, + topk, + softmax_scale, + block_counts, + ) + + def _sm100_sparse_attention( q: torch.Tensor, k: torch.Tensor, @@ -226,14 +258,31 @@ def _sm120_sparse_attention( @functools.lru_cache(maxsize=None) -def _get_subblock_sparse_attention_runner(device: torch.device): - """Resolve the architecture-specific kernel once per CUDA device.""" +def _get_subblock_sparse_attention_runner( + device: torch.device, compute_mode: str = DEFAULT_COMPUTE_MODE +): + """Resolve the architecture/mode-specific kernel once per CUDA device.""" capability = torch.cuda.get_device_capability(device) + if compute_mode not in ("bf16", "sage_fp8"): + raise ValueError(f"unknown SubBlock compute mode {compute_mode!r}") if capability == (9, 0): - return _sm90_sparse_attention + if compute_mode == "bf16": + return _sm90_sparse_attention + if compute_mode == "sage_fp8": + return _sm90_sage_fp8_sparse_attention 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." + ) 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." + ) return _sm120_sparse_attention raise RuntimeError( "SubBlock sparse attention supports compute capability 9.0, 10.0, or 12.0; " @@ -249,15 +298,16 @@ def _run_subblock_sparse_attention( topk: int, softmax_scale: float, block_counts: torch.Tensor | None = None, + compute_mode: str = DEFAULT_COMPUTE_MODE, ) -> torch.Tensor: - """Dispatch a prepared 64x64 routing plan to Hopper or Blackwell. + """Dispatch a prepared routing plan to Hopper or Blackwell. SM90 requires every active index prefix to be sorted in ascending order; SM100 and SM120 accept the router's original order. Heterogeneous callers must sort compact sparse prefixes before expanding them to full-width dense rows. """ - runner = _get_subblock_sparse_attention_runner(q.device) + runner = _get_subblock_sparse_attention_runner(q.device, compute_mode) return runner( q, k, @@ -320,12 +370,17 @@ class SubBlockSparseSchedule(msgspec.Struct, frozen=True): n_k: int n_q: int min_seq_len: int + compute_mode: str @classmethod def from_server_args(cls) -> SubBlockSparseSchedule: from sglang.multimodal_gen.runtime.server_args import get_global_server_args 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 schedule = SubBlockSparseSchedule( sparsity=float(config.get("sparsity", DEFAULT_SPARSITY)), skip_first_steps=int( @@ -334,9 +389,10 @@ class SubBlockSparseSchedule(msgspec.Struct, frozen=True): skip_first_layers=int( config.get("skip_first_layers", DEFAULT_SKIP_FIRST_LAYERS) ), - n_k=int(config.get("n_k", DEFAULT_N_K)), + n_k=int(config.get("n_k", default_n_k)), n_q=int(config.get("n_q", DEFAULT_N_Q)), min_seq_len=int(config.get("min_seq_len", DEFAULT_MIN_SEQ_LEN)), + compute_mode=compute_mode, ) if not 0.0 <= schedule.sparsity < 1.0: raise ValueError( @@ -347,6 +403,11 @@ class SubBlockSparseSchedule(msgspec.Struct, frozen=True): raise ValueError(f"subblock {name} must be 1, 2, 4 or 8, got {value}") if schedule.skip_first_steps < 0 or schedule.skip_first_layers < 0: raise ValueError("subblock skip_first_* must be non-negative") + if schedule.compute_mode not in ("bf16", "sage_fp8"): + raise ValueError( + "subblock compute_mode must be 'bf16' or 'sage_fp8', got " + f"{schedule.compute_mode!r}" + ) return schedule @@ -388,7 +449,18 @@ class SubBlockSparseAttentionImpl(AttentionImpl): and self.schedule.sparsity > 0.0 ) self.router = ( - SubBlockRouter(n_k=self.schedule.n_k, n_q=self.schedule.n_q) + SubBlockRouter( + n_k=self.schedule.n_k, + n_q=self.schedule.n_q, + block_size_k=( + SAGE_FP8_SM90_KEY_BLOCK_SIZE + if self.schedule.compute_mode == "sage_fp8" + else SUBBLOCK_SPARSE_BLOCK_SIZE + ), + budget_granularity=( + 1 if self.schedule.compute_mode == "sage_fp8" else 8 + ), + ) if self.layer_enabled else None ) @@ -396,7 +468,8 @@ class SubBlockSparseAttentionImpl(AttentionImpl): if self.layer_enabled: logger.info_once( f"SubBlock sparse attention: sparsity={self.schedule.sparsity:.3f} " - f"n_k={self.schedule.n_k} n_q={self.schedule.n_q}, dense for the first " + f"n_k={self.schedule.n_k} n_q={self.schedule.n_q} " + f"compute_mode={self.schedule.compute_mode}, dense for the first " f"{self.schedule.skip_first_steps} denoise steps and the first " f"{self.schedule.skip_first_layers} DiT layers" ) @@ -475,7 +548,9 @@ class SubBlockSparseAttentionImpl(AttentionImpl): "query blocks are dense" ) block_counts = None - runner = _get_subblock_sparse_attention_runner(q.device) + runner = _get_subblock_sparse_attention_runner( + q.device, self.schedule.compute_mode + ) block_index = ( plan.index.sort(dim=-1).values if runner is _sm90_sparse_attention @@ -522,6 +597,7 @@ class SubBlockSparseAttentionImpl(AttentionImpl): kernel_topk, self.softmax_scale, block_counts, + self.schedule.compute_mode, ) def forward( diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index 206a66434..def45664d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py @@ -581,6 +581,7 @@ def test_validate_server_args_accepts_transformer_backend_override(): server_args = SimpleNamespace( component_attention_backends={"transformer": "subblock_sparse_attn"}, attention_backend="fa", + attention_backend_config={}, ring_degree=1, resolve_component_attention_backend=lambda *_names: ( AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN, 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 9dddc90d4..4dd2b0311 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 @@ -21,6 +21,7 @@ from unittest.mock import Mock, patch import torch from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse.router import ( + SubBlockRouter, _snap_up_to_8, ) from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( @@ -163,9 +164,25 @@ class TestSubBlockSparseSchedule(unittest.TestCase): self.assertEqual(schedule.skip_first_layers, 0) self.assertEqual(schedule.n_k, 4) self.assertEqual(schedule.n_q, 4) + 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) + + def test_explicit_sage_fp8_n_k_is_respected(self): + with _patch_schedule({"compute_mode": "sage_fp8", "n_k": 4}): + schedule = SubBlockSparseSchedule.from_server_args() + self.assertEqual(schedule.n_k, 4) def test_rejects_out_of_range_values(self): - for config in ({"sparsity": 1.0}, {"n_k": 3}, {"skip_first_steps": -1}): + for config in ( + {"sparsity": 1.0}, + {"n_k": 3}, + {"skip_first_steps": -1}, + {"compute_mode": "fp8"}, + ): with self.subTest(config=config), _patch_schedule(config): with self.assertRaises(ValueError): SubBlockSparseSchedule.from_server_args() @@ -185,6 +202,23 @@ class TestBudgetGranularity(unittest.TestCase): self.assertEqual(_snap_up_to_8(3, 5), 5) +class TestRouterGeometry(unittest.TestCase): + def test_sm90_sage_fp8_preserves_16_token_pooling_cells(self): + router = SubBlockRouter( + n_q=4, + n_k=8, + block_size_k=128, + budget_granularity=1, + ) + self.assertEqual(64 // router.n_q, 16) + self.assertEqual(router.block_size_k // router.n_k, 16) + self.assertEqual(router.budget_granularity, 1) + + def test_rejects_non_divisible_block_geometry(self): + with self.assertRaisesRegex(ValueError, "divisible"): + SubBlockRouter(n_q=4, n_k=8, block_size_k=100) + + class TestSubBlockSparseBackend(unittest.TestCase): def test_sm90_adapter_uses_presorted_indices_and_64x64_blocks(self): captured = {} @@ -286,6 +320,12 @@ 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) + @requires_subblock_kernel class TestSubBlockNumerics(unittest.TestCase): diff --git a/test/registered/cpu/test_subblock_sparse_attention.py b/test/registered/cpu/test_subblock_sparse_attention.py index ecfcf926f..5a8daa8bc 100644 --- a/test/registered/cpu/test_subblock_sparse_attention.py +++ b/test/registered/cpu/test_subblock_sparse_attention.py @@ -1,17 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 +import sys import unittest -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import Mock, patch import torch +from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + _load_sparge_attention_sm90_ops, + _routing_plan_to_block_map, +) from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( MiniMaxH3PipelineConfig, ) from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( SubBlockSparseAttentionImpl, _get_subblock_sparse_attention_runner, + _sm90_sage_fp8_sparse_attention, _sm90_sparse_attention, _sm100_sparse_attention, _sm120_sparse_attention, @@ -40,12 +46,87 @@ from sglang.multimodal_gen.runtime.platforms import ( from sglang.multimodal_gen.runtime.platforms.cuda import ( _SubBlockSparseAttentionBackendResolver, ) +from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=9, suite="stage-a-test-cpu-intel") +class TestSubBlockSageFp8PlanAdapter(CustomTestCase): + def test_variable_counts_ignore_each_row_suffix(self): + index = torch.tensor([[[[2, 0, 1, 3], [3, 1, 0, 2]]]], dtype=torch.int32) + counts = torch.tensor([[[3, 1]]], dtype=torch.int32) + + block_map = _routing_plan_to_block_map(index, counts, 4) + + torch.testing.assert_close( + block_map, + torch.tensor([[[[True, True, True, False], [False, False, False, True]]]]), + ) + + +class TestSubBlockSageFp8DependencyLoader(CustomTestCase): + def setUp(self): + _load_sparge_attention_sm90_ops.cache_clear() + self.addCleanup(_load_sparge_attention_sm90_ops.cache_clear) + + def test_missing_dependency_has_install_help(self): + missing_modules = { + name: None + for name in ( + "spas_sage_attn", + "spas_sage_attn._fused", + "spas_sage_attn._qattn", + "spas_sage_attn.utils", + ) + } + with ( + patch.dict(sys.modules, missing_modules), + self.assertRaisesRegex(ImportError, "pip install.*SpargeAttn"), + ): + _load_sparge_attention_sm90_ops() + + def test_complete_dependency_exposes_required_ops(self): + package = ModuleType("spas_sage_attn") + package.__path__ = [] + fused = ModuleType("spas_sage_attn._fused") + qattn = ModuleType("spas_sage_attn._qattn") + utils = ModuleType("spas_sage_attn.utils") + fused.transpose_pad_permute_cuda = Mock() + fused.scale_fuse_quant_cuda = Mock() + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90 = ( + Mock() + ) + utils.block_map_lut_triton = Mock() + utils.get_vanilla_qk_quant = Mock() + package._fused = fused + package._qattn = qattn + package.utils = utils + + with patch.dict( + sys.modules, + { + "spas_sage_attn": package, + "spas_sage_attn._fused": fused, + "spas_sage_attn._qattn": qattn, + "spas_sage_attn.utils": utils, + }, + ): + ops = _load_sparge_attention_sm90_ops() + + self.assertEqual( + ops, + ( + utils.get_vanilla_qk_quant, + utils.block_map_lut_triton, + fused.transpose_pad_permute_cuda, + fused.scale_fuse_quant_cuda, + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90, + ), + ) + + class TestSubBlockSparseAttentionDispatch(CustomTestCase): def setUp(self): _get_subblock_sparse_attention_runner.cache_clear() @@ -70,6 +151,23 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase): self.assertIs(runner, _sm100_sparse_attention) + def test_dispatches_sm90_sage_fp8_independently_from_bf16(self): + device = torch.device("cuda:0") + with patch("torch.cuda.get_device_capability", return_value=(9, 0)): + bf16_runner = _get_subblock_sparse_attention_runner(device, "bf16") + sage_runner = _get_subblock_sparse_attention_runner(device, "sage_fp8") + + 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): + device = torch.device("cuda:0") + with ( + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + self.assertRaisesRegex(RuntimeError, "currently targets SM90"), + ): + _get_subblock_sparse_attention_runner(device, "sage_fp8") + def test_dispatches_sm120(self): device = torch.device("cuda:0") with patch("torch.cuda.get_device_capability", return_value=(12, 0)): @@ -134,6 +232,14 @@ 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)): @@ -145,6 +251,128 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase): class TestSubBlockSparseAttentionModalities(CustomTestCase): + @staticmethod + def _subblock_server_args(compute_mode: str): + return SimpleNamespace( + attention_backend="subblock_sparse_attn", + attention_backend_config={"compute_mode": compute_mode}, + ring_degree=1, + resolve_component_attention_backend=lambda *_names: ( + AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN, + "transformer", + ), + ) + + def test_sage_fp8_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(9, 0), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.kernels.ops.attention.subblock_sage_fp8_sm90." + "_load_sparge_attention_sm90_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_sage_fp8_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(9, 0), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.kernels.ops.attention.subblock_sage_fp8_sm90." + "_load_sparge_attention_sm90_ops", + side_effect=ImportError("Install SpargeAttention"), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_attn_backend" + ) as get_backend, + self.assertRaisesRegex(ImportError, "Install SpargeAttention"), + ): + 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") + loader = Mock() + with ( + patch.object(current_platform, "is_mps", return_value=False), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.kernels.ops.attention.subblock_sage_fp8_sm90." + "_load_sparge_attention_sm90_ops", + loader, + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_attn_backend" + ), + ): + config.validate_server_args(server_args) + + loader.assert_not_called() + + def test_sage_fp8_rejects_non_sm90_during_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(10, 0), + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_global_forced_attn_backend", + return_value=None, + ), + patch( + "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "get_attn_backend" + ) as get_backend, + self.assertRaisesRegex(ValueError, "requires SM90.*found 10.0"), + ): + config.validate_server_args(server_args) + + get_backend.assert_not_called() + def test_transformer_subblock_with_ring_fails_admission(self): config = MiniMaxH3PipelineConfig() server_args = SimpleNamespace( @@ -400,7 +628,7 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase): impl = object.__new__(SubBlockSparseAttentionImpl) impl.softmax_scale = 2**-0.5 impl.causal = False - impl.schedule = SimpleNamespace(sparsity=0.75) + impl.schedule = SimpleNamespace(sparsity=0.75, compute_mode="bf16") plan = SimpleNamespace( index=torch.tensor( [[[[7, 1, 4], [6, 2, 0], [5, 0, 3]]]], dtype=torch.int32 @@ -419,6 +647,7 @@ class TestSubBlockSparseAttentionModalities(CustomTestCase): for runner, sparse_rows in ( (_sm90_sparse_attention, ([1, 4, 7], [0, 3, 5])), + (_sm90_sage_fp8_sparse_attention, ([7, 1, 4], [5, 0, 3])), (_sm100_sparse_attention, ([7, 1, 4], [5, 0, 3])), (_sm120_sparse_attention, ([7, 1, 4], [5, 0, 3])), ): diff --git a/test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py b/test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py new file mode 100644 index 000000000..1f5828e4f --- /dev/null +++ b/test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Correctness tests for native SM90 SubBlock Sage FP8 attention.""" + +import math +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=45, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def _native_sm90_sage_available() -> bool: + try: + import spas_sage_attn._qattn as qattn + except (ImportError, OSError): + return False + return hasattr( + qattn, + "qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90", + ) + + +requires_native_sm90_sage = unittest.skipUnless( + torch.cuda.is_available() + and torch.cuda.get_device_capability() == (9, 0) + and _native_sm90_sage_available(), + "requires SM90 and a compiled SpargeAttention installation", +) + + +def _masked_reference(q, k, v, index, counts, scale, key_block_size=128): + logits = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * scale + mask = torch.zeros_like(logits, dtype=torch.bool) + for b in range(q.shape[0]): + for h in range(q.shape[2]): + for qb in range(index.shape[2]): + q_slice = slice(qb * 64, min((qb + 1) * 64, q.shape[1])) + for slot in range(int(counts[b, h, qb])): + kb = int(index[b, h, qb, slot]) + k_slice = slice( + kb * key_block_size, + min((kb + 1) * key_block_size, k.shape[1]), + ) + mask[b, h, q_slice, k_slice] = True + logits.masked_fill_(~mask, -float("inf")) + p = torch.softmax(logits, dim=-1) + return torch.einsum("bhqk,bkhd->bqhd", p, v.float()).to(torch.bfloat16) + + +def _cosine(a, b): + return float( + torch.nn.functional.cosine_similarity( + a.float().flatten(), b.float().flatten(), dim=0 + ) + ) + + +@requires_native_sm90_sage +class TestSubBlockSageFp8NativeSm90(CustomTestCase): + def test_full_budget_ragged_tail_reproduces_dense(self): + from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + subblock_sage_fp8_sm90_attention, + ) + + torch.manual_seed(19) + seq_len = 1024 + 37 + heads = 4 + shape = (1, seq_len, heads, 128) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + scale = 1.0 / math.sqrt(128) + + query_blocks = math.ceil(seq_len / 64) + key_blocks = math.ceil(seq_len / 128) + index = ( + torch.arange(key_blocks, device="cuda", dtype=torch.int32) + .view(1, 1, 1, key_blocks) + .expand(1, heads, query_blocks, key_blocks) + .contiguous() + ) + output = subblock_sage_fp8_sm90_attention(q, k, v, index, key_blocks, scale) + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + scale=scale, + ).transpose(1, 2) + + self.assertTrue(torch.isfinite(output.float()).all()) + self.assertGreater(_cosine(output, reference), 0.998) + + def test_sparse_k128_plan_and_variable_counts(self): + from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import ( + subblock_sage_fp8_sm90_attention, + ) + + torch.manual_seed(23) + seq_len = 1024 + 37 + heads = 4 + shape = (1, seq_len, heads, 128) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + scale = 1.0 / math.sqrt(128) + + query_blocks = math.ceil(seq_len / 64) + key_blocks = math.ceil(seq_len / 128) + width = 4 + index = torch.stack( + [ + torch.roll(torch.arange(key_blocks), shifts=query_block)[:width] + for query_block in range(query_blocks) + ] + ) + index = ( + index.to(device="cuda", dtype=torch.int32) + .view(1, 1, query_blocks, width) + .expand(1, heads, query_blocks, width) + .contiguous() + ) + counts = ( + ((torch.arange(query_blocks, device="cuda", dtype=torch.int32) % width) + 1) + .view(1, 1, query_blocks) + .expand(1, heads, query_blocks) + .contiguous() + ) + + output = subblock_sage_fp8_sm90_attention(q, k, v, index, width, scale, counts) + reference = _masked_reference(q, k, v, index, counts, scale) + + self.assertTrue(torch.isfinite(output.float()).all()) + self.assertGreater(_cosine(output, reference), 0.997) + + +if __name__ == "__main__": + unittest.main(verbosity=3)