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

[MiniMax-H3] SubBlock: training-free block-sparse attention for the DiT (#34148)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
forrestl
2026-08-11 11:11:11 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent dd20826e0a
commit 704808ed27
8 changed files with 1545 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md
+169
View File
@@ -0,0 +1,169 @@
# SubBlock sparse attention — training-free block sparsity for the MiniMax-H3 DiT
Routes FlashInfer's 64-token block-sparse kernel (`bsa_attn_blk64_fwd`) with a
sub-block score. Nothing is trained and no weights change: a cheap estimator
runs before attention and hands the kernel a `q2k_block_index`.
Spelled out in full, with every key at its default — which is the recommended
configuration and what every number below was measured at:
```bash
sglang serve --model-path MiniMaxAI/MiniMax-H3 --model-variant fl2va \
--num-gpus 8 --ulysses-degree 8 --performance-mode speed \
--attention-backend subblock_sparse_attn \
--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}'
```
**`text_encoder=fa` is not optional.** `--attention-backend` applies to every
component, and the Qwen3-VL text encoder admits only `fa` / `torch_sdpa` /
`sage_attn_3`; without the override it raises and the server never starts. Put
the override on the *encoder*, not the DiT — `transformer=subblock_sparse_attn`
appears to work and silently does nothing, because H3 resolves the DiT backend
lazily on the first forward, outside the component-loading context that the
override applies to.
`--attention-backend-config` is optional and overrides only the keys it names,
so `'{"sparsity": 0.85}'` alone trades quality for another 6%. Inline JSON gets
mangled by `shlex.split`; pass a **file path** instead if the shell eats the
quotes.
## What it runs on
Everything below comes from `bsa_attn_blk64_fwd`, not from this backend.
| | |
| --- | --- |
| GPU | **compute capability 10.0 only** — B200 / GB200 class. The kernel is built `-gencode=arch=compute_100a,code=sm_100a`, which is arch-specific and does not forward-run on 10.3 (B300 / GB300) or 12.x (RTX PRO 6000, RTX 50xx). |
| dtype | bfloat16 |
| head_dim | 128 |
| attention | non-causal, one contiguous sequence per call |
Inside the DiT, anything the kernel cannot serve — cross attention, the token
refiner, sequences under `min_seq_len`, non-bf16 activations, head_dim != 128 —
falls back to dense for that call, so no layer has to be excluded by hand.
**On an unsupported GPU it is not a fallback, it is an error at startup.** The
resolver checks the compute capability before anything loads and refuses
anything but 10.0, so an H100 or a B300 fails at launch rather than after ten
dense denoise steps. Do not rely on the kernel's own guard for this: it compares
only the major version, so it would accept 10.3 and then fail with no cubin.
## How the score works
The usual proxy for a 64x64 block is `mean(Q_block) · mean(K_block)`. Averaging
64 keys into one vector destroys exactly the variation that decides which keys a
query wants. So each block is cut into sub-blocks — `n_k` on the key side, `n_q`
on the query side — and every sub-block pair is scored and combined with a
log-sum-exp:
```
score(i, j) = log Σ_{a,b} exp( mean(Q_{i,a}) · mean(K_{j,b}) · softmax_scale )
```
which estimates the block's un-normalised softmax mass directly — the quantity
that says how much is lost by skipping the block.
Splitting the query side *alone* is worse than not splitting: a block's mass sums
over its query rows, so with one key vector to score against the query detail
averages out. Splitting both together is a different proposition, and the only
estimator change in this family that has held up end to end. `router.py` carries
the recall table behind `n_q = n_k = 4` and the record of what was tried and
rejected.
## Configuration
| 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_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 |
**`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.
**`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
score is the log-sum-exp over all `n_q * n_k` sub-block pairs. Raising them
sharpens the estimate of which blocks carry mass, at `n_q * n_k` times the score
matrix — 0.5% of denoise time at the default, so cost is not the constraint.
Raise `n_q` and `n_k` **together**: splitting the query side alone is worse than
not splitting at all.
**The two schedule cutoffs are asymmetric on purpose.** `skip_first_steps` keeps
the leading denoise forwards dense; those steps settle the layout of the sample
and visibly re-frame the shot when approximated — lowering it from 10 to 5 halves
cosine against the dense render. Depth does not behave that way, so
`skip_first_layers` defaults to 0 and every DiT layer is sparse. Do not lower
`skip_first_steps` without looking at the output.
**`min_seq_len` is a floor, not a tuning knob.** Below it the whole call runs
dense, and in packed varlen batches the test is per document, so H3's padding
tail goes dense while the 37.7k-token media document is routed. Two things break
down on short sequences: the router is four fixed Triton launches against an
attention cost that falls as S², so the overhead stops paying for itself; and the
budget goes coarse — 4096 keys is only 64 blocks, and at 1024 keys the
multiple-of-8 floor already keeps half of them. 4096 sits well below any real
video sequence and well above where either effect bites; it was chosen on that
reasoning rather than from a measured threshold sweep.
## Measured
MiniMax-H3 t2va, 1344x768 / 5 s / 50 steps, 8x B200, Ulysses-8, bf16, at the
shipped defaults (152 of 590 key blocks per query block). All arms in one session
on one node, cold sample dropped; spread within an arm is under 0.07 s.
| | DiT denoise | vs dense |
| --- | ---: | ---: |
| dense (FlashAttention) | 18.270 s | 1.000x |
| SubBlock sparse | 16.061 s | **1.138x** |
| SubBlock sparse + [flashinfer#4397][fi] | 15.012 s | **1.217x** |
[flashinfer-ai/flashinfer#4397][fi] rebuilds the kernel's internal Q/K/V tile
layout in one pass instead of three. It is bit-identical and **not required**:
worth 1.070x on its own.
[fi]: https://github.com/flashinfer-ai/flashinfer/pull/4397
Sparsity is the speed lever and it saturates — 0.75 gives 1.136x, 0.80 gives
1.178x, 0.85 gives 1.211x. Cutting the budget 40% past 0.75 buys 6%, because at
37.7k tokens attention is no longer the bulk of the step, and 0.85 rendered worst
of the three on cosine against dense. `n_k` moves the denoise time by 0.3% across
its whole range: it is a quality knob, not a speed one.
**The speedup is bounded by sequence length, not by the method.** The same config
measured 1.13x at 37.7k tokens, 1.20x at 52k and 1.47x at 96k — the backend only
touches attention, and attention's share of the DiT grows with S. Treat 1.2x as
the 768p/5 s number, not the ceiling.
The same effect shows up in the sequence-parallel degree, since that sets how
much of the sequence each GPU holds: on 4x B200 at Ulysses-4 the identical
config gives **1.168x** on denoise and **1.138x** end to end, against 1.138x on
denoise at Ulysses-8.
Peak memory is unchanged (99,356 vs 99,358 MiB/GPU): block sparsity saves
compute, not activations, and the `[B,H,Gq,Gk]` score matrix is ~20 MB at
S=37.7k. Absolute times are node-specific; only ratios measured in one session
are comparable.
## Files
| | |
| --- | --- |
| `router.py` | `SubBlockRouter` — pooling, scoring, selection, `RoutingPlan` |
| `kernels.py` | Triton pooling / segmented-LSE / fused top-k |
| `../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.
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/__init__.py
+15
View File
@@ -0,0 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
"""SubBlock -- training-free block-sparse attention routing for video DiTs.
Originally vendored from the standalone SubBlock repository; ``router.py`` and
``kernels.py`` have since diverged from it.
``router.py`` scores every (query block, key block) pair from sub-block-pooled
Q/K and turns the scores into the ``q2k_block_index`` that FlashInfer's
``bsa_attn_blk64_fwd`` consumes (SM100, bf16, head_dim 128). The estimator and
the measurements behind its defaults are documented there.
"""
from .router import SubBlockRouter, load_bsa_attn_blk64_fwd
__all__ = ["SubBlockRouter", "load_bsa_attn_blk64_fwd"]
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/kernels.py
+300
View File
@@ -0,0 +1,300 @@
# SPDX-License-Identifier: Apache-2.0
"""Fused sub-block score kernel: bf16 GEMM + segmented log-sum-exp in one pass.
score[i, j] = log2 sum_{b < n_k} 2 ** ( qbar_i . kbar_{j,b} )
The PyTorch router materialises ``[B, H, Gq, Gk*n_k]`` fp32 (254 MB at S=96k, n_k=4),
reduces it, and throws it away -- the GEMM is 4% of its time, the rest is that tensor's
round trip to HBM. Here the reduction happens in registers before anything is written, so
only ``[B, H, Gq, Gk]`` ever reaches memory: n_k times less traffic.
exp2/log2 are used internally (they are the hardware instructions; the caller folds
``softmax_scale * log2(e)`` into Q), and the result is converted back to natural-log units
so it matches the reference implementation exactly, not just up to ranking.
Padding: sub-cells are ordered, and validity is monotone, so a single ``n_valid`` scalar
(the number of key sub-cells holding at least one real token) is enough -- everything at or
past it is forced to -inf so it can never win a slot.
"""
import math
import torch
import triton
import triton.language as tl
_NEG = tl.constexpr(-1.0e30) # Triton only lets @jit read constexpr globals
_LN2 = tl.constexpr(0.6931471805599453)
@triton.jit
def _score_kernel(
Q,
K,
O,
stride_qm,
stride_ql,
stride_kn,
stride_kl,
stride_om,
stride_on,
stride_ol,
M,
M_VALID,
N_VALID,
NOUT,
MOUT,
BLK_M: tl.constexpr,
BLK_N: tl.constexpr,
NK: tl.constexpr,
NQR: tl.constexpr,
D: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
pid_l = tl.program_id(2)
offs_m = pid_m * BLK_M + tl.arange(0, BLK_M)
offs_n = pid_n * BLK_N + tl.arange(0, BLK_N)
offs_d = tl.arange(0, D)
q = tl.load(
Q + pid_l * stride_ql + offs_m[:, None] * stride_qm + offs_d[None, :],
mask=offs_m[:, None] < M,
other=0.0,
)
k = tl.load(
K + pid_l * stride_kl + offs_n[:, None] * stride_kn + offs_d[None, :],
mask=offs_n[:, None] < N_VALID,
other=0.0,
)
acc = tl.dot(q, tl.trans(k), out_dtype=tl.float32) # [BLK_M, BLK_N]
# a sub-cell past the last real one must not contribute to its group's log-sum-exp
acc = tl.where(offs_n[None, :] < N_VALID, acc, _NEG)
# same on the query side: with NQ > 1 the last query block can own sub-cells that
# are entirely padding, and those pool to zero -- an exp2(0) = 1 term that would
# otherwise be folded into the block's score.
acc = tl.where(offs_m[:, None] < M_VALID, acc, _NEG)
acc = tl.reshape(acc, (BLK_M, BLK_N // NK, NK))
m = tl.max(acc, axis=2)
s = tl.sum(tl.exp2(acc - m[:, :, None]), axis=2)
lse = m + tl.log2(s)
lse = tl.where(m > _NEG / 2, lse, _NEG) # whole group was padding
# Fold the NQR query sub-cells of a query block together. Log-sum-exp is
# associative, so reducing NK then NQ is the same one log-sum-exp over all
# NQ*NK sub-block pairs -- and two stages keeps both reductions on an axis
# that is already contiguous in registers.
if NQR > 1:
lse = tl.reshape(lse, (BLK_M // NQR, NQR, BLK_N // NK))
m2 = tl.max(lse, axis=1)
s2 = tl.sum(tl.exp2(lse - m2[:, None, :]), axis=1)
lse = tl.where(m2 > _NEG / 2, m2 + tl.log2(s2), _NEG)
# exp2/log2 internally (they map to the hardware instructions), then back to natural
# log units so the fused and reference backends return the same numbers, not just the
# same ranking. One multiply in registers.
out = lse * _LN2
out = out.to(O.dtype.element_ty) # bf16 halves what the selection step has to read
offs_o = pid_n * (BLK_N // NK) + tl.arange(0, BLK_N // NK)
offs_q = pid_m * (BLK_M // NQR) + tl.arange(0, BLK_M // NQR)
tl.store(
O
+ pid_l * stride_ol
+ offs_q[:, None] * stride_om
+ offs_o[None, :] * stride_on,
out,
mask=(offs_q[:, None] < MOUT) & (offs_o[None, :] < NOUT),
)
BLK_M = BLK_N = 128 # score tile; must hold whole blocks, so a multiple of n_q and n_k
def fused_scores(qp, kp, out, *, n_k, n_valid, n_q, m_valid):
"""qp: [L, Gq*n_q, D] bf16 (already carrying softmax_scale*log2e),
kp: [L, Gk*n_k, D] bf16 -> out: [L, Gq, Gk], natural-log scores.
The n_q query sub-cells of a query block are folded together by log-sum-exp.
``n_valid`` / ``m_valid`` are the counts of key / query sub-cells holding at
least one real token; the rest pooled to zero and must not contribute.
"""
L, M, D = qp.shape
N = kp.shape[1]
Mout, Nout = out.shape[1], out.shape[2]
grid = (triton.cdiv(M, BLK_M), triton.cdiv(N, BLK_N), L)
_score_kernel[grid](
qp,
kp,
out,
qp.stride(1),
qp.stride(0),
kp.stride(1),
kp.stride(0),
out.stride(1),
out.stride(2),
out.stride(0),
M,
m_valid,
n_valid,
Nout,
Mout,
BLK_M=BLK_M,
BLK_N=BLK_N,
NK=n_k,
NQR=n_q,
D=D,
num_warps=4,
num_stages=3,
)
return out
@triton.jit
def _pool_kernel(
X,
Y,
stride_xb,
stride_xt,
stride_xh,
stride_yl,
stride_yn,
S,
H,
SUB: tl.constexpr,
D: tl.constexpr,
SCALE,
):
"""[B, S, H, D] -> [B*H, n_cells, D]: masked mean of every SUB consecutive tokens.
Fused so the activation is read once and the pooled result is written straight in
bf16; the PyTorch version needs an fp32 temporary plus a transpose.
"""
cell = tl.program_id(0)
l = tl.program_id(1)
b = l // H
h = l % H
offs_t = cell * SUB + tl.arange(0, SUB)
offs_d = tl.arange(0, D)
mask = offs_t < S
x = tl.load(
X
+ b * stride_xb
+ offs_t[:, None] * stride_xt
+ h * stride_xh
+ offs_d[None, :],
mask=mask[:, None],
other=0.0,
).to(tl.float32)
cnt = tl.sum(mask.to(tl.float32), axis=0)
acc = tl.sum(x, axis=0) / tl.maximum(cnt, 1.0) * SCALE
tl.store(Y + l * stride_yl + cell * stride_yn + offs_d, acc.to(tl.bfloat16))
def fused_pool(x, n_cells, sub, out, scale=1.0):
"""x: [B, S, H, D] bf16 -> out: [B*H, n_cells, D] bf16"""
B, S, H, D = x.shape
_pool_kernel[(n_cells, B * H)](
x,
out,
x.stride(0),
x.stride(1),
x.stride(2),
out.stride(0),
out.stride(1),
S,
H,
SUB=sub,
D=D,
SCALE=scale,
# one warp, not four: the tile is only SUB x 128, so extra warps buy no parallelism
# and cost scheduling. Measured 0.272 -> 0.074 ms at S=96k, 1.27 -> 4.7 TB/s.
num_warps=1,
)
return out
@triton.jit
def _topk_kernel(S, OUT, G, K, BLK: tl.constexpr, ITERS: tl.constexpr):
"""Exact-enough per-row top-K in a single pass over global memory.
A score row is only Gk values -- 3 KB in bf16 at S=96k -- so it fits in registers. Load
it once, then do everything on chip. torch.topk instead makes several passes over the
whole matrix, which is why it costs half the router.
The threshold search interpolates on the count rather than halving the interval: the
count-vs-threshold curve is the row's empirical CDF, so a secant step lands far closer
than a bisection step. Full-row reductions are what this kernel pays for, so fewer
steps is the whole game.
The invariant is count(s >= lo) >= K, so the compaction can only over-fill, never
under-fill; `pos < K` truncates the boundary group. At ITERS=12 about 1.6% of rows end
up with a different set than an exact top-K, but only among blocks that tie at the
threshold -- the total selected score differs by ~1e-6 relative, which is nothing.
"""
row = tl.program_id(0)
offs = tl.arange(0, BLK)
m = offs < G
s = tl.load(S + row * G + offs, mask=m, other=-float("inf")).to(tl.float32)
lo = tl.min(tl.where(m, s, float("inf")))
hi = tl.max(tl.where(m, s, -float("inf"))) + 1.0
clo = tl.sum(m.to(tl.int32), axis=0).to(tl.float32)
chi = 0.0
for _ in tl.static_range(ITERS):
den = clo - chi
t = (clo - K) / tl.where(den > 0.5, den, 1.0)
t = tl.minimum(tl.maximum(t, 0.05), 0.95) # keep the step inside the bracket
mid = lo + (hi - lo) * t
cnt = tl.sum(((s >= mid) & m).to(tl.int32), axis=0).to(tl.float32)
take = cnt >= K
lo = tl.where(take, mid, lo)
clo = tl.where(take, cnt, clo)
hi = tl.where(take, hi, mid)
chi = tl.where(take, chi, cnt)
sel = (s >= lo) & m
pos = tl.cumsum(sel.to(tl.int32), axis=0) - 1
tl.store(OUT + row * K + pos, offs.to(tl.int32), mask=sel & (pos < K))
def topk_iters(G, k):
"""Threshold-search steps needed to match ``torch.topk``.
The search interpolates on the *count* above a trial threshold, which assumes the count
is linear in the threshold. That holds near the median but not in the tail, where the
score density decays roughly exponentially and the secant undershoots -- so the further
into the tail k sits, the more steps are needed. Measured on 21 real files as the
fraction of query rows selecting a different block set than exact top-k:
log2(G/k) 1.0 2.1 3.3 4.3 5.6 (sparsity .50 .76 .90 .95 .98)
16 steps 0.09% 0.14% 1.62% 5.17% 1.88%
24 steps 0.02% 0.01% 0.14% 0.61% 0.21%
32 steps 0.02% 0.01% 0.02% 0.19% 0.05%
A flat 16 is fine at the usual operating points and silently wrong past sparsity 0.9
(+3.2% relative L2 at 0.90, +7.1% at 0.95). These cutoffs hold every regime under 0.2%.
Interpolating on log(count) instead linearises the tail and does fix sparsity >= 0.95,
but it is far worse where the tail model does not apply (27% differing rows at sparsity
0.5), so the step count is the robust knob, not the model.
"""
L = math.log2(max(G, 1) / max(k, 1))
return 16 if L <= 2.5 else 24 if L <= 3.5 else 32
def fused_topk(scores2d, k):
"""scores2d: [rows, G] contiguous -> [rows, k] int32 column ids (unsorted)."""
rows, G = scores2d.shape
out = torch.empty(rows, k, dtype=torch.int32, device=scores2d.device)
_topk_kernel[(rows,)](
scores2d,
out,
G,
k,
BLK=triton.next_power_of_2(G),
ITERS=topk_iters(G, k),
num_warps=4 if G >= 1024 else 2, # short rows do not fill four warps
)
return out
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py
+256
View File
@@ -0,0 +1,256 @@
# SPDX-License-Identifier: Apache-2.0
"""Sub-block block-sparse routing for FlashInfer's ``bsa_attn_blk64_fwd``.
Training-free. Runs *before* attention, produces the ``q2k_block_index`` tensor the
64-token block-sparse kernel consumes.
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:
score(i, j) = log sum_{a<n_q, b<n_k} exp( qbar_{i,a} . kbar_{j,b} * softmax_scale )
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:
n_q=1 n_k=1 .6513 4 u <- plain avg pooling
n_q=1 n_k=2 .6598 8 u
n_q=1 n_k=4 .6655 16 u
n_q=1 n_k=8 .6697 32 u
n_q=8 n_k=1 .6494 32 u <- splitting Q *alone* is worse than not splitting
n_q=8 n_k=8 .6793 256 u <- but splitting both is the best of them
oracle .7355 -
(1 u = one ``[S/128, 128] x [128, S/128]`` GEMM = 1/16384 of the dense attention it gates.)
Splitting Q alone loses: a block's mass sums over its query rows, so with one key vector
to score against, the query detail averages out. Splitting both together is a different
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.
Not worth retrying without new evidence
---------------------------------------
Summing un-normalised sub-block mass over the query axis lets the hottest query sub-block
own a block's score, and the true per-row attention carries a ``1/Z_r`` the raw sum drops,
which over-weights exactly the rows whose attention is spread widest -- the rows that lose
least from dropping any one block. Turning each query sub-block into a distribution over
key blocks first fixes that, and on 1092 real (cell, head, query block) samples at
n_q=n_k=4, sparsity 0.9 it measured better on both proxies: block mass recall .6741 ->
.6779 and relative L2 of the rebuilt attention output .2043 -> .1982, paired t = +8.0 and
-5.7.
It is worse in the pixels, on **0 of 15** prompts, by 0.107 cosine against the dense render
(paired t = -6.4). Single-layer output error, even measured directly, does not order these
estimators the way 40 denoise steps through 50 layers do. Nothing short of an end-to-end
render has predicted this correctly yet -- neither block mass recall nor single-step output
L2.
Worth trying, not yet exposed
-----------------------------
A per-head budget beats any estimator upgrade measured here: at a fixed mean sparsity,
spending more blocks on diffuse heads and fewer on peaked ones lifts 5th-percentile mass
recall from .52 to .90. It needs a rule for setting the per-head split, which nothing in
the pipeline currently produces.
Usage
-----
router = SubBlockRouter(n_k=4, n_q=4)
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),
q2k_block_nums=None)
"""
from __future__ import annotations
import functools
import importlib.util
import math
import sys
import types
from pathlib import Path
import msgspec
import torch
from .kernels import fused_pool, fused_scores, fused_topk
@functools.lru_cache(maxsize=1)
def load_bsa_attn_blk64_fwd():
"""FlashInfer's 64-block sparse attention entry point.
``flashinfer.cute_dsl.sparse.__init__`` also pulls in the blk128 CuTe-DSL backend,
which breaks in ways blk64 does not care about: it hard-requires the ``quack``
package, and it tracks a moving ``cutlass.cute`` API (0.6.15.post1 raises
``AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'``). blk64 is
plain CUDA and needs none of it, so whatever the package import trips over we load
``bsa_attn_blk64.py`` under a synthetic parent package instead -- same file, same
kernel. If blk64 itself is broken or absent, that load raises and the caller sees it.
"""
try:
from flashinfer.cute_dsl.sparse import bsa_attn_blk64_fwd
return bsa_attn_blk64_fwd
except Exception:
pass
import flashinfer
base = Path(flashinfer.__file__).resolve().parent / "cute_dsl" / "sparse"
pkg = "_flashinfer_sparse_blk64_only"
def _load(name: str, path: Path, is_pkg: bool):
spec = importlib.util.spec_from_file_location(
name,
path,
submodule_search_locations=[str(path.parent)] if is_pkg else None,
)
mod = importlib.util.module_from_spec(spec)
sys.modules[name] = mod
spec.loader.exec_module(mod)
return mod
if pkg not in sys.modules:
parent = types.ModuleType(pkg)
parent.__path__ = [str(base)]
sys.modules[pkg] = parent
_load(f"{pkg}.blk64", base / "blk64" / "__init__.py", True)
mod = _load(f"{pkg}.bsa_attn_blk64", base / "bsa_attn_blk64.py", False)
return mod.bsa_attn_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
def _snap_up_to_8(topk: int, num_blocks: int) -> int:
"""Round a block budget up to what the kernel is going to charge for anyway.
``bsa_attn_blk64_fwd`` pads each query row's block count up to a multiple of
``BUDGET_GRANULARITY`` with phantom slots that repeat the last real block and
are then masked out of the softmax. Asking for 148 blocks therefore costs
exactly what 152 costs, with four of the slots computed and thrown away.
Measured at S=37.7k on B200, 2 prompts in one session: 152 blocks take
16.055 s against 16.061 s for 148, and 120 take 15.490 s against 15.496 s
for 118 -- free, inside the noise. So take the blocks.
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)
class RoutingPlan(msgspec.Struct, frozen=True):
"""What the kernel needs, plus the budget that produced it."""
index: torch.Tensor # [B, H, Gq, topk] int32
topk: int # key blocks kept per query block
num_blocks: int # key blocks available
@property
def density(self) -> float:
return self.topk / self.num_blocks
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.
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
the diagonal changed relative L2 by 0.2% and the sink only helped in DiT layers 2-32,
which did not survive to the pixels.
"""
def __init__(self, n_k: int = 4, n_q: int = 4) -> 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}"
)
self.n_k, self.n_q = n_k, n_q
@torch.no_grad()
def scores(
self, q: torch.Tensor, k: torch.Tensor, softmax_scale: float
) -> torch.Tensor:
"""``[B, S, H, D] -> [B, H, Gq, Gk]`` block scores (log-space, higher = keep).
Two Triton kernels: pool, then GEMM + segmented log-sum-exp in registers, so the
``[B, H, Gq*n_q, Gk*n_k]`` intermediate never reaches memory.
``softmax_scale * log2(e)`` is folded into Q so the kernel can use the exp2/log2
hardware instructions; it multiplies by ln 2 on the way out, so scores come back
in natural-log units. Selection is a top-k and any monotone rescale leaves that
alone, so the units only matter to a reader of the magnitudes.
The scores stay **float32**. bf16 would halve what selection reads, but with 8
mantissa bits many blocks tie exactly at the threshold and the fused selector
breaks ties by column index -- which systematically prefers early key blocks, one
region of the video. Measured +3.9% relative L2 at S=96k.
"""
b, s, h, d = q.shape
sk = k.shape[1]
gq, gk = -(-s // BLOCK), -(-sk // BLOCK)
nq, nk = self.n_q, self.n_k
sub_q, sub_k = BLOCK // nq, BLOCK // 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.
# 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.
pooled_q = torch.empty(b * h, gq * nq, d, device=q.device, dtype=torch.bfloat16)
pooled_k = torch.empty(b * h, gk * nk, d, device=k.device, dtype=torch.bfloat16)
fused_pool(q, gq * nq, sub_q, pooled_q, scale=softmax_scale * LOG2E)
fused_pool(k, gk * nk, sub_k, pooled_k)
out = torch.empty(b * h, gq, gk, device=q.device, dtype=torch.float32)
fused_scores(
pooled_q,
pooled_k,
out,
n_k=nk,
n_valid=-(-sk // sub_k),
n_q=nq,
m_valid=-(-s // sub_q),
)
return out.view(b, h, gq, gk)
@torch.no_grad()
def route(
self, q: torch.Tensor, k: torch.Tensor, sparsity: float, softmax_scale: float
) -> RoutingPlan:
"""Select the top ``(1 - sparsity)`` fraction of key blocks per query block."""
b, s, h, d = q.shape
gk = -(-k.shape[1] // BLOCK)
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)
# One pass over the score matrix instead of torch.topk's several; the kernel
# accepts the blocks in any order, so nothing sorts them.
index = fused_topk(scores.reshape(-1, gk), topk).view(b, h, gq, topk)
return RoutingPlan(index=index, topk=topk, num_blocks=gk)
@staticmethod
def block_sizes(seq_len: int, device) -> torch.Tensor:
"""Real token count per 64-block, for the kernel's tail masking."""
g = -(-seq_len // BLOCK)
start = torch.arange(g, device=device, dtype=torch.int32) * BLOCK
return (seq_len - start).clamp(0, BLOCK).to(torch.int32)
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py
+407
View File
@@ -0,0 +1,407 @@
# SPDX-License-Identifier: Apache-2.0
"""SubBlock block-sparse attention backend.
Routes FlashInfer's 64-token block-sparse kernel with a K-side sub-block
log-sum-exp score (see ``backends/subblock_sparse/``). Everything is training-free:
the router runs before attention and produces the ``q2k_block_index`` the
kernel consumes.
Sparsity is not applied everywhere. The early denoise steps settle the layout
of the sample and tolerate approximation badly, so the backend falls back to
dense attention for them. Depth turns out not to matter the same way, which is
why the layer cutoff defaults to zero -- see the defaults below. The schedule
is configured through ``--attention-backend-config``, which overrides
individual keys of the defaults below::
--attention-backend subblock_sparse_attn \
--attention-backend-config '{"sparsity": 0.85}'
Requirements inherited from the kernel: compute capability 10.0 (B200 / GB200
class -- it is built for ``sm_100a``, which does not forward-run on 10.3 or
12.x), bf16, head_dim 128. Inside the DiT, any call the kernel cannot serve --
cross/refiner attention, short sequences, non-bf16 -- runs dense instead. On any
other GPU the resolver refuses the backend at startup rather than falling back.
``--attention-backend`` reaches every component, and the text encoder admits
only fa / torch_sdpa / sage_attn_3, so pair it with
``--component-attention-backends text_encoder=fa``; see the README.
"""
from __future__ import annotations
import functools
import re
from dataclasses import dataclass
from typing import Any
import msgspec
import torch
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
AttentionBackend,
AttentionImpl,
AttentionMetadata,
AttentionMetadataBuilder,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
SubBlockRouter,
load_bsa_attn_blk64_fwd,
)
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# The kernel is fixed at 64-token blocks and 128-wide heads.
SUBBLOCK_SPARSE_HEAD_DIM = 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
# tokens, 0.75 gives 1.14x, 0.80 gives 1.18x and 0.85 gives 1.21x -- cutting the
# block budget by 40% past 0.75 buys 6%, because attention is no longer the bulk
# of the step. 0.85 was the worst arm on cosine-vs-dense on both clips rendered
# across all three grades, and 0.80 costs 0.017 / 0.006 cos_c against 0.75 on
# those same two clips for 3.5% of the time, so the default takes the quality.
DEFAULT_SPARSITY = 0.75
# The two cutoffs were swept independently on MiniMax-H3 t2va (1344x768, 5 s,
# 50 steps, n_k=4, sparsity 0.75) and behave nothing alike. Lowering the step
# cutoff from 10 to 5 halves cosine-vs-dense (0.558 -> 0.310 on two clips) and
# visibly re-frames the shot; going to 0 leaves the sample essentially
# uncorrelated with dense for 1.20x -> 1.30x. Lowering the layer cutoff from 2
# to 0 costs 0.0013 of that cosine -- inside the 0.02 run-to-run noise floor --
# and is worth ~1%, so the first DiT blocks get no special treatment.
DEFAULT_SKIP_FIRST_STEPS = 10
DEFAULT_SKIP_FIRST_LAYERS = 0
DEFAULT_N_K = 4
# Query-side splitting. Splitting Q *alone* is worse than not splitting -- with
# one key vector to score against, the query detail averages out -- which is
# where the "n_q is worthless" reading came from. Splitting both sides together
# is a different estimator: the log-sum-exp then runs over query-key sub-block
# pairs. It is the only estimator change in this family that has reproduced end
# to end, and it costs 0.5% of the denoise time. Measured against n_q=1 on
# fifteen t2va prompts, every arm rendered in one session against that session's
# own dense render, as cosine of the decoded video:
# sparsity 0.90 +0.062 paired t = +2.6 better on 13/15
# sparsity 0.75 +0.008 paired t = +2.1 better on 10/15
# The margin shrinks as the budget loosens, which is the pattern every estimator
# comparison here has followed: at the shipped 148 of 590 blocks the rules mostly
# agree on what to keep.
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
# ``blocks.<idx>.attn`` is a DiT layer; ``token_refiner.blocks.<idx>.attn`` and
# anything else is not and stays dense.
_DIT_LAYER_PREFIX = re.compile(r"^blocks\.(\d+)\.")
def _dit_layer_index(prefix: str) -> int | None:
match = _DIT_LAYER_PREFIX.match(prefix)
return int(match.group(1)) if match else None
@functools.lru_cache(maxsize=8)
def _cached_block_sizes(seq_len: int, device: torch.device) -> torch.Tensor:
"""Per-block real token counts; identical for every layer and step.
Rebuilding it per call costs an arange plus a clamp launch on the critical
path for a tensor that only depends on the sequence length.
"""
return SubBlockRouter.block_sizes(seq_len, device)
class SubBlockSparseAttentionBackend(AttentionBackend):
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [SUBBLOCK_SPARSE_HEAD_DIM]
@staticmethod
def get_enum() -> AttentionBackendEnum:
return AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
@staticmethod
def get_impl_cls() -> type[SubBlockSparseAttentionImpl]:
return SubBlockSparseAttentionImpl
@staticmethod
def get_metadata_cls() -> type[SubBlockSparseAttentionMetadata]:
return SubBlockSparseAttentionMetadata
@staticmethod
def get_builder_cls() -> type[SubBlockSparseAttentionMetadataBuilder]:
return SubBlockSparseAttentionMetadataBuilder
@dataclass
class SubBlockSparseAttentionMetadata(AttentionMetadata):
current_timestep: int
class SubBlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
# The base class declares __init__ abstract, so a builder that does not
# override it cannot be instantiated at all.
def __init__(self) -> None:
pass
def prepare(self) -> None:
pass
def build( # type: ignore[override]
self, current_timestep: int, **kwargs: dict[str, Any]
) -> SubBlockSparseAttentionMetadata:
return SubBlockSparseAttentionMetadata(current_timestep=current_timestep)
class SubBlockSparseSchedule(msgspec.Struct, frozen=True):
"""When sparsity is allowed to apply, and how much of it."""
sparsity: float
skip_first_steps: int
skip_first_layers: int
n_k: int
n_q: int
min_seq_len: int
@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 {}
schedule = SubBlockSparseSchedule(
sparsity=float(config.get("sparsity", DEFAULT_SPARSITY)),
skip_first_steps=int(
config.get("skip_first_steps", DEFAULT_SKIP_FIRST_STEPS)
),
skip_first_layers=int(
config.get("skip_first_layers", DEFAULT_SKIP_FIRST_LAYERS)
),
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)),
)
if not 0.0 <= schedule.sparsity < 1.0:
raise ValueError(
f"subblock sparsity must be in [0, 1), got {schedule.sparsity}"
)
for name, value in (("n_k", schedule.n_k), ("n_q", schedule.n_q)):
if value not in (1, 2, 4, 8):
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")
return schedule
class SubBlockSparseAttentionImpl(AttentionImpl):
"""Block-sparse attention with a dense fallback for the excluded region.
One impl instance is built per attention module, so ``prefix`` fixes the
layer for the lifetime of the object; only the denoise step varies per
call and it comes from the forward context.
"""
def __init__(
self,
num_heads: int,
head_size: int,
causal: bool = False,
softmax_scale: float | None = None,
num_kv_heads: int | None = None,
prefix: str = "",
**extra_impl_args,
) -> None:
self.prefix = prefix
self.num_heads = num_heads
self.head_size = head_size
self.causal = causal
self.softmax_scale = (
softmax_scale if softmax_scale is not None else head_size**-0.5
)
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
self.schedule = SubBlockSparseSchedule.from_server_args()
self.layer_idx = _dit_layer_index(prefix)
# A layer outside the DiT stack (token refiner, cross attention) never
# runs sparse: its sequences are short and its budget meaningless.
self.layer_enabled = (
self.layer_idx is not None
and self.layer_idx >= self.schedule.skip_first_layers
and head_size == SUBBLOCK_SPARSE_HEAD_DIM
and self.schedule.sparsity > 0.0
)
self.router = (
SubBlockRouter(n_k=self.schedule.n_k, n_q=self.schedule.n_q)
if self.layer_enabled
else None
)
self.dense_impl = self._build_dense_impl(causal=causal)
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"{self.schedule.skip_first_steps} denoise steps and the first "
f"{self.schedule.skip_first_layers} DiT layers"
)
def _build_dense_impl(self, *, causal: bool) -> AttentionImpl:
"""Flash attention, used wherever the schedule excludes sparsity."""
from sglang.multimodal_gen.runtime.layers.attention.selector import (
get_attn_backend,
)
backend = get_attn_backend(
self.head_size,
torch.bfloat16,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
},
selected_attention_backend=AttentionBackendEnum.FA,
)
return backend.get_impl_cls()(
num_heads=self.num_heads,
head_size=self.head_size,
causal=causal,
softmax_scale=self.softmax_scale,
num_kv_heads=self.num_kv_heads,
prefix=f"{self.prefix}.dense",
)
def _step_enabled(self) -> bool:
return get_forward_context().current_timestep >= self.schedule.skip_first_steps
def _sparse_ready(self, q: torch.Tensor, k: torch.Tensor) -> bool:
return (
self.layer_enabled
and self._step_enabled()
and q.dtype == torch.bfloat16
and k.shape[-3] >= self.schedule.min_seq_len
and not self.causal
)
def _sparse_attention(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
) -> torch.Tensor:
"""q, k, v: ``[1, S, H, 128]`` bf16 -> same shape."""
bsa_attn_blk64_fwd = load_bsa_attn_blk64_fwd()
plan = self.router.route(
q, k, sparsity=self.schedule.sparsity, softmax_scale=self.softmax_scale
)
# Proof that the sparse path actually ran, with the shape it ran on --
# the construction-time log above only says the layer was eligible.
logger.info_once(
f"SubBlock sparse attention active: S={k.shape[1]} heads={q.shape[2]} "
f"keeping {plan.topk}/{plan.num_blocks} key blocks per query block "
f"(sparsity {1 - plan.density:.4f})"
)
out = bsa_attn_blk64_fwd(
q,
k,
v,
plan.index,
plan.topk,
block_sizes=_cached_block_sizes(k.shape[1], k.device),
q2k_block_nums=None, # the budget is uniform across rows
softmax_scale=self.softmax_scale,
)
return out[0] if isinstance(out, tuple) else out
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: SubBlockSparseAttentionMetadata | None = None,
) -> torch.Tensor:
"""query/key/value: ``[B, S, H, D]``."""
if not self._sparse_ready(query, key):
return self.dense_impl.forward(query, key, value, attn_metadata)
return self._sparse_attention(query, key, value)
def forward_varlen(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
*,
cu_seqlens: torch.Tensor,
max_seqlen: int,
cu_seqlens_host: tuple[int, ...] | None = None,
) -> torch.Tensor:
"""Packed ``[T, H, D]`` rows split into documents by ``cu_seqlens``.
The block-sparse kernel takes one contiguous sequence, so each packed
document is routed on its own. Documents shorter than ``min_seq_len``
-- in MiniMax H3 the padding tail -- go through the dense kernel.
"""
def all_dense() -> torch.Tensor:
return self.dense_impl.forward_varlen(
query,
key,
value,
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
cu_seqlens_host=cu_seqlens_host,
)
if cu_seqlens_host is None or not self._sparse_ready(query, key):
return all_dense()
segments = [
(start, stop)
for start, stop in zip(cu_seqlens_host[:-1], cu_seqlens_host[1:])
if stop > start
]
sparse_segments = {
(start, stop)
for start, stop in segments
if stop - start >= self.schedule.min_seq_len
}
if not sparse_segments:
return all_dense()
out = torch.empty_like(query)
# cu_seqlens covers every packed row in practice; a caller that leaves
# a tail outside the last document would otherwise read uninitialized
# memory back out.
if segments[-1][1] < query.shape[0]:
out[segments[-1][1] :].zero_()
for start, stop in segments:
# Deliberately not `.contiguous()`. After the Ulysses all-to-all,
# q/k/v are last-dim slices of one packed buffer, so they are
# strided; both the block-sparse kernel and SDPA permute them
# anyway, and forcing contiguity here measured as a wasted
# full-tensor copy (0.46 ms per call at S=37.7k on B200).
q_seg = query[start:stop].unsqueeze(0)
k_seg = key[start:stop].unsqueeze(0)
v_seg = value[start:stop].unsqueeze(0)
if (start, stop) in sparse_segments:
seg_out = self._sparse_attention(q_seg, k_seg, v_seg)
else:
seg_out = self._dense_segment(q_seg, k_seg, v_seg)
out[start:stop] = seg_out[0]
return out
def _dense_segment(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
) -> torch.Tensor:
"""Dense attention over one packed document, ``[1, S, H, D]``."""
return torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
is_causal=self.causal,
scale=self.softmax_scale,
).transpose(1, 2)
__all__ = [
"SubBlockSparseAttentionBackend",
"SubBlockSparseAttentionImpl",
"SubBlockSparseAttentionMetadata",
"SubBlockSparseAttentionMetadataBuilder",
"SubBlockSparseSchedule",
]
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+39
View File
@@ -297,6 +297,44 @@ class _VMOBAAttentionBackendResolver(_CudaAttentionBackendResolver):
raise ImportError("Video MoBA Attention backend is not installed. ") from e raise ImportError("Video MoBA Attention backend is not installed. ") from e
class _SubBlockSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
backend = AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
# The blk64 kernel is built `-gencode=arch=compute_100a,code=sm_100a`, which
# is arch-specific: 10.3 (B300 / GB300) and 12.x have no cubin. Its own guard
# only compares the major version, so it would accept 10.3 and fail later.
required_capability = (10, 0)
@classmethod
def resolve(cls, platform) -> str:
capability = platform.get_device_capability()
if capability is None or capability != cls.required_capability:
found = capability.as_version_str() if capability else "unknown"
raise ValueError(
"SubBlock sparse attention needs compute capability "
f"{'.'.join(map(str, cls.required_capability))} (B200 / GB200); "
f"this device reports {found}."
)
try:
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import ( # noqa: F401
load_bsa_attn_blk64_fwd,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import ( # noqa: F401
SubBlockSparseAttentionBackend,
)
# Importing the entry point catches a missing or broken FlashInfer;
# the CUDA extension itself is built lazily on the first call.
load_bsa_attn_blk64_fwd()
return "sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn.SubBlockSparseAttentionBackend"
except Exception as e:
logger.error("Failed to import SubBlock sparse attention: %s", str(e))
raise ImportError(
"SubBlock sparse attention needs FlashInfer with the blk64 "
"block-sparse kernel (flashinfer.cute_dsl.sparse.bsa_attn_blk64_fwd)."
) from e
class _FlashAttention2BackendResolver(_CudaAttentionBackendResolver): class _FlashAttention2BackendResolver(_CudaAttentionBackendResolver):
backend = AttentionBackendEnum.FA2 backend = AttentionBackendEnum.FA2
@@ -338,6 +376,7 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = {
_SparseVideoGen2AttentionBackendResolver, _SparseVideoGen2AttentionBackendResolver,
_SolAttnBackendResolver, _SolAttnBackendResolver,
_VMOBAAttentionBackendResolver, _VMOBAAttentionBackendResolver,
_SubBlockSparseAttentionBackendResolver,
_FlashAttention2BackendResolver, _FlashAttention2BackendResolver,
_FlashAttentionBackendResolver, _FlashAttentionBackendResolver,
) )
python/sglang/multimodal_gen/runtime/platforms/interface.py
+2
View File
@@ -44,6 +44,7 @@ class AttentionBackendEnum(enum.Enum):
BLOCK_SPARSE_ATTN = enum.auto() BLOCK_SPARSE_ATTN = enum.auto()
RAIN_FUSION_ATTN = enum.auto() RAIN_FUSION_ATTN = enum.auto()
SOL_ATTN = enum.auto() SOL_ATTN = enum.auto()
SUBBLOCK_SPARSE_ATTN = enum.auto()
NO_ATTENTION = enum.auto() NO_ATTENTION = enum.auto()
def __str__(self): def __str__(self):
@@ -62,6 +63,7 @@ class AttentionBackendEnum(enum.Enum):
AttentionBackendEnum.BLOCK_SPARSE_ATTN, AttentionBackendEnum.BLOCK_SPARSE_ATTN,
AttentionBackendEnum.RAIN_FUSION_ATTN, AttentionBackendEnum.RAIN_FUSION_ATTN,
AttentionBackendEnum.SOL_ATTN, AttentionBackendEnum.SOL_ATTN,
AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
} }
python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py
+357
View File
@@ -0,0 +1,357 @@
# SPDX-License-Identifier: Apache-2.0
"""SubBlock block-sparse attention backend.
The schedule tests are pure CPU. The numerical tests need an SM100 GPU with
FlashInfer's ``bsa_attn_blk64_fwd`` and are skipped otherwise.
The trick that makes the sparse kernel checkable against dense attention: at
``sparsity`` just above 0 every block is inside the budget, so the block-sparse
result must reproduce dense attention up to bf16 rounding. That covers the
routing indices, the tail block sizes, and the softmax scale in one assertion,
none of which an accuracy-only comparison at real sparsity would pin down.
"""
from __future__ import annotations
import unittest
from unittest.mock import patch
import torch
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse.router import (
_snap_up_to_8,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
SubBlockSparseAttentionBackend,
SubBlockSparseAttentionImpl,
SubBlockSparseSchedule,
_dit_layer_index,
)
HEAD_DIM = 128
NUM_HEADS = 4
def _sm100_available() -> bool:
if not torch.cuda.is_available():
return False
# Exactly 10.0: the kernel is built for sm_100a, and 10.3 has no cubin.
if torch.cuda.get_device_capability(0) != (10, 0):
return False
try:
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
load_bsa_attn_blk64_fwd,
)
load_bsa_attn_blk64_fwd()
except Exception:
return False
return True
requires_sm100 = unittest.skipUnless(
_sm100_available(), "needs SM100 and FlashInfer bsa_attn_blk64_fwd"
)
class _FakeServerArgs:
def __init__(self, config):
self.attention_backend_config = config
def _patch_schedule(config):
return patch(
"sglang.multimodal_gen.runtime.server_args.get_global_server_args",
return_value=_FakeServerArgs(config),
)
def _patch_step(step: int):
class _Ctx:
current_timestep = step
return patch(
"sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn.get_forward_context",
return_value=_Ctx(),
)
def _dense_reference(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float
) -> torch.Tensor:
"""q, k, v: [1, S, H, D]."""
return torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale=scale
).transpose(1, 2)
def _structured_qkv(
seq_len: int, device: torch.device, n_topic: int = 64, seed: int = 0
):
"""Attention-like q/k/v: ``[1, S, H, 128]`` bf16.
Random q/k produces a near-uniform attention map, and under a uniform map
no block-selection rule can work -- dropping blocks drops mass wherever you
cut. Real video attention concentrates on a limited region of keys, so each
query here prefers one narrow key topic (a couple of 64-token blocks), and
each topic carries its own value, so attending to the wrong blocks gives a
visibly wrong answer instead of the same blurred average.
"""
gen = torch.Generator(device=device).manual_seed(seed)
shape = (n_topic, NUM_HEADS, HEAD_DIM)
key_centers = torch.randn(shape, device=device, generator=gen)
value_centers = torch.randn(shape, device=device, generator=gen)
topic = torch.arange(seq_len, device=device) * n_topic // seq_len
def _noise(scale: float) -> torch.Tensor:
return scale * torch.randn(
seq_len, NUM_HEADS, HEAD_DIM, device=device, generator=gen
)
q = key_centers[topic] + _noise(0.3)
# the topic a query wants sits half a sequence away, so a rule that simply
# keeps the diagonal cannot pass
k = key_centers[(topic + n_topic // 2) % n_topic] + _noise(0.3)
v = value_centers[topic] + _noise(0.1)
return (
q[None].to(torch.bfloat16),
k[None].to(torch.bfloat16),
v[None].to(torch.bfloat16),
)
def _cosine(a: torch.Tensor, b: torch.Tensor) -> float:
return float(
torch.nn.functional.cosine_similarity(
a.float().flatten(), b.float().flatten(), dim=0
)
)
class TestSubBlockSparseSchedule(unittest.TestCase):
def test_dit_layer_index_only_matches_top_level_blocks(self):
self.assertEqual(_dit_layer_index("blocks.7.attn"), 7)
self.assertEqual(_dit_layer_index("blocks.0.attn"), 0)
self.assertIsNone(_dit_layer_index("token_refiner.blocks.1.attn"))
self.assertIsNone(_dit_layer_index(""))
self.assertIsNone(_dit_layer_index("blocks_extra.3.attn"))
def test_defaults_when_config_is_empty(self):
with _patch_schedule({}):
schedule = SubBlockSparseSchedule.from_server_args()
self.assertEqual(schedule.sparsity, 0.75)
self.assertEqual(schedule.skip_first_steps, 10)
# Depth is not protected by default; the early steps are. See the
# sweep recorded next to the constants.
self.assertEqual(schedule.skip_first_layers, 0)
self.assertEqual(schedule.n_k, 4)
self.assertEqual(schedule.n_q, 4)
def test_rejects_out_of_range_values(self):
for config in ({"sparsity": 1.0}, {"n_k": 3}, {"skip_first_steps": -1}):
with self.subTest(config=config), _patch_schedule(config):
with self.assertRaises(ValueError):
SubBlockSparseSchedule.from_server_args()
class TestBudgetGranularity(unittest.TestCase):
"""The kernel bills in groups of 8 blocks; the budget should collect them."""
def test_snaps_up_to_the_billed_count(self):
for topk, expected in ((148, 152), (118, 120), (1, 8), (0, 8)):
with self.subTest(topk=topk):
self.assertEqual(_snap_up_to_8(topk, 590), expected)
def test_never_exceeds_the_blocks_that_exist(self):
"""The cap wins over the granularity: 590 blocks means at most 590."""
self.assertEqual(_snap_up_to_8(586, 590), 590)
self.assertEqual(_snap_up_to_8(3, 5), 5)
class TestSubBlockSparseBackend(unittest.TestCase):
def test_the_advertised_builder_can_be_built(self):
"""`AttentionMetadataBuilder.__init__` is abstract; a builder that does
not override it makes `get_builder_cls()()` a TypeError."""
builder = SubBlockSparseAttentionBackend.get_builder_cls()()
builder.prepare()
metadata = builder.build(current_timestep=7)
self.assertIsInstance(
metadata, SubBlockSparseAttentionBackend.get_metadata_cls()
)
self.assertEqual(metadata.current_timestep, 7)
class TestSubBlockGating(unittest.TestCase):
"""The schedule must decide sparsity from the layer and the step alone."""
def _impl(self, prefix: str, **config) -> SubBlockSparseAttentionImpl:
with _patch_schedule(config), patch.object(
SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
):
return SubBlockSparseAttentionImpl(
num_heads=NUM_HEADS,
head_size=HEAD_DIM,
causal=False,
softmax_scale=HEAD_DIM**-0.5,
prefix=prefix,
)
def test_skip_first_layers_gates_the_bottom_of_the_stack(self):
for prefix, expected in (
("blocks.0.attn", False),
("blocks.1.attn", False),
("blocks.2.attn", True),
):
with self.subTest(prefix=prefix):
impl = self._impl(prefix, skip_first_layers=2)
self.assertEqual(impl.layer_enabled, expected)
def test_token_refiner_is_dense(self):
self.assertFalse(self._impl("token_refiner.blocks.0.attn").layer_enabled)
def test_head_dim_other_than_128_is_dense(self):
with _patch_schedule({}), patch.object(
SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
):
impl = SubBlockSparseAttentionImpl(
num_heads=NUM_HEADS,
head_size=64,
causal=False,
softmax_scale=64**-0.5,
prefix="blocks.9.attn",
)
self.assertFalse(impl.layer_enabled)
def test_first_steps_are_dense(self):
impl = self._impl("blocks.9.attn")
for step, expected in ((0, False), (9, False), (10, True), (49, True)):
with self.subTest(step=step), _patch_step(step):
self.assertEqual(impl._step_enabled(), expected)
def test_short_sequences_are_dense(self):
impl = self._impl("blocks.9.attn")
q = torch.empty(1, 1024, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16)
with _patch_step(20):
self.assertFalse(impl._sparse_ready(q, q))
def test_fp32_is_dense(self):
impl = self._impl("blocks.9.attn")
q = torch.empty(1, 8192, NUM_HEADS, HEAD_DIM, dtype=torch.float32)
with _patch_step(20):
self.assertFalse(impl._sparse_ready(q, q))
@requires_sm100
class TestSubBlockNumerics(unittest.TestCase):
seq_len = 8192
def _impl(self, **config) -> SubBlockSparseAttentionImpl:
with _patch_schedule(config):
return SubBlockSparseAttentionImpl(
num_heads=NUM_HEADS,
head_size=HEAD_DIM,
causal=False,
softmax_scale=HEAD_DIM**-0.5,
prefix="blocks.9.attn",
)
def test_full_budget_reproduces_dense(self):
device = torch.device("cuda")
q, k, v = _structured_qkv(self.seq_len, device)
impl = self._impl(sparsity=1e-6)
with _patch_step(20):
out = impl.forward(q, k, v, None)
ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
self.assertGreater(_cosine(out, ref), 0.999)
def test_ragged_tail_reproduces_dense(self):
"""A sequence that is not a multiple of the 64-token block."""
device = torch.device("cuda")
seq_len = self.seq_len + 37
q, k, v = _structured_qkv(seq_len, device)
impl = self._impl(sparsity=1e-6)
with _patch_step(20):
out = impl.forward(q, k, v, None)
ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
self.assertGreater(_cosine(out, ref), 0.999)
def test_routing_finds_the_blocks_that_carry_the_mass(self):
"""At 0.75 sparsity the router must keep the blocks that matter.
The random-budget control is the point of this test: with the same
number of blocks but chosen at random the output collapses, so a high
cosine here measures the routing, not a forgiving fixture.
"""
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse import (
SubBlockRouter,
load_bsa_attn_blk64_fwd,
)
device = torch.device("cuda")
q, k, v = _structured_qkv(self.seq_len, device)
ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
impl = self._impl(sparsity=0.75)
with _patch_step(20):
routed = impl.forward(q, k, v, None)
self.assertTrue(torch.isfinite(routed.float()).all())
self.assertGreater(_cosine(routed, ref), 0.99)
num_blocks = (self.seq_len + 63) // 64
topk = impl.router.route(q, k, sparsity=0.75, softmax_scale=HEAD_DIM**-0.5).topk
# A random permutation per row, not `randint`: sampling with replacement
# would leave the control holding duplicate blocks, so it would attend
# fewer distinct blocks than the router at the same budget, and the
# repeats would distort the softmax mass on top of that.
random_index = (
torch.rand(1, NUM_HEADS, num_blocks, num_blocks, device=device)
.argsort(dim=-1)[..., :topk]
.to(torch.int32)
)
random_out = load_bsa_attn_blk64_fwd()(
q,
k,
v,
random_index,
topk,
block_sizes=SubBlockRouter.block_sizes(self.seq_len, device),
q2k_block_nums=None,
softmax_scale=HEAD_DIM**-0.5,
)
random_out = random_out[0] if isinstance(random_out, tuple) else random_out
self.assertLess(_cosine(random_out, ref), 0.9)
def test_skipped_step_is_bitwise_dense(self):
device = torch.device("cuda")
q, k, v = _structured_qkv(self.seq_len, device)
impl = self._impl(sparsity=0.75, skip_first_steps=10)
with _patch_step(3):
skipped = impl.forward(q, k, v, None)
dense = impl.dense_impl.forward(q, k, v, None)
torch.testing.assert_close(skipped, dense, rtol=0, atol=0)
def test_varlen_routes_each_document(self):
"""Packed [real | padding] layout, exactly MiniMax H3's cu_seqlens."""
device = torch.device("cuda")
used, total = self.seq_len, self.seq_len + 512
q, k, v = _structured_qkv(total, device)
q, k, v = q[0], k[0], v[0]
cu_host = (0, used, total)
cu = torch.tensor(cu_host, dtype=torch.int32, device=device)
impl = self._impl(sparsity=1e-6)
with _patch_step(20):
out = impl.forward_varlen(
q, k, v, cu_seqlens=cu, max_seqlen=used, cu_seqlens_host=cu_host
)
for start, stop in ((0, used), (used, total)):
ref = _dense_reference(
q[start:stop][None],
k[start:stop][None],
v[start:stop][None],
HEAD_DIM**-0.5,
)[0]
self.assertGreater(_cosine(out[start:stop], ref), 0.999)
if __name__ == "__main__":
unittest.main()