用途:本教程面向 sgl-project/sglang 仓库,演示如何向 sgl-kernel 添加一个重量级 AOT(提前编译)CUDA/C++ 算子,并以 element-wise 的 scale(x, factor) = x * factor 为例走通完整流程。
主要能力与规则:支持 FP16、BF16、FP32,通过 DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 宏分发。两条经验规则:若算子不依赖 CUTLASS 等大型 C++ 项目,优先用 python/sglang/kernels/jit;若依赖 CUTLASS 或需进入 AOT wheel / torch op 注册流程,则用 sgl-kernel。例外是依赖 flashinfer 或经 flashinfer 提供的 CUTLASS 时仍可做 jit_kernel。每个新算子必须附带 pytest 测试和基准脚本(marker.do_bench)。
使用方式:需改动实现文件 csrc/elementwise/scale.cu、声明 include/sgl_kernel_ops.h、注册 csrc/common_extension.cc、构建 CMakeLists.txt(SOURCES 按字母序)、Python API 与 init.py 再导出、测试与基准文件。还需在 python/sglang/kernels/ops/<group>/ 下加 lazy wrapper 与 KernelSpec,backend=KernelBackend.AOT。构建命令:cd python/sglang/kernels/aot && make build -j16;资源受限时用 make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"。验证:pytest python/sglang/kernels/aot/tests/test_scale.py -q 与 python python/sglang/kernels/aot/benchmark/bench_scale.py。
适用场景与限制:适用于需 AOT wheel 或 torch op 注册的 CUDA 算子开发。基准须考虑 L2 缓存复用。CI 会运行 pr-test-sgl-kernel.yml 及 B200 任务。排错可用 CUDA_LAUNCH_BLOCKING=1、compute-sanitizer、analyze_whl_kernel_sizes.py。若 .cu 未加入 SOURCES 会链接报错。具体收费、平台等信息暂未提供。