企业官网建设流程全解析

从零为 sgl-kernel 添加 AOT CUDA/C++ 内核：完整教程（含测试与基准测试）

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

本教程面向需要在 sgl-kernel（即本仓库python/sglang/kernels/aot/下的 AOT 内核库，Python 导入路径为sgl_kernel）中新增"重量级"预编译 CUDA/C++ 内核的开发者。文章以新增一个逐元素缩放算子scale(x, factor) = x * factor为完整示例，覆盖 C++ 实现、Torch 算子注册、CMake 构建、Python API 暴露、pytest 测试与 Triton 基准测试的全流程。读完本文你将能独立判断"何时该走 JIT 路径、何时该走 AOT 路径"，并把一个新算子完整、规范地合入sgl-kernel。

前置背景：什么是 AOT sgl-kernel，与 JIT 内核如何取舍

在python/sglang/kernels/aot/下维护着sglang-kernel（历史名称sgl-kernel）内核库，它面向 LLM 推理引擎提供优化的计算原语。源码树位于 python/sglang/kernels/aot，对应的 Python 导入名是sgl_kernel。与通过 Triton 在运行时即时编译的 JIT 内核不同，AOT 内核是随 wheel 包一起编译、经由 PyTorch 扩展机制注册（torch.ops.sgl_kernel.*）的 CUDA/C++ 算子，天然适用于依赖 CUTLASS 等大型 C++ 工程的"重量级"实现，且构建一次、处处加载。

仓库中同时存在轻量内核的默认路径python/sglang/kernels/jit（配套的.claude/skills/add-jit-kernel/SKILL.md技能文档），因此社区贡献新内核时首先需要遵循两条黄金法则（must follow）：

1. 优先选择 python/sglang/kernels/jit：当内核不依赖CUTLASS 或其他大型 C++ 工程时，JIT 是默认路径，适合迭代快速的轻量内核。
2. 优先选择 sgl-kernel（AOT）：当内核依赖CUTLASS 或其他大型 C++ 工程，或希望纳入 AOT wheel 与 torch op 注册流程时。
3. 例外情况：如果依赖的是flashinfer，或经由flashinfer已经提供的 CUTLASS，内核仍可作为jit_kernel实现。

此外，每一个新内核都必须配套交付两项产出：

• 测试（pytest）
• 基准测试脚本（triton.testing）

仓库集成地图：新增 AOT 内核会触及的文件

以下文件/区域是新内核合入时通常要改动的全部落点，也是后文每个 Step 的索引：

• 实现：python/sglang/kernels/aot/csrc/elementwise/scale.cu（按类别选择正确的子目录）
• 公开声明：python/sglang/kernels/aot/include/sgl_kernel_ops.h
• Torch 扩展注册：python/sglang/kernels/aot/csrc/common_extension.cc
• 构建：python/sglang/kernels/aot/CMakeLists.txt（set(SOURCES ...)）
• Python API：python/sglang/kernels/aot/python/sgl_kernel/与python/sglang/kernels/aot/python/sgl_kernel/__init__.py
• 测试：python/sglang/kernels/aot/tests/test_scale.py
• 基准测试：python/sglang/kernels/aot/benchmark/bench_scale.py

从源码结构看，python/sglang/kernels/aot/csrc 下按功能域组织子目录：allreduce/、attention/、elementwise/、gemm/、moe/、mamba/、grammar/、quantization/、speculative/等，新增算子应归入语义最贴切的子目录。

Step 1：在csrc/中实现 CUDA 内核与 launch 封装

先按算子类别选择正确子目录：逐元素算子放csrc/elementwise/，其余如csrc/gemm/、csrc/attention/、csrc/moe/各归其位。本示例在python/sglang/kernels/aot/csrc/elementwise/scale.cu中新增一个简单的逐元素缩放算子：

#include <ATen/cuda/CUDAContext.h> #include <c10/cuda/CUDAGuard.h> #include <torch/all.h> #include "utils.h" // DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 // scale_kernel: out[i] = input[i] * factor // Supports float, half (__half), __nv_bfloat16 via template T template <typename T> __global__ void scale_kernel(T* __restrict__ out, const T* __restrict__ input, float factor, int64_t n) { int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { out[idx] = static_cast<T>(static_cast<float>(input[idx]) * factor); } } void scale(at::Tensor& out, const at::Tensor& input, double factor) { TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor"); TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); TORCH_CHECK(out.sizes() == input.sizes(), "out and input must have the same shape"); TORCH_CHECK(out.scalar_type() == input.scalar_type(), "out and input must have the same dtype"); const int64_t n = input.numel(); const int threads = 256; const int blocks = (n + threads - 1) / threads; const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); // Dispatches over float, float16, bfloat16 DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] { scale_kernel<c_type><<<blocks, threads, 0, stream>>>( static_cast<c_type*>(out.data_ptr()), static_cast<const c_type*>(input.data_ptr()), static_cast<float>(factor), n); cudaError_t status = cudaGetLastError(); TORCH_CHECK(status == cudaSuccess, "scale_kernel launch failed: ", cudaGetErrorString(status)); return true; }); }

关键要点与源码佐证

• 接口与校验：host 侧函数接收at::Tensor，用TORCH_CHECK完成设备/连续性/shape/dtype 前置校验，获取当前 CUDA 流用at::cuda::getCurrentCUDAStream()。Python 包装层保持极薄，shape、dtype、device 校验尽量放在紧邻 launch 的 C++ 代码中完成。
• dtype 分派宏：DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16覆盖float（FP32）、half（FP16）、__nv_bfloat16（BF16），其实现定义于 python/sglang/kernels/aot/include/utils.h：按at::ScalarTypeswitch 到float、_DISPATCH_CASE_F16、_DISPATCH_CASE_BF16，其余类型走TORCH_CHECK(false, ...)报错。本教程示例算子支持 FP16（torch.float16）、BF16（torch.bfloat16）、FP32（torch.float32）三种 dtype，正是由该宏驱动的模板实例化完成的。
• 错误检查：每次 kernel launch 后都要取cudaGetLastError()并以TORCH_CHECK兜底，避免异步错误被延迟暴露。
• 架构约束：若内核仅在某些架构上可用，应在 host 侧用TORCH_CHECK强制约束，并在测试中通过 skip 逻辑跳过不支持的架构（仓库惯例见下文 Step 6 的@pytest.mark.skipif用法）。

作为真实代码参照，逐元素激活相关算子的实现位于 python/sglang/kernels/aot/csrc/elementwise/activation.cu，其在 launch 前同样执行at::cuda::getCurrentCUDAStream()、OptionalCUDAGuard，并用同一分派宏处理 FP16/BF16/FP32。

Step 2：在include/sgl_kernel_ops.h添加 C++ 声明

编辑python/sglang/kernels/aot/include/sgl_kernel_ops.h，在 elementwise 区块（源码中以/* ... */注释分区，如既有注释* From csrc/elementwise）内加入：

void scale(at::Tensor& out, const at::Tensor& input, double factor);

该头文件是全部 AOT 算子 C++ 接口的"总目录"，python/sglang/kernels/aot/include/sgl_kernel_ops.h 中的每个函数声明都按来源子目录组织并有明确注释，新声明加入对应分区即可保持可维护性。

Step 3：在csrc/common_extension.cc注册算子

编辑python/sglang/kernels/aot/csrc/common_extension.cc，在TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)块内新增 schema 定义与设备实现绑定：

// From csrc/elementwise m.def("scale(Tensor! out, Tensor input, float factor) -> ()"); m.impl("scale", torch::kCUDA, &scale);

关键要点

• Tensor!语义：Tensor!表示 in-place / 可变的输出参数，这保证了调用签名与torch.compile的可理解性。
• schema 的重要性：schema 对torch.compile和一致的调用签名至关重要；仓库 README 明确要求以m.def（带 schema 用于 torch.compile）+m.impl（设备绑定）的方式注册扩展。
• scalar 类型约定：torch schema 中按 PyTorch 标量类型书写（此处为float），但 C++ launcher 的函数签名仍需对接受 scalar 实参使用double（torch::Library的类型映射约定，与 Python 侧int/float到int64_t/double的映射一致）。若 C++ 实现层使用了int/float这类第三方库原生类型，可用 python/sglang/kernels/aot/include/sgl_kernel_torch_shim.h 中的make_pytorch_shim自动做类型转换。

真实注册示例可以在 common_extension.cc 中看到silu_and_mul的注册形式：m.def("silu_and_mul(Tensor! out, Tensor input) -> ()");后接m.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);。

Step 4：把新源文件加入CMakeLists.txt

编辑python/sglang/kernels/aot/CMakeLists.txt，在set(SOURCES ...)列表中加入：

csrc/elementwise/scale.cu

关键要点

• 字母序要求：该文件在set(SOURCES ...)上方显式注明NOTE: Please sort the filenames alphabetically，新增条目必须按字母序插入（见 CMakeLists.txt 中activation.cu、concat_mla.cu、copy.cu、dsv4_norm_rope.cu等已按序排列的条目）。
• 架构约束落地：若内核有架构限制，需在测试与基准脚本中通过 skip 逻辑体现。
• 遗漏后果：.cu文件若未进入SOURCES，链接期会出现符号未定义（undefined symbol）错误。

Step 5：在python/sgl_kernel/下暴露 Python API

优先沿用现有模块组织方式。对逐元素内核，惯例是：

• 在python/sglang/kernels/aot/python/sgl_kernel/elementwise.py中实现 Python 包装；
• 然后在python/sglang/kernels/aot/python/sgl_kernel/__init__.py中 re-export。

例如在python/sglang/kernels/aot/python/sgl_kernel/elementwise.py中新增：

import torch def scale( input: torch.Tensor, factor: float, out: torch.Tensor | None = None, ) -> torch.Tensor: """ Element-wise scale: out = input * factor. Supported dtypes: torch.float16, torch.bfloat16, torch.float32. Parameters ---------- input : CUDA input tensor factor : scale factor (float) out : optional pre-allocated CUDA output tensor (same shape/dtype as input) """ if out is None: out = torch.empty_like(input) torch.ops.sgl_kernel.scale.default(out, input, factor) return out

随后参照现有内核的导入风格，把scale加入 python/sgl_kernel/init.py 的 re-export。真实源码中该文件通过from sgl_kernel.elementwise import (...)显式列名导入（例如concat_mla_k、copy_to_gpu_no_ce、rmsnorm等），而elementwise.py内部则通过torch.ops.sgl_kernel.<name>.default(...)调用底层注册算子，例如torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)。

Step 6：编写 pytest 测试（必需）

创建python/sglang/kernels/aot/tests/test_scale.py：

import pytest import torch import sgl_kernel @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("size", [128, 1024, 4096, 65536]) @pytest.mark.parametrize("factor", [0.5, 1.0, 2.0]) def test_scale_correctness(dtype, size, factor): input = torch.randn(size, dtype=dtype, device="cuda") out = torch.empty_like(input) result = sgl_kernel.scale(input, factor, out=out) assert result is out expected = input * factor rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-2) torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) def test_scale_shape_mismatch(): input = torch.randn(128, dtype=torch.float16, device="cuda") out = torch.empty(256, dtype=torch.float16, device="cuda") with pytest.raises(RuntimeError, match="same shape"): sgl_kernel.scale(input, 2.0, out=out) def test_scale_cpu_input(): input = torch.randn(128, dtype=torch.float16) # CPU out = torch.empty_like(input) with pytest.raises(RuntimeError, match="CUDA"): sgl_kernel.scale(input, 2.0, out=out) if __name__ == "__main__": import sys sys.exit(pytest.main([__file__, "-q"]))

测试约定（与仓库既有测试一致）：

• 测试统一放在 python/sglang/kernels/aot/tests 目录下（目录内含conftest.py、utils.py及test_activation.py、test_norm.py、test_copy.py、test_topk.py等大量既有用例可参考）；
• 若某用例需要按环境/架构跳过，使用@pytest.mark.skipif(condition, reason="...")，例如仓库惯例中的架构能力判断（如 Nvfp4 需要 compute capability >= 10）；
• 正确性用例同时覆盖"结果正确"与"异常路径"，通过pytest.raises验证 host 侧TORCH_CHECK抛出的错误信息。

Step 7：添加 Triton 基准测试（必需）

创建python/sglang/kernels/aot/benchmark/bench_scale.py：

import itertools import torch import triton import triton.testing import sgl_kernel from sglang.utils import is_in_ci IS_CI = is_in_ci() dtypes = [torch.float16] if IS_CI else [torch.float16, torch.bfloat16, torch.float32] sizes = [4096] if IS_CI else [2**n for n in range(10, 20)] # 1K … 512K factors = [2.0] configs = list(itertools.product(dtypes, sizes)) def torch_scale(input: torch.Tensor, factor: float) -> torch.Tensor: return input * factor @triton.testing.perf_report( triton.testing.Benchmark( x_names=["dtype", "size"], x_vals=configs, line_arg="provider", line_vals=["sglang", "torch"], line_names=["SGL Kernel", "PyTorch"], styles=[("green", "-"), ("red", "--")], ylabel="µs (median)", plot_name="scale-performance", args={}, ) ) def benchmark(dtype, size, provider): input = torch.randn(size, dtype=dtype, device="cuda") out = torch.empty_like(input) factor = 2.0 if provider == "sglang": fn = lambda: sgl_kernel.scale(input, factor, out=out) else: fn = lambda: torch_scale(input, factor) ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( fn, quantiles=[0.5, 0.2, 0.8] ) return 1000 * ms, 1000 * max_ms, 1000 * min_ms if __name__ == "__main__": benchmark.run(print_data=True)

基准测试约定与仓库补充说明：

• 基准脚本统一放在 python/sglang/kernels/aot/benchmark 目录下（如bench_activation.py、bench_rmsnorm.py等），文件名遵循bench_*.py命名；
• 仓库 README 建议优先使用triton.testing.do_bench_cudagraph进行内核基准：相比do_bench，它能降低 CPU 开销对内核性能测量精度的影响、把 PDL（Programmatic Dependent Launch）效应计入单个内核结果，并在支持 PDL 的架构（SM >= 90）上给出更贴近真实的性能数据；
• 用sglang.utils.is_in_ci()（定义于 python/sglang/utils.py）在 CI 场景收敛配置矩阵，缩短运行时间；
• 展示维度通常为dtype × size，以 PyTorch 原生实现作为 baseline 做中位延迟对比。

Step 8：构建

在python/sglang/kernels/aot目录下执行：

cd python/sglang/kernels/aot make build -j16

如需限制宿主机资源占用：

cd python/sglang/kernels/aot make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"

构建相关说明（与 README 及 Makefile 一致）：

• make build默认占用全部可用 CPU 核；可通过MAX_JOBS控制 make 与 CMake 并行度，通过CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"额外限制 NVCC 内部线程数，从而降低 CPU 占用与峰值内存；
• Makefile 还提供check-deps/install-deps（安装scikit-build-core、isort、black）、install（pip install -e . --no-build-isolation开发模式安装）、rebuild（clean 后重建）、test（跑全部测试）等目标；
• 构建/安装前置依赖参考仓库 README：需要 CMake >= 3.31、Python >= 3.10、scikit-build-core，并要求torch == 2.13.0；也可以直接pip3 install sglang-kernel --upgrade安装已发布版本。

Step 9：验证

构建成功后运行测试与基准脚本：

pytest python/sglang/kernels/aot/tests/test_scale.py -q python python/sglang/kernels/aot/benchmark/bench_scale.py

PR CI 也会运行pr-test-sgl-kernel.yml：当检测到内核变更时，会额外触发 B200 任务sgl-kernel-b200-test（对应工作流文件为 .github/workflows/pr-test-sgl-kernel.yml）。该任务可作为 AOTsgl-kernel变更在 Blackwell 架构上的覆盖信号，合入前建议以其通过作为依据。

Troubleshooting：常见问题速查

参考文档

• python/sglang/kernels/aot/README.md — sgl-kernel 源码构建、安装与开发指引
• python/sglang/kernels/aot/include/sgl_kernel_ops.h — C++ 接口总目录
• python/sglang/kernels/aot/csrc/common_extension.cc —TORCH_LIBRARY_FRAGMENT注册实现
• python/sglang/kernels/aot/CMakeLists.txt — 源文件列表与构建配置
• python/sglang/kernels/aot/include/utils.h —DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16等分派宏定义
• python/sglang/kernels/aot/csrc/elementwise/activation.cu — FP16/BF16/FP32 分派模式的参考实现
• .claude/skills/add-jit-kernel/SKILL.md — 轻量 JIT 内核的配套新增教程，用于与 AOT 路径对照决策

小结：新增/修改文件清单

python/sglang/kernels/aot/csrc/elementwise/scale.cu # NEW: CUDA kernel + launcher python/sglang/kernels/aot/include/sgl_kernel_ops.h # MODIFIED: C++ declaration python/sglang/kernels/aot/csrc/common_extension.cc # MODIFIED: schema + dispatch registration python/sglang/kernels/aot/CMakeLists.txt # MODIFIED: add source file (alphabetical) python/sglang/kernels/aot/python/sgl_kernel/elementwise.py # MODIFIED: Python wrapper python/sglang/kernels/aot/python/sgl_kernel/__init__.py # MODIFIED: re-export Python API python/sglang/kernels/aot/tests/test_scale.py # NEW: tests python/sglang/kernels/aot/benchmark/bench_scale.py # NEW: benchmark

以上 8 个文件构成了一个 AOT 内核从 CUDA 实现到可验证产物的完整闭环：核心逻辑在csrc/中实现并通过宏完成 FP16/BF16/FP32 分派，接口经sgl_kernel_ops.h声明、common_extension.cc以m.def/m.impl注册为torch.ops.sgl_kernel.*算子，由CMakeLists.txt纳入编译；对外通过sgl_kernel/Python 包提供薄封装；最后以 pytest 校验正确性与异常路径、以 Tritondo_bench_cudagraph基准对比 PyTorch baseline，并依托sgl-kernel-b200-testCI 任务完成 Blackwell 架构覆盖。

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