minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
HuangJi
2026-09-20 16:40:28 +08:00
committed by GitHub
parent 414adef060
commit 2a0cb2f04e
6 changed files with 277 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
+13 -6
View File
@@ -267,7 +267,7 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
)
if compute_mode == "sage_fp8":
capability = current_platform.get_device_capability()
if capability is None or capability.to_int() != 90:
if capability is None or capability.to_int() not in (90, 120):
found = (
capability.as_version_str()
if capability is not None
@@ -275,14 +275,21 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
)
raise ValueError(
"MiniMax-H3 SubBlock compute_mode='sage_fp8' currently "
"requires SM90 (compute capability 9.0); "
"requires SM90 or SM120 (compute capability 9.0 or 12.0); "
f"found {found}."
)
from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
_load_sparge_attention_sm90_ops,
)
if capability.to_int() == 90:
from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
_load_sparge_attention_sm90_ops,
)
_load_sparge_attention_sm90_ops()
_load_sparge_attention_sm90_ops()
else:
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
_load_sm120_sage_ops,
)
_load_sm120_sage_ops()
if selected_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3:
if server_args.ring_degree > 1:
raise ValueError(
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md
+20 -4
View File
@@ -4,8 +4,8 @@ Routes a SubBlock plan to SGLang's CuTe-DSL block-sparse FlashAttention kernel
on SM90 or FlashInfer's architecture-specific blk64 kernels on SM100 and SM120.
The architecture-neutral `"compute_mode": "sage_fp8"` selects online Sage-style
FP8 compute; its current SM90 implementation uses the native SageAttention2
INT8-QK/FP8-PV Hopper kernel. Future SM100 and SM120 Sage FP8 implementations
can register under the same configuration value.
INT8-QK/FP8-PV Hopper kernel. SM120 uses FlashInfer's CuTe-DSL Sage kernel
with Q64 x K64 blocks. SM100 supports BF16 only.
Nothing is trained and no weights change: a cheap estimator runs before
attention and hands the selected kernel a `q2k_block_index`.
@@ -46,7 +46,7 @@ are listed below.
| | |
| --- | --- |
| GPU | **compute capability 9.0, 10.0, or 12.0** — H100 / H200 use SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel or the current `sage_fp8` implementation; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` BF16 kernel; SM120 devices use FlashInfer's `bsa_attn_sm120_blk64_fwd` CuTe-DSL BF16 kernel. Other capabilities, including 10.3 (B300 / GB300), are rejected. |
| GPU | **compute capability 9.0, 10.0, or 12.0** — H100 / H200 use SGLang's CuTe-DSL SM90 block-sparse FlashAttention kernel or the current `sage_fp8` implementation; B200 / GB200 use FlashInfer's architecture-specific `sm_100a` BF16 kernel; SM120 devices use FlashInfer's CuTe-DSL BF16 or Sage FP8 kernel. Other capabilities, including 10.3 (B300 / GB300), are rejected. |
| dtype | bfloat16 |
| head_dim | 128 |
| attention | non-causal, one contiguous sequence per call |
@@ -106,7 +106,17 @@ the normal router's 16-token key pooling cells. Install the optional dependency:
pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation
```
Enable it on SM90 with:
SM120 requires FlashInfer's CuTe-DSL SM120 Sage backend and
`quantize_sage_qkv_sm120`, added in
[FlashInfer #4691](https://github.com/flashinfer-ai/flashinfer/pull/4691).
The default FlashInfer `0.6.18` pin does not include these APIs. After installing
SGLang, install the FlashInfer source revision used for SM120 validation:
```bash
pip install "git+https://github.com/flashinfer-ai/flashinfer.git@6a84331eb6013e5e61018dc2be532ae90520d30f"
```
Enable it on SM90 or SM120 with:
```bash
--attention-backend-config '{"compute_mode":"sage_fp8"}'
@@ -195,3 +205,9 @@ Tests: `test/unit/test_subblock_sparse_attention.py` and
test covers the native production dispatch. Running at a full block budget must
reproduce dense attention up to the expected quantization error, pinning routing
indices, ragged tails, scale domains and the softmax scale in one check.
Run the SM120 Sage numerical test with:
```bash
python test/manual/attention/test_subblock_sage_fp8_sm120.py
```
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py
+63 -15
View File
@@ -21,10 +21,9 @@ individual keys of the defaults below::
Requirements inherited from the kernels: compute capability 9.0 (Hopper) or
10.0/12.0 (Blackwell), bf16, head_dim 128. Hopper uses SGLang's CuTe-DSL SM90
block-sparse FlashAttention kernel by default; ``compute_mode="sage_fp8"``
selects its native SageAttention2 INT8-QK/FP8-PV kernel. The mode name is stable
across architectures, so future SM100/SM120 implementations can use their own
Sage FP8 arithmetic without changing server configuration. B200 and SM120
devices currently use FlashInfer's architecture-specific blk64 BF16 kernels.
selects native SageAttention2 on SM90 and FlashInfer's CuTe-DSL Sage kernel
on SM120, both using INT8 QK and FP8 PV. SM100 uses the BF16 kernel because
its Sage path does not support SubBlock's variable counts and partial blocks.
Inside the DiT, unsupported calls run dense instead; unsupported GPU
architectures are rejected.
@@ -110,6 +109,13 @@ DEFAULT_COMPUTE_MODE = "bf16"
_DIT_LAYER_PREFIX = re.compile(r"^blocks\.(\d+)\.")
def _sage_key_block_size() -> int:
# Resolve on the worker's current device, after device initialization.
if torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0):
return SUBBLOCK_SPARSE_BLOCK_SIZE
return SAGE_FP8_SM90_KEY_BLOCK_SIZE
def _dit_layer_index(prefix: str) -> int | None:
match = _DIT_LAYER_PREFIX.match(prefix)
return int(match.group(1)) if match else None
@@ -233,6 +239,50 @@ def _sm100_sparse_attention(
return out[0] if isinstance(out, tuple) else out
@functools.lru_cache(maxsize=1)
def _load_sm120_sage_ops():
try:
from flashinfer.cute_dsl.sparse.bsa_attn_sm120 import (
bsa_attn_sm120_blk64_sage_fwd,
)
from flashinfer.cute_dsl.sparse.bsa_utils.sage_quant_sm120 import (
quantize_sage_qkv_sm120,
)
except (ImportError, OSError) as exc:
raise ImportError(
"SM120 SubBlock sage_fp8 requires FlashInfer with the CuTe-DSL "
"SM120 Sage backend and quantize_sage_qkv_sm120."
) from exc
return quantize_sage_qkv_sm120, bsa_attn_sm120_blk64_sage_fwd
def _sm120_sage_fp8_sparse_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_block_index: torch.Tensor,
topk: int,
softmax_scale: float,
block_counts: torch.Tensor | None = None,
) -> torch.Tensor:
"""Quantize BSHD activations online and execute a Q64 x K64 Sage plan."""
quantize, attention = _load_sm120_sage_ops()
q_hnd, k_hnd, v_hnd = (x.transpose(1, 2).contiguous() for x in (q, k, v))
quantized = quantize(q_hnd, k_hnd, v_hnd)
out = attention(
*quantized,
q2k_block_index.contiguous(),
topk,
block_sizes=_cached_block_sizes(k.shape[1], k.device),
q2k_block_nums=(
block_counts.contiguous() if block_counts is not None else None
),
softmax_scale=softmax_scale,
backend="cute_dsl",
)
return out.transpose(1, 2).contiguous()
def _sm120_sparse_attention(
q: torch.Tensor,
k: torch.Tensor,
@@ -273,16 +323,13 @@ def _get_subblock_sparse_attention_runner(
if capability == (10, 0):
if compute_mode != "bf16":
raise RuntimeError(
f"SubBlock compute_mode={compute_mode!r} currently targets SM90; "
"the SM100 FlashInfer Sage adapter is not wired in SGLang yet."
"SM100 FlashInfer Sage does not support SubBlock variable block "
"counts and partial blocks; use compute_mode='bf16'."
)
return _sm100_sparse_attention
if capability == (12, 0):
if compute_mode != "bf16":
raise RuntimeError(
f"SubBlock compute_mode={compute_mode!r} currently targets SM90; "
"the SM120 FlashInfer Sage adapter is not wired in SGLang yet."
)
if compute_mode == "sage_fp8":
return _sm120_sage_fp8_sparse_attention
return _sm120_sparse_attention
raise RuntimeError(
"SubBlock sparse attention supports compute capability 9.0, 10.0, or 12.0; "
@@ -378,9 +425,10 @@ class SubBlockSparseSchedule(msgspec.Struct, frozen=True):
config = get_global_server_args().attention_backend_config or {}
compute_mode = str(config.get("compute_mode", DEFAULT_COMPUTE_MODE))
# Hopper's native kernel consumes 128-token K blocks. Eight sub-blocks
# preserve the default router's 16-token key pooling cells.
default_n_k = 8 if compute_mode == "sage_fp8" else DEFAULT_N_K
# Use 16-token pooling cells for both K128 (SM90) and K64 (SM120).
default_n_k = (
_sage_key_block_size() // 16 if compute_mode == "sage_fp8" else DEFAULT_N_K
)
schedule = SubBlockSparseSchedule(
sparsity=float(config.get("sparsity", DEFAULT_SPARSITY)),
skip_first_steps=int(
@@ -453,7 +501,7 @@ class SubBlockSparseAttentionImpl(AttentionImpl):
n_k=self.schedule.n_k,
n_q=self.schedule.n_q,
block_size_k=(
SAGE_FP8_SM90_KEY_BLOCK_SIZE
_sage_key_block_size()
if self.schedule.compute_mode == "sage_fp8"
else SUBBLOCK_SPARSE_BLOCK_SIZE
),
python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py
+21 -8
View File
@@ -167,9 +167,15 @@ class TestSubBlockSparseSchedule(unittest.TestCase):
self.assertEqual(schedule.compute_mode, "bf16")
def test_sage_fp8_uses_16_token_key_subblocks_by_default(self):
with _patch_schedule({"compute_mode": "sage_fp8"}):
schedule = SubBlockSparseSchedule.from_server_args()
self.assertEqual(schedule.n_k, 8)
for capability, expected_n_k in (((9, 0), 8), ((12, 0), 4)):
with (
self.subTest(capability=capability),
patch("torch.cuda.is_available", return_value=True),
patch("torch.cuda.get_device_capability", return_value=capability),
_patch_schedule({"compute_mode": "sage_fp8"}),
):
schedule = SubBlockSparseSchedule.from_server_args()
self.assertEqual(schedule.n_k, expected_n_k)
def test_explicit_sage_fp8_n_k_is_respected(self):
with _patch_schedule({"compute_mode": "sage_fp8", "n_k": 4}):
@@ -320,11 +326,18 @@ class TestSubBlockGating(unittest.TestCase):
with _patch_step(20):
self.assertFalse(impl._sparse_ready(q, q))
def test_sage_fp8_builds_the_sm90_64x128_router(self):
impl = self._impl("blocks.9.attn", compute_mode="sage_fp8")
self.assertEqual(impl.router.block_size_k, 128)
self.assertEqual(impl.router.n_k, 8)
self.assertEqual(impl.router.budget_granularity, 1)
def test_sage_fp8_builds_architecture_specific_router(self):
for capability, key_block_size, n_k in (((9, 0), 128, 8), ((12, 0), 64, 4)):
with (
self.subTest(capability=capability),
patch("torch.cuda.is_available", return_value=True),
patch("torch.cuda.get_device_capability", return_value=capability),
):
impl = self._impl("blocks.9.attn", compute_mode="sage_fp8")
self.assertEqual(impl.router.block_size_k, key_block_size)
self.assertEqual(impl.router.n_k, n_k)
self.assertEqual(impl.router.block_size_k // impl.router.n_k, 16)
self.assertEqual(impl.router.budget_granularity, 1)
@requires_subblock_kernel