mcpField
A
add-sgl-kernel

add-sgl-kernel

github.com/sgl-project開発38.3k40 回インストール今日更新
リポジトリ

add-sgl-kernel について

`sgl-kernel` に AOT CUDA カーネルを追加し、SGLang に統合するためのチュートリアルスキルです。要素ごとのスケーリングを例に、実装、登録、Python API、テスト、ベンチマーク、ビルド、検証を扱います。

主な機能

  1. 01JIT と AOT の配置方針を説明
  2. 02CUDA 実装と登録を案内
  3. 03Python API と SGLang ラッパーを追加
  4. 04pytest テストとベンチマークを作成
  5. 05カーネルをビルドして検証

活用シーン

  1. 01要素ごとの CUDA 演算の追加
  2. 02SGLang ランタイムへのカーネル統合
  3. 03AOT カーネル変更のテストと性能測定

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/sgl-project/sglang --skill add-sgl-kernel

インストール後、Claude Code で/add-sgl-kernelを入力して読み込まれたか確認します。

詳細

チュートリアル：sgl-kernel に新しいカーネルを追加する（AOT／重量級）

パブリックオペレーターの名前空間、論理的なグループ化、遅延レジストリメタデータ、テストの配置には、kernel-organization を適用してください。以下の実装チュートリアルは、その API 契約に代わるものではありません。

このチュートリアルでは、単純な要素ごとのスケール演算を AOT カーネルとして追加する方法を説明します。完全なワークフローを示すために、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 で定義）を介してディスパッチされます

2 つの基本ルール（必ず守ってください）

  1. カーネルが CUTLASS やその他の大規模な C++ プロジェクトに依存しない場合は、まず python/sglang/kernels/jit を優先してください。素早い反復が有効な軽量カーネルでは、これがデフォルトの方法です。
  2. カーネルが CUTLASS やその他の大規模な C++ プロジェクトに依存する場合、または AOT wheel／torch op 登録フローに含める場合は、sgl-kernel を優先してください。
  3. 例外：依存先が flashinfer、または flashinfer 経由ですでに提供されている CUTLASS の場合でも、カーネルは jit_kernel として実装できます。

さらに、新しいカーネルにはすべて、次のものを含める必要があります。

  • テスト（pytest）
  • ベンチマークスクリプト（sglang.kernels.jit.benchmark の marker.do_bench。JIT カーネルに限らず、あらゆる呼び出し可能オブジェクトで動作します）

リポジトリ統合マップ

通常、以下のファイル／領域を変更します。

  • 実装：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/ にカーネルを実装する

適切なサブディレクトリを選びます。

  • 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 ラッパーは薄く保ち、シェイプ／dtype／デバイスの検証は起動処理の直前に C++ で行います
  • DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 は float、half（FP16）、__nv_bfloat16（BF16）を扱います
  • カーネルを起動するたびにデバイスエラーチェックを追加します
  • 特定のアーキテクチャでしか動作しないカーネルの場合は、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! はインプレース／変更可能な出力引数を意味します
  • スキーマは torch.compile と一貫した呼び出しシグネチャにとって重要です
  • Torch スキーマでは PyTorch のスカラー型（ここでは float）を使います。ただし、torch::Library で受け取るスカラー引数については、C++ ランチャーのシグネチャでは引き続き 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 ラッパーを python/sglang/kernels/aot/python/sgl_kernel/elementwise.py に実装します
  • 次に、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 ホイールのシンボルを公開した後、python/sglang/kernels/ops/<group>/ に遅延ラッパーと KernelSpec を追加します。backend=KernelBackend.AOT を設定し、実際にサポートされるデバイス／アーキテクチャを記録してください。また、sgl_kernel のインポートは実装パス内にとどめてください。SGLang の実行時処理と統合テストでは、このラッパーをインポートします。この例では、sglang.kernels.ops.elementwise.scale を使用します。

以下のホイールレベルのテストでは、単独での API／ビルドを検証します。これらは、CI に登録する SGLang の正当性テスト test/registered/kernels/ops/elementwise/test_scale.py やベンチマーク test/registered/kernels/benchmark/elementwise/bench_scale.py の代わりにはなりません。登録方法と CI の実行予算については write-sglang-test に従ってください。配布対象の 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

コメント

Markdown 対応