Phase 2.5 of RFC #29630 旨在将散落在 sglang 包各处的约 280 个 Triton 内核及 CuTe DSL / TileLang 内核迁移到 sglang.kernels 包下的三个规范位置。本 PR 作为第 6/7 轮，覆盖线性注意力家族、MiniMax 稀疏操作和 multimodal_gen 中的遗留内核，包括唯一一个位于规范内核树之外的 .cu 源文件。
值得仔细阅读。该 PR 展示了大规模代码组织重构的最佳实践：使用逐字节搬迁保持历史可追溯性、统一注册入口、系统更新所有导入引用。对于关注代码架构和内核管理的工程师有参考价值。
无实质性讨论。PR 由 BBuf 自动化生成并直接合并，仅触发 CI 重跑。
lightning_attn.py、seg_la.py 以及 gdn_blackwell/、kda_blackwell/ 子包（CuTe DSL）从 sglang.srt.layers.attention.linear 整体移到 sglang.kernels.ops.attention.linear，保持逐字节相同。minimax_sparse_ops/（common/decode/prefill 共 12 个 Triton 内核）从 sglang.srt.layers.attention 移到 sglang.kernels.ops.attention.minimax_sparse。multimodal_gen/csrc/render/ 下的 .cu 文件（hunyuan3d_rasterizer）和 cpp_extension 加载器移到 sglang.kernels.ops.diffusion.render。sparse_linear_attn.py 中的 get_block_map、mean_pool、compress_kernel、_attn_fwd 四个函数提取到新文件 sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py，原文件改为从新位置导入。ops/attention/__init__.py 中添加 linear 和 minimax_sparse 子包的导入，使新迁移的内核可通过 KernelSpec 清单发现。| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py |
扩散注意力 | modified | 8.42 |
python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py |
注意力内核 | added | 7.24 |
python/sglang/srt/layers/attention/linear/lightning_backend.py |
线性注意力 | modified | 5.19 |
benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py |
KDA 基准 | modified | 5.0 |
python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py |
3D 渲染 | modified | 4.79 |
python/sglang/kernels/ops/attention/__init__.py |
内核注册 | modified | 4.11 |
python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py
core-logic
核心变更：移除了 4 个 Triton 内核定义（get_block_map, mean_pool, compress_kernel, _attn_fwd），改为从新的 sglang.kernels 子包导入。这是多模态扩散模型中 Sparse Linear Attention 的后端文件，也是本 PR 中唯一有实质内容删除的文件。
"""Sparse Linear Attention backend - 内核已搬迁至 sglang.kernels."""
import torch
import triton
# 注意：不再直接 import triton.language as tl，因为内核已移出
from sglang.kernels.ops.diffusion.sparse_linear_attn_kernels import (
_attn_fwd,
get_block_map,
)
# 原文件中定义的 get_block_map、mean_pool、compress_kernel、_attn_fwd
# 已整块搬迁至 sparse_linear_attn_kernels.py，此处仅保留导入
# 后端类保持不变，注意力计算逻辑通过导入的内核函数完成
class SparseLinearAttentionBackend(AttentionBackend):
# ... 类实现细节未改动 ...
pass
python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py
infrastructure
新增目标文件，包含所有从原 sparse_linear_attn.py 提取的内核函数。作为迁移的目的地，定义了 get_block_map、mean_pool、compress_kernel 和 _attn_fwd。
"""扩散模型稀疏线性注意力内核（从 sglang.multimodal_gen 迁移而来）。"""
import torch
import triton
import triton.language as tl
def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
# smooth-k 技术（来自 SageAttention）：对 k 减去均值
arg_k = k - torch.mean(k, dim=-2, keepdim=True)
# 对 Q 和 K 进行块均衡池化
pooled_qblocks = mean_pool(q, BLKQ)
pooled_kblocks = mean_pool(arg_k, BLKK)
pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
K = pooled_score.shape[-1]
topk = min(K, int(topk_ratio * K))
lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
sparse_map.scatter_(-1, lut, 1)
return sparse_map, lut, topk
def mean_pool(x, BLK):
x = x.contiguous()
B, H, L, D = x.shape
L_BLOCKS = (L + BLK - 1) // BLK
x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
grid = (L_BLOCKS, B * H)
compress_kernel[grid](x, x_mean, L, D, BLK)
return x_mean
@triton.jit
def compress_kernel(
X, XM, L: tl.constexpr, D: tl.constexpr, BLOCK_L: tl.constexpr,
):
idx_l = tl.program_id(0)
idx_bh = tl.program_id(1)
offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
offs_d = tl.arange(0, D)
x_offset = idx_bh * L * D
xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
x = tl.load(
X + x_offset + offs_l[:, None] * D + offs_d[None, :],
mask=offs_l[:, None] < L,
)
nx = min(BLOCK_L, L - idx_l * BLOCK_L)
x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
tl.store(
XM + xm_offset + idx_l * D + offs_d,
x_mean.to(XM.dtype.element_ty),
)
# _attn_fwd 等其他内核在此文件后续定义，与原实现逐字节相同
当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。
风险较低。所有内核移动均为逐字节搬迁（git rename R100 或仅导入路径差异），不改变任何执行逻辑。主要风险包括：(1) 导入路径重写遗漏可能导致运行时错误，但已通过仓库范围的 py_compile 和测试套件验证；(2) 与同系列其他 PR 存在 init.py 追加冲突，通过多次合并 main 解决。整体回归风险很小。
影响 44 个文件的导入路径，涉及 sglang.srt、sglang.multimodal_gen 等模块。对最终用户透明，所有功能保持不变。为后续内核版本管理奠定基础。
参与讨论