add-sgl-kernel

教程：向 sgl-kernel 添加新 Kernel（AOT / 重量级）

对于公共算子命名空间、逻辑分组、惰性注册表元数据以及测试放置位置，请应用 kernel-organization。 下面的实现教程并不替代该 API 约定。

本教程将带你完成添加一个简单的逐元素缩放操作作为 AOT kernel 的完整流程。我们将实现 scale(x, factor) = x * factor 来演示完整的工作流。

目标

添加一个新操作，将张量的每个元素乘以一个标量因子：

• 输入：张量 x（CUDA）和标量 factor（float）
• 输出：x * factor（逐元素，原地或写入预先分配的 out）
• 支持的 dtype：FP16（torch.float16）、BF16（torch.bfloat16）、FP32（torch.float32）
  • 通过 DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 宏进行分发（定义于 python/sglang/kernels/aot/include/utils.h）

两条经验法则（必须遵循）

1. 当 kernel 不依赖 CUTLASS 或其他大型 C++ 项目时，优先使用 python/sglang/kernels/jit。这是轻量级 kernel 的默认路径，便于快速迭代。
2. 当 kernel 确实依赖 CUTLASS 或其他大型 C++ 项目，或应作为 AOT wheel / torch op 注册流程的一部分时，优先使用 sgl-kernel。
3. 例外：如果依赖是 flashinfer，或是已通过 flashinfer 提供的 CUTLASS，则 kernel 仍可实现为 jit_kernel。

此外，每个新 kernel 都必须附带：

• 测试（pytest）
• 基准测试脚本（来自 sglang.kernels.jit.benchmark 的 marker.do_bench —— 它适用于任何可调用对象，不仅限于 JIT kernel）

---

仓库集成图

你通常需要修改以下文件/区域：

• 实现：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

---

第 1 步：在 csrc/ 中实现 kernel

选择正确的子目录：

• 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;
  });
}

要点：

• 使用 at::Tensor（PyTorch 张量）、TORCH_CHECK 进行校验、at::cuda::getCurrentCUDAStream() 获取流
• 保持 Python 封装层精简；在 C++ 中紧邻启动路径处进行形状/dtype/设备校验
• DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 覆盖 float、half（FP16）、__nv_bfloat16（BF16）
• 每次 kernel 启动后添加设备错误检查
• 如果某个 kernel 仅在特定架构上工作，请用 TORCH_CHECK 强制执行，并在测试中跳过相应逻辑

---

第 2 步：在 include/sgl_kernel_ops.h 中添加 C++ 声明

编辑 python/sglang/kernels/aot/include/sgl_kernel_ops.h，添加到 elementwise 部分：

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

---

第 3 步：在 csrc/common_extension.cc 中注册 op

编辑 python/sglang/kernels/aot/csrc/common_extension.cc，在 TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) 内：

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

要点：

• Tensor! 表示原地 / 可变输出参数
• schema 对 torch.compile 以及一致的调用签名非常重要
• 保持 torch schema 使用 PyTorch 标量类型（此处为 float），但请注意，C++ 启动器签名对于 torch::Library 接受的标量参数仍需要 double

---

步骤 4：将新源文件添加到 CMakeLists.txt

编辑 python/sglang/kernels/aot/CMakeLists.txt，添加到 set(SOURCES ...)：

csrc/elementwise/scale.cu

要点：

• 保持列表按字母顺序排序（文件明确要求这样做）
• 如果内核有架构限制，请通过跳过逻辑在测试/基准测试中体现出来

---

步骤 5：在 python/sglang/kernels/aot/python/sgl_kernel/ 下暴露 Python API

优先遵循现有的模块组织方式。对于逐元素内核，通常的模式是：

• 在 python/sglang/kernels/aot/python/sgl_kernel/elementwise.py 中实现 Python 包装器
• 然后从 python/sglang/kernels/aot/python/sgl_kernel/__init__.py 中重新导出

例如，在 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

然后按照其他内核使用的现有导入风格，从 python/sglang/kernels/aot/python/sgl_kernel/__init__.py 中重新导出。

---

SGLang 集成（运行时使用所必需）

在暴露 AOT wheel 符号后，在 python/sglang/kernels/ops/<group>/ 下添加一个懒加载包装器和一个 KernelSpec。设置 backend=KernelBackend.AOT，记录 实际的设备/架构支持情况，并将 sgl_kernel 导入保留在 实现路径内部。SGLang 运行时和集成测试会导入该包装器。 对于此示例，使用 sglang.kernels.ops.elementwise.scale。

下面的 wheel 级测试验证其独立 API/构建。它们不能 替代 CI 注册的 SGLang 正确性测试（位于 test/registered/kernels/ops/elementwise/test_scale.py）和基准测试（位于 test/registered/kernels/benchmark/elementwise/bench_scale.py）。遵循 write-sglang-test 进行注册和 CI 预算；不要将测试注册 在已发布的 python/sglang/ 包下。

步骤 6：编写测试（必需）

创建 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"]))

---

步骤 7：添加基准测试（必需）

每个基准测试都必须考虑 L2 缓存复用 —— 参见 rules/kernel-benchmark.md。

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

import torch

import sgl_kernel
from sglang.kernels.jit.benchmark import marker


def sglang_scale(input: torch.Tensor, factor: float, out: torch.Tensor) -> None:
    sgl_kernel.scale(input, factor, out=out)


def torch_scale(input: torch.Tensor, factor: float, out: torch.Tensor) -> None:
    torch.mul(input, factor, out=out)


FN_MAP = {"sglang": sglang_scale, "torch": torch_scale}


@marker.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], [torch.float16])
@marker.parametrize("size", [2**n for n in range(10, 20)], [4096])  # 1K .. 512K
@marker.benchmark("provider", ["sglang", "torch"])
def benchmark(dtype: torch.dtype, size: int, provider: str):
    input = torch.randn(size, dtype=dtype, device="cuda")
    out = torch.empty_like(input)
    return marker.do_bench(
        FN_MAP[provider],
        # Pass every tensor through input_args (not a closure) so marker rotates
        # them across CUDA-graph calls to defeat L2 reuse.
        input_args=(input, 2.0, out),
        # Bandwidth = bytes(input) + bytes(out), both already in input_args.
        memory_output=None,
    )


if __name__ == "__main__":
    benchmark.run()

---

步骤 8：构建

构建：

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"

---

步骤 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。将该作业用作 AOT sgl-kernel 更改的 Blackwell 覆盖率信号。

---

故障排除

• 异步 CUDA 错误：CUDA_LAUNCH_BLOCKING=1
• 内存错误：compute-sanitizer --tool memcheck python ...
• 构建太慢 / OOM：减少 MAX_JOBS 和 SGL_KERNEL_COMPILE_THREADS
• 二进制膨胀：使用 python/sglang/kernels/aot/analyze_whl_kernel_sizes.py
• CMake 源文件列表：如果你的 .cu 文件未出现在 SOURCES 中，符号将在链接时未定义

---

参考资料

• python/sglang/kernels/aot/README.md
• python/sglang/kernels/aot/include/sgl_kernel_ops.h
• python/sglang/kernels/aot/csrc/common_extension.cc
• 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 分发模式的参考

创建/修改的文件摘要

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