SuperAPI AI 工具集

add-jit-kernel

通用与其它

查看项目

更新于 10-16 04:31

技能介绍

用途:该教程面向 SGLang 贡献者,演示如何向 sglang.kernels 的 JIT 基础设施及公共算子组添加一个新的轻量级 JIT CUDA 内核。示例实现逐元素缩放 scale(x, factor) = x * factor,输入为 CUDA 张量 x 与运行时传入的 float 标量 factor,输出为内部新分配的 x * factor,支持 FP16、BF16、FP32。

主要能力与约定:教程给出 JIT 与 AOT(sgl-kernel)的选择标准——不依赖 CUTLASS 等大型 C++ 项目的轻量内核优先用 jit_kernel;依赖 CUTLASS 或需参与 wheel 构建与 torch op 注册的用 AOT;依赖 flashinfer 或经 flashinfer 提供的 CUTLASS 仍可用 jit_kernel。约定包括:JIT 代码置于 namespace sglang;调优策略放在 Python,C++ 只接收结果;检查按 static_assert > C++ host 检查 > 缓存 Python > 每次调用 Python 的成本排序;使用固定宽度整数类型;C++ 用 Doxygen 注释;源码仅 ASCII;details 命名空间为私有;只读指针用 const T* restrict;关注寄存器预算,可用 SGLANG_JIT_LOG_RESOURCE_USAGE=1 SGLANG_JIT_FORCE_RECOMPILE=1 记录寄存器、溢出与共享内存;向量化内核建议以 .ensure_alignment(w) 收尾。

公共抽象:教程介绍 python/sglang/kernels/jit/include/sgl_kernel/ 下的工具,包括 utils.h(CHECK_HOST、host::RuntimeCheck、host::Panic、host::div_ceil、host::irange、host::pointer::offset)、utils.cuh(fp16_t 等类型别名、SGL_DEVICE、device::kWarpThreads、device::get_lane_id、device::load_as/store_as、host::LaunchKernel、PDL 相关、CHECK_CUDA)、tensor.h(SymbolicSize、SymbolicDType、SymbolicDevice、TensorMatcher 及 verify)、ffi.h(host::alloc_workspace_tensor、host::ffi::empty/empty_like/from_blob)、type.cuh(DTypeTrait、packed_t、device::cast/unpack、ReductionOp)、vec.cuh(AlignedVector、LoadStoreBytes)、tile.cuh(tile::Memory)等。

适用场景与限制:适合需要快速迭代、首次使用时编译的轻量 CUDA 内核开发。资料未提供具体安装命令、完整代码示例、性能数据或收费信息。