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

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"]

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(

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.

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 (

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<n_q, b<n_k} exp( qbar_{i,a} . kbar_{j,b} * softmax_scale )
 
@@ -17,8 +20,9 @@ which is a direct estimate of the block's un-normalised softmax mass
 ``sum_{r in i, c in j} exp(q_r . k_c * scale)`` -- the quantity that decides how much
 attention mass is lost when the block is skipped.
 
-Measured on 567 (task x denoise-step x layer x head) samples of MiniMax-H3 DiT attention,
-mean recall of the retained softmax mass at 0.9 block sparsity:
+Measured with Q64 x K64 geometry on 567 (task x denoise-step x layer x head)
+samples of MiniMax-H3 DiT attention, mean recall of the retained softmax mass at
+0.9 block sparsity:
 
     n_q=1 n_k=1     .6513     4 u        <- plain avg pooling
     n_q=1 n_k=2     .6598     8 u
@@ -35,7 +39,9 @@ to score against, the query detail averages out. Splitting both together is a di
 proposition -- the log-sum-exp then runs over query-key sub-block *pairs*, and "some part
 of this query block wants some part of that key block" is a signal that survives the
 averaging. That is the best row in the table, and it is the only estimator change in this
-family that has separated from anything else end to end. ``n_q = n_k = 4`` ships.
+family that has separated from anything else end to end. The Q64 x K64 BF16 paths
+ship with ``n_q = n_k = 4``. SM90 ``sage_fp8`` uses ``n_q = 4, n_k = 8`` with
+Q64 x K128, preserving the same 16-token pooling cells.
 
 Not worth retrying without new evidence
 ---------------------------------------
@@ -63,7 +69,7 @@ the pipeline currently produces.
 
 Usage
 -----
-    router = SubBlockRouter(n_k=4, n_q=4)
+    router = SubBlockRouter(n_k=4, n_q=4)  # current BF16 geometry
     plan = router.route(q, k, sparsity=0.8, softmax_scale=d**-0.5)   # q, k: [B, S, H, D]
     out, _ = bsa_attn_blk64_fwd(q, k, v, plan.index, plan.topk,
                                 block_sizes=SubBlockRouter.block_sizes(S, q.device),
@@ -145,9 +151,9 @@ def load_bsa_attn_sm120_blk64_fwd():
 
 
 LOG2E = 1.4426950408889634
-BLOCK = 64  # the kernel's block granularity (kSparseBlockSize=64)
-BUDGET_GRANULARITY = 8  # blocks per query row the kernel bills in, padding to fit
-VALID_N = (1, 2, 4, 8)  # sub-blocks per 64-token block -> 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)
 

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.<idx>.attn`` is a DiT layer; ``token_refiner.blocks.<idx>.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(

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,

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