从零为 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优先选择 python/sglang/kernels/jit当内核不依赖CUTLASS 或其他大型 C 工程时JIT 是默认路径适合迭代快速的轻量内核。优先选择 sgl-kernelAOT当内核依赖CUTLASS 或其他大型 C 工程或希望纳入 AOT wheel 与 torch op 注册流程时。例外情况如果依赖的是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.hTorch 扩展注册python/sglang/kernels/aot/csrc/common_extension.cc构建python/sglang/kernels/aot/CMakeLists.txtset(SOURCES ...)Python APIpython/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_castint64_t(blockIdx.x) * blockDim.x threadIdx.x; if (idx n) { out[idx] static_castT(static_castfloat(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_kernelc_typeblocks, threads, 0, stream( static_castc_type*(out.data_ptr()), static_castconst c_type*(input.data_ptr()), static_castfloat(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覆盖floatFP32、halfFP16、__nv_bfloat16BF16其实现定义于 python/sglang/kernels/aot/include/utils.h按at::ScalarTypeswitch 到float、_DISPATCH_CASE_F16、_DISPATCH_CASE_BF16其余类型走TORCH_CHECK(false, ...)报错。本教程示例算子支持 FP16torch.float16、BF16torch.bfloat16、FP32torch.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.compilem.impl设备绑定的方式注册扩展。scalar 类型约定torch schema 中按 PyTorch 标量类型书写此处为float但 C launcher 的函数签名仍需对接受 scalar 实参使用doubletorch::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.pyimport 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, dtypedtype, devicecuda) out torch.empty_like(input) result sgl_kernel.scale(input, factor, outout) 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, rtolrtol, atolatol) def test_scale_shape_mismatch(): input torch.randn(128, dtypetorch.float16, devicecuda) out torch.empty(256, dtypetorch.float16, devicecuda) with pytest.raises(RuntimeError, matchsame shape): sgl_kernel.scale(input, 2.0, outout) def test_scale_cpu_input(): input torch.randn(128, dtypetorch.float16) # CPU out torch.empty_like(input) with pytest.raises(RuntimeError, matchCUDA): sgl_kernel.scale(input, 2.0, outout) 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.pyimport 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_valsconfigs, line_argprovider, line_vals[sglang, torch], line_names[SGL Kernel, PyTorch], styles[(green, -), (red, --)], ylabelµs (median), plot_namescale-performance, args{}, ) ) def benchmark(dtype, size, provider): input torch.randn(size, dtypedtype, devicecuda) out torch.empty_like(input) factor 2.0 if provider sglang: fn lambda: sgl_kernel.scale(input, factor, outout) 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_dataTrue)基准测试约定与仓库补充说明基准脚本统一放在 python/sglang/kernels/aot/benchmark 目录下如bench_activation.py、bench_rmsnorm.py等文件名遵循bench_*.py命名仓库 README 建议优先使用triton.testing.do_bench_cudagraph进行内核基准相比do_bench它能降低 CPU 开销对内核性能测量精度的影响、把 PDLProgrammatic 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_JOBS2 CMAKE_ARGS-DSGL_KERNEL_COMPILE_THREADS1构建相关说明与 README 及 Makefile 一致make build默认占用全部可用 CPU 核可通过MAX_JOBS控制 make 与 CMake 并行度通过CMAKE_ARGS-DSGL_KERNEL_COMPILE_THREADS1额外限制 NVCC 内部线程数从而降低 CPU 占用与峰值内存Makefile 还提供check-deps/install-deps安装scikit-build-core、isort、black、installpip install -e . --no-build-isolation开发模式安装、rebuildclean 后重建、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.pyPR CI 也会运行pr-test-sgl-kernel.yml当检测到内核变更时会额外触发 B200 任务sgl-kernel-b200-test对应工作流文件为 .github/workflows/pr-test-sgl-kernel.yml。该任务可作为 AOTsgl-kernel变更在 Blackwell 架构上的覆盖信号合入前建议以其通过作为依据。Troubleshooting常见问题速查症状处置手段异步 CUDA 错误设置环境变量CUDA_LAUNCH_BLOCKING1让 launch 同步便于定位出错的内核显存类错误使用compute-sanitizer --tool memcheck python ...做内存检查构建过慢 / OOM降低MAX_JOBS与SGL_KERNEL_COMPILE_THREADS的取值如make build -j1 MAX_JOBS2 CMAKE_ARGS-DSGL_KERNEL_COMPILE_THREADS1wheel 体积膨胀使用 python/sglang/kernels/aot/analyze_whl_kernel_sizes.py 分析内核尺寸需pip install cubloaty定位过大的内核与模板实例化膨胀CMake SOURCES 遗漏.cu文件若未加入SOURCES符号在链接期会未定义先检查set(SOURCES ...)列表参考文档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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考