企业官网建设流程全解析

SGLang SubBlock 稀疏注意力：为 MiniMax-H3 视频 DiT 提供免训练块稀疏加速

【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang

SubBlock（子块）稀疏注意力是 SGLang 多模态生成运行时（python/sglang/multimodal_gen）中为 MiniMax-H3 视频 Diffusion Transformer（DiT）专门设计的一种免训练块稀疏注意力后端：它不修改任何权重，仅用一个极廉价的估计器在注意力之前选出值得保留的 64-token 块，再把计划交给架构专属的块稀疏 FlashAttention 内核执行。本文以仓库内 subblock_sparse/README.md 为核心，结合 router.py、kernels.py 与 subblock_sparse_attn.py 的源码，完整讲解如何启动、如何调参、评分如何工作、底层如何分派到 SM90/SM100/SM120 三种内核，以及实测性能边界，让你能把这一方案直接复现到自己的 MiniMax-H3 视频推理上。

一、背景：为什么需要一个"子块"级别的稀疏方案

视频 DiT 的注意力序列极长（MiniMax-H3 t2va 单条媒体序列可达 3.7 万 token），注意力是去噪步中的主要开销。块稀疏 FlashAttention 把序列切成固定大小的块，只计算"值得计算"的块对，从而跳过大量计算。但传统代理分数——对每个 64×64 块取mean(Q_block) · mean(K_block)——有个致命弱点：把 64 个 key 平均成一个向量，恰好抹掉了"一个 query 到底想要哪些 key"所依赖的差异。

SubBlock 的思路是：在块内部再切子块，对每个子块对单独打分，再用 log-sum-exp 合并成块的分数。这个分数直接估计该块的未归一化 softmax 质量——即跳过这个块会丢失多少注意力质量，因此排序更可靠。

整个方案"训练免费"（training-free）：没有训练过程，不改变任何权重，路由器在注意力之前运行，把选中的 key 块索引以q2k_block_index张量的形式交给内核。从 router.py 的模块文档可以看到，SubBlock 最初是从独立仓库 vendored 进 SGLang 的，router.py与kernels.py此后已与上游分叉，演进出了本文所述的实现。

二、快速启动：推荐的完整命令行

README 给出了经过完整拼写的推荐配置，也就是所有下文测量数据的基准配置：

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}'

关键提醒 1：text_encoder覆盖不是可选项

--attention-backend会作用于每一个组件，而 Qwen3-VL 文本编码器只接受fa、torch_sdpa或sage_attn_3。若不提供--component-attention-backends text_encoder=fa覆盖，服务器会直接抛错、无法启动。

覆盖必须加在编码器上，而不是 DiT 上——transformer=subblock_sparse_attn看起来能工作，但实际上会静默地什么都不做，因为 H3 是在第一次 forward 时才惰性解析 DiT 的后端，此时已不在组件加载的上下文中。源码 subblock_sparse_attn.py 的模块文档明确记录了这一点。

关键提醒 2：SM120 平台使用text_encoder=torch_sdpa

在 SM120（如 GB300 一代的 Blackwell Ultra 后续架构）上，CUDA 平台为稠密注意力选择 Torch SDPA，且组件级后端请求会被严格校验，因此要改用--component-attention-backends text_encoder=torch_sdpa。

关键提醒 3：配置 JSON 的覆盖语义与 shell 转义

--attention-backend-config是可选的，且只覆盖它点名的键。例如只传'{"sparsity": 0.85}'就能单独把稀疏度调高，换取约 6% 的额外加速。但要注意：内联 JSON 会被shlex.split处理，shell 吞掉引号时配置会被破坏；README 建议改用文件路径传递配置，而不是内联 JSON。

三、运行前提与约束

后端会按架构选择内核，其共享约束如下表（摘自 README）：

在 DiT 内部，任何内核无法服务的调用——交叉注意力、token refiner、短于min_seq_len的序列、非 bf16 激活、head_dim ≠ 128——都会对该次调用回退到稠密注意力，因此不需要手工排除任何层。这一点在 subblock_sparse_attn.py 中由layer_enabled逻辑实现：不在 DiT 栈内的层（token refiner、交叉注意力）永远不会走稀疏路径。

在不支持的 GPU 上，这不是回退，而是启动期报错。解析器在加载所选内核之前严格接受计算能力 9.0、10.0 或 12.0，因此 B300 等不支持的能力值会在启动时失败，而不是在跑完十个稠密去噪步之后才失败。精确检查 10.0 是必须的：FlashInfer 的内核是为sm_100a构建的，没有前向兼容的 10.3 cubin。subblock_sparse_attn.py 中的_get_subblock_sparse_attention_runner就是这一校验的实现，且测试 test_subblock_sparse_attention.py 明确断言 10.3 设备会抛出RuntimeError。

四、评分原理：从均值池化到子块 log-sum-exp

4.1 分数公式

64×64 块的传统代理分数是mean(Q_block) · mean(K_block)。SubBlock 把每个块切成子块——key 侧n_k个、query 侧n_q个——对每一对子块打分，再用 log-sum-exp 合并：

score(i, j) = log Σ_{a,b} exp( mean(Q_{i,a}) · mean(K_{j,b}) · softmax_scale )

这个分数直接估计块的未归一化 softmax 质量——也就是跳过该块会损失多少注意力质量。router.py 的模块注释给出了更严格的表述：它是对sum_{r in i, c in j} exp(q_r . k_c * scale)的直接估计。

4.2 只拆 query 侧反而更差

关键经验来自实测召回率表（router.py，在 MiniMax-H3 DiT 注意力的 567 个样本上测得，块稀疏度 0.9）：

（1 u = 一个[S/128, 128] × [128, S/128]GEMM，等于其所门控的稠密注意力的 1/16384。）

单独拆分 query 侧比不拆更差：一个块的质量对其 query 行求和，只有一个 key 向量可打分时，query 细节会被平均掉。两侧一起拆分则完全不同——log-sum-exp 现在遍历 query-key 子块对，"这个 query 块的某部分想要那个 key 块的某部分"这一信号能扛过平均化。这是该估计器家族中唯一端到端站得住脚的行，n_q = n_k = 4因此成为默认值。

4.3 实现中的数值细节

• 分数保持在 float32。router.py 注释说明：bf16 只有 8 个尾数位，很多块会在阈值处恰好打平，而融合选择器按列索引打破平局——这会系统性偏向早期的 key 块（视频的某一个区域），实测在 S=96k 时相对 L2 增加 3.9%。
• 内核内部用 exp2/log2（硬件指令），softmax_scale * log2(e)在池化时折叠进 Q，输出再乘以 ln 2 回到自然对数单位，与参考实现逐位一致而不只是排序一致，见 kernels.py。
• 尾块处理：池化直接作用在池化张量上处理参差不齐的尾部，先把 q/k 补齐到 G×64 会复制整个 300+ MB 激活。子单元越过最后一个真实 token 的部分池化为零，*_valid标量让分数内核强制其为 -inf，防止 exp(0)=1 的项既抬高分数又压平排序差异，见 router.py。

五、配置参数详解

这些默认值在源码中以DEFAULT_*常量定义（subblock_sparse_attn.py），并附有选择理由的测量记录。SubBlockSparseSchedule.from_server_args()负责从attention_backend_config读取并校验：sparsity必须在 [0, 1) 内、n_k/n_q必须为 1/2/4/8、skip_first_*必须非负，非法值在启动时抛ValueError（subblock_sparse_attn.py）。

5.1sparsity是上界，不是精确值

内核会把每个 query 行的块数向上填充到 8 的倍数，用幻影槽位（phantom slot）补齐后再在 softmax 中掩掉，所以 148 块的代价与 152 块完全相同。路由器的_snap_up_to_8就利用这一点主动取满：router.py 实测 152 块耗时 16.055 s vs 148 块 16.061 s、120 块 15.490 s vs 118 块 15.496 s——完全在噪声内。在 590 块时请求 0.75 实际交付 0.7424，启动日志会报告实际保留的块数。sparsity是速度调节杆，也是大多数用户唯一需要动的旋钮。

5.2n_k/n_q买的是评分精度，不是速度

它们决定一个块在打分前被切得多细：n_k=4意味着四个 16-token key 子块，块的分数是对全部n_q * n_k个子块对的 log-sum-exp。调高它们会锐化"哪些块承载质量"的估计，代价是n_q * n_k倍的分数矩阵——在默认值下只占去噪时间的 0.5%，所以成本不是约束。要同时调高n_q和n_k：单独拆 query 侧比不拆更差（见 4.2）。

5.3 两个调度截止点刻意不对称

• skip_first_steps：开头的去噪 forward 保持稠密。这些步稳定样本的布局，被近似时会明显重构图（re-frame）。把它从 10 降到 5，余弦相似度相对稠密渲染减半；源码注释记录了更细的扫描：5 → 0 时样本与稠密几乎不相关（0.558 → 0.310 的余弦，两段 clip 上），而稠密时间比从 1.20x 升到 1.30x。不要把skip_first_steps调低而不看输出。
• skip_first_layers：深度没有这种表现，所以默认 0，每个 DiT 层都稀疏。把层截止从 2 降到 0 只损失 0.0013 的余弦（在 0.02 的 run-to-run 噪声地板内），却能换回约 1% 时间。

5.4min_seq_len是下限，不是调参旋钮

低于它，整个调用走稠密；在打包的 varlen 批次中，这个测试是**按文档（per document）**进行的，所以 H3 的填充尾走稠密，而 37.7k token 的媒体文档被路由。短序列上有两件事会崩坏：其一，路由器是四个固定的 Triton launch，而对注意力成本随 S² 下降，开销很快就回不了本；其二，预算变粗——4096 个 key 只有 64 块，而 1024 个 key 时 8 的倍数地板已经要保留其中一半。4096 远低于任何真实视频序列、又远高于两种效应咬合的位置，它是基于这一推理而非阈值扫描选出来的。

六、源码实现：三个文件的职责划分

后端由三个文件组成：

6.1 路由器：三个融合 Triton 内核

SubBlockRouter.route()的工作流是：scores()先对 Q/K 做子块池化并算出[B, H, Gq, Gk]的块分数，再用_snap_up_to_8计算预算并执行一次 top-k 选择，产出RoutingPlan（含index、topk、num_blocks，且提供density属性）。三个融合内核分别是：

1. fused_pool（kernels.py）：[B, S, H, D] → [B*H, n_cells, D]的掩码均值，激活只读一次、直接以 bf16 写出。单 warp 而非四 warp——瓦片只有 SUB×128，额外 warp 买不到并行度，反而增加调度开销（S=96k 时实测 0.272 → 0.074 ms，1.27 → 4.7 TB/s）。
2. fused_scores（kernels.py）：bf16 GEMM + 分段 log-sum-exp 单遍完成，[B, H, Gq*n_q, Gk*n_k]中间量永不落内存。对照：PyTorch 参考实现会物化 254 MB 的 fp32 中间量（S=96k、n_k=4），其 GEMM 只占 4% 时间，其余全是该张量的 HBM 往返。归约先沿 NK 轴再沿 NQ 轴，log-sum-exp 的结合律保证这与一次性的 NQ×NK 归约结果相同，且两个阶段都落在寄存器中连续的轴上。
3. fused_topk（kernels.py）：单遍全局内存的"足够精确"逐行 top-K。一行分数在 bf16 下只有 3 KB（S=96k），一次载入寄存器后全部片上完成；torch.topk则要对整个矩阵做多遍扫描，这正是它占路由器一半开销的原因。阈值搜索在计数上做插值（割线法）而非二分区间——计数-阈值曲线就是该行的经验 CDF，割线步远比二分步接近答案。迭代次数由topk_iters按log2(G/k)自适应（16/24/32 三档），保证各稀疏度区间下选择集与精确 top-K 的差异都低于 0.2%。

6.2 后端：调度、门控与架构分派

SubBlockSparseAttentionBackend是标准的AttentionBackend实现，get_supported_head_sizes()只返回[128]。核心逻辑在SubBlockSparseAttentionImpl：

• 层启用判定（构造时）：prefix用^blocks\.(\d+)\.正则匹配 DiT 层索引；层在 DiT 栈内、层号 ≥skip_first_layers、head_dim == 128 且sparsity > 0才启用稀疏（subblock_sparse_attn.py）。
• 步启用判定（每次调用）：current_timestep >= skip_first_steps，步号来自 forward context（subblock_sparse_attn.py）。
• 稠密回退：_sparse_ready()要求layer_enabled、步启用、bf16、k.shape[-3] >= min_seq_len且非因果；任一不满足即走dense_impl（构造时选定 FA 后端，subblock_sparse_attn.py）。因此没有任何层需要手工排除。
• 内核分派：_get_subblock_sparse_attention_runner按设备计算能力一次性解析（带 lru_cache）：(9,0) → SM90 CuTe-DSL 内核（懒加载以避免把庞大的 CuTe 依赖树拖进 SM100 路径）；(10,0) → FlashInferbsa_attn_blk64_fwd；(12,0) → FlashInferbsa_attn_sm120_blk64_fwd；其他能力值抛RuntimeError。
• 索引顺序差异：SM90 要求每个活跃稀疏前缀按升序排序（plan.index.sort(dim=-1)），SM100/SM120 接受路由器的原始顺序（subblock_sparse_attn.py）。测试 test_subblock_sparse_attention.py 对三种 runner 分别验证了这一行为，并确认混合稀疏调用只发一次异构 BSA 调用。
• varlen 打包路径：forward_varlen按cu_seqlens把打包行拆成文档，短于min_seq_len的文档（H3 的填充尾）走稠密段路径，媒体文档走稀疏；分段切分时刻意不调用.contiguous()——Ulysses all-to-all 之后 q/k/v 是打包缓冲区的末维切片（strided 视图），强制连续化实测是浪费一次整张量拷贝（S=37.7k 时每次调用 0.46 ms）。
• FlashInfer 加载的健壮性：load_bsa_attn_blk64_fwd尝试从flashinfer.cute_dsl.sparse导入；若失败（该包还拖入 blk128 后端，硬依赖quack且追踪易碎的cutlass.cuteAPI），则退而在合成父包下直接加载bsa_attn_blk64.py——同一份文件、同一个内核，绕开无关依赖（router.py）。

6.3 已试过并被否决的变体

router.py 记录了完整的试错史，这本身就构成对默认配置的背书：

• 先把每个 query 子块转成对 key 块的分布再打分：在 1092 个真实样本上两个代理指标都更好（块质量召回 .6741 → .6779，重建输出相对 L2 .2043 → .1982，配对 t 分别为 +8.0 和 -5.7），但在 15 个提示的像素端全面更差（0.107 的余弦损失，配对 t = -6.4）。结论：单层输出误差甚至不能正确排序这些估计器，只有端到端渲染才能——这解释了为何默认选择如此保守。
• 结构性的块保留（attention sink，或强制对角线 j == i）：在 200 个真实 H3 注意力单元上测得对角线只改变相对 L2 0.2%，sink 只在 DiT 2-32 层有帮助且未能活到像素端，因此刻意缺席。
• 每头预算（per-head budget）：优于任何估计器升级，但需要设置每头分配的规则，当前管线没有产出，因此"值得尝试但未暴露"。

七、实测性能

以下数据来自 README 的 "Measured" 章节：MiniMax-H3 t2va，1344×768 / 5 s / 50 步，8× B200，Ulysses-8，bf16，使用发布的默认配置（每个 query 块保留 590 中的 152 个 key 块）。所有分支在同一节点同一会话中测得，丢弃冷采样，单分支内波动低于 0.07 s。

（PR #4397 把内核内部 Q/K/V 瓦片布局的重建从三遍改为一遍，逐位一致且非必需，单独值 1.070x。）

• 稀疏度是速度杠杆且会饱和：0.75 → 1.136x，0.80 → 1.178x，0.85 → 1.211x。把预算在 0.75 基础上再砍 40% 只买回 6%，因为在 37.7k token 时注意力已不再是步的主体，且 0.85 在三个档位中与稠密的余弦最差。
• n_k是质量旋钮而非速度旋钮：它在其整个取值范围内只移动去噪时间 0.3%。
• 加速比受序列长度约束，而非方法本身：同一配置在 37.7k token 测 1.13x、52k 测 1.20x、96k 测 1.47x——后端只碰注意力，而注意力在 DiT 中的占比随 S 增长。把 1.2x 当作 768p/5 s 的数字，而不是天花板。
• 序列并行度也有同样的效应（它决定每张 GPU 持有多少序列）：4× B200 上 Ulysses-4 的同一配置，去噪 1.168x、端到端 1.138x，对比 Ulysses-8 下去噪 1.138x。
• 峰值内存不变（99,356 vs 99,358 MiB/GPU）：块稀疏省计算、不省激活；[B,H,Gq,Gk]分数矩阵在 S=37.7k 时约 20 MB。绝对时间是节点相关的，只有同一会话内的比率可比。

八、测试：如何把稀疏内核钉在稠密结果上

单元测试位于 test/registered/cpu/test_subblock_sparse_attention.py。其中有一个巧妙的验证手法：把稀疏度设在略高于 0 处运行——此时每个块都在预算内，结果必须复现稠密注意力直至 bf16 舍入误差，一条断言同时钉死路由索引、参差尾块大小和 softmax 尺度。

该测试文件还覆盖了：每设备一次的内核分派解析（_get_subblock_sparse_attention_runner缓存）、三种架构 runner 的转发行为（含block_sizes与q2k_block_nums透传）、10.3 能力值拒绝、ring_degree与 SubBlock 的准入冲突、varlen 打包下缺失 query 掩码的三态行为（稠密回退允许、短段允许、真跑稀疏时报错）、以及 ref2va 场景下"参考视频与目标视频稀疏、图像与 padding 保持稠密"的视频位置掩码语义。

九、小结与实操建议

SubBlock 稀疏注意力在 SGLang 中的定位非常清晰：它是 MiniMax-H3 视频 DiT 的免训练、零权重改动加速方案，把"哪些 64-token 块值得算"交给一个子块级 log-sum-exp 估计器，把"算"交给 SM90/SM100/SM120 三套架构专属的块稀疏内核。上手时记住四条：

1. 永远带上--component-attention-backends text_encoder=fa（SM120 上换torch_sdpa），否则服务器起不来；
2. 默认配置就是推荐配置，日常只动sparsity一个旋钮（它是上界，日志会告诉你实际值）；
3. 不要单独调n_q，skip_first_steps调低必须看渲染输出；
4. 不在 9.0/10.0/12.0 能力值的 GPU 上会启动即报错，这是设计行为而非故障。

【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang