chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
607
third_party/sglang/.claude/skills/add-jit-kernel/SKILL.md
vendored
Normal file
607
third_party/sglang/.claude/skills/add-jit-kernel/SKILL.md
vendored
Normal file
@@ -0,0 +1,607 @@
|
||||
---
|
||||
name: add-jit-kernel
|
||||
description: Step-by-step tutorial for adding a new lightweight JIT CUDA kernel to sglang's jit_kernel module
|
||||
---
|
||||
|
||||
# Tutorial: Adding a New JIT Kernel to SGLang
|
||||
|
||||
This tutorial walks through adding a simple element-wise scale operation as a JIT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow.
|
||||
|
||||
## Goal
|
||||
|
||||
Add a new operation that scales each element of a tensor by a scalar factor:
|
||||
|
||||
- Input: tensor `x` (CUDA) and scalar `factor` (float, passed at runtime)
|
||||
- Output: `x * factor` (element-wise), allocated internally
|
||||
- Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)**
|
||||
|
||||
## When to use JIT vs AOT (`sgl-kernel`)
|
||||
|
||||
- **JIT (`jit_kernel`)**: prefer this first for kernels that do **not** depend on CUTLASS or another large C++ project. It is the default choice for lightweight kernels that benefit from rapid iteration and first-use compilation.
|
||||
- **AOT (`sgl-kernel`)**: prefer this when the kernel **does** depend on CUTLASS or another large C++ project, or when it should live in `sgl-kernel/` and participate in the wheel build / torch op registration flow.
|
||||
- **Exception**: kernels that depend on `flashinfer`, or on CUTLASS that is already provided through `flashinfer`, can still be implemented as `jit_kernel`.
|
||||
|
||||
---
|
||||
|
||||
## Common Abstractions in `python/sglang/jit_kernel/include/sgl_kernel/`
|
||||
|
||||
**Always prefer these abstractions over raw CUDA primitives.** They provide safety, readability, and consistency with the rest of the codebase.
|
||||
|
||||
**Important include rule:** for every `#include <sgl_kernel/...>` line, add a short trailing comment explaining why that header is included (for example `// For TensorMatcher, SymbolicSize, SymbolicDevice`). This matches the current JIT kernel style and keeps include usage self-documenting.
|
||||
|
||||
### `utils.h` — Host-side utilities
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/utils.h>
|
||||
```
|
||||
|
||||
- **`host::RuntimeCheck(cond, args...)`** — Assert a condition at runtime; throws `PanicError` with file/line info on failure. Prefer this over bare `assert`.
|
||||
- **`host::Panic(args...)`** — Unconditionally throw a `PanicError` with a descriptive message.
|
||||
- **`host::div_ceil(a, b)`** — Integer ceiling division `(a + b - 1) / b`.
|
||||
- **`host::irange(n)`** / **`host::irange(start, end)`** — Range views for cleaner loops.
|
||||
- **`host::pointer::offset(ptr, offsets...)`** — Byte-safe pointer arithmetic on `void*`. Use this instead of raw casts.
|
||||
|
||||
### `utils.cuh` — Device-side utilities + `LaunchKernel`
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
```
|
||||
|
||||
- **Type aliases**: `fp16_t`, `bf16_t`, `fp32_t`, `fp8_e4m3_t`, `fp8_e5m2_t` and their packed variants `fp16x2_t`, `bf16x2_t`, `fp32x2_t`, etc.
|
||||
- **`SGL_DEVICE`** — Expands to `__forceinline__ __device__`. Use on all device functions.
|
||||
- **`device::kWarpThreads`** — Constant `32`.
|
||||
- **`device::load_as<T>(ptr, offset)`** / **`device::store_as<T>(ptr, val, offset)`** — Type-safe loads/stores from `void*`.
|
||||
- **`device::pointer::offset(ptr, offsets...)`** — Pointer arithmetic on device.
|
||||
- **`host::LaunchKernel(grid, block, device_or_stream [, smem])`** — RAII kernel launcher that:
|
||||
- Resolves the CUDA stream from a `DLDevice` via TVM-FFI automatically.
|
||||
- Checks the CUDA error with file/line info after launch via `operator()(kernel, args...)`.
|
||||
- Supports `.enable_pdl(bool)` for PDL (Programmatic Dependent Launch, SM90+).
|
||||
- **`host::RuntimeDeviceCheck(cudaError_t)`** — Check a CUDA error; throw on failure.
|
||||
|
||||
### `tensor.h` — Tensor validation (`TensorMatcher`, Symbolic types)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
```
|
||||
|
||||
This is the **primary validation API** for all kernel launchers. Use it to validate every `tvm::ffi::TensorView` argument.
|
||||
|
||||
- **`host::SymbolicSize{"name"}`** — A named symbolic dimension. Call `.set_value(n)` to pin it, `.unwrap()` to extract after verification.
|
||||
- **`host::SymbolicDType`** — Symbolic dtype. Use `.set_options<Ts...>()` to restrict allowed types.
|
||||
- **`host::SymbolicDevice`** — Symbolic device. Use `.set_options<kDLCUDA>()` to restrict to CUDA.
|
||||
- **`host::TensorMatcher({dims...})`** — Fluent builder for tensor validation:
|
||||
- `.with_dtype<T>()` — require a specific C++ type (e.g. `fp16_t`)
|
||||
- `.with_dtype<T1, T2, ...>()` — allow a set of types
|
||||
- `.with_device<kDLCUDA>(device_sym)` — require CUDA and bind the checked device to a `SymbolicDevice`
|
||||
- `.with_strides({strides...})` — validate strides (omit to require contiguous)
|
||||
- `.verify(tensor_view)` — execute the check; throws `PanicError` with full context on failure; **chainable** (`verify(a).verify(b)` to check multiple tensors with the same shape)
|
||||
|
||||
**Typical pattern:**
|
||||
```cpp
|
||||
auto N = SymbolicSize{"num_elements"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<fp16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(dst)
|
||||
.verify(src); // same shape, dtype, device as dst
|
||||
const size_t n = N.unwrap();
|
||||
const DLDevice dev = device.unwrap();
|
||||
```
|
||||
|
||||
### `type.cuh` — `dtype_trait<T>` and `packed_t<T>`
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/type.cuh>
|
||||
```
|
||||
|
||||
- **`dtype_trait<T>`** — Static trait struct for each scalar type. Provides:
|
||||
- `dtype_trait<T>::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`)
|
||||
- `dtype_trait<T>::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`)
|
||||
- `dtype_trait<T>::max/min(x, y)` — type-dispatched binary math (primarily for `fp32_t`)
|
||||
- **`packed_t<T>`** — Two-element packed alias: `packed_t<fp16_t>` = `fp16x2_t`, `packed_t<bf16_t>` = `bf16x2_t`, `packed_t<fp32_t>` = `fp32x2_t`. Use for vectorized loads/stores.
|
||||
- **`device::cast<To, From>(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast<fp32x2_t, fp16x2_t>(v)`.
|
||||
|
||||
### `vec.cuh` — Vectorized memory access (`AlignedVector`)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
```
|
||||
|
||||
- **`device::AlignedVector<T, N>`** — Aligned storage for N elements of type T. N must be a power of two, `sizeof(T)*N <= 32`. Enables vectorized loads/stores for bandwidth efficiency. In terms of API/codegen constraints, the upper bound is 256-bit; in practice, 128-bit is the portable default, while 256-bit vectorization is typically only viable on `SM100+` and should be gated by an architecture check when needed.
|
||||
- `.load(ptr, offset)` — vectorized load from `ptr[offset]`
|
||||
- `.store(ptr, offset)` — vectorized store to `ptr[offset]`
|
||||
- `.fill(value)` — fill all lanes
|
||||
- `operator[](i)` — element access
|
||||
|
||||
### `tile.cuh` — `tile::Memory` (strided memory access pattern)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
```
|
||||
|
||||
- `tile::Memory<T>` is fundamentally a **1D cooperative accessor** over a contiguous region.
|
||||
- **`device::tile::Memory<T>::cta(blockDim.x)`** — Creates a tile accessor where each thread handles `tid = threadIdx.x` with stride `tsize` (for `cta(blockDim.x)`, this is `blockDim.x`). Common for loops over a 1D array.
|
||||
- **`.load(ptr, offset)`** — loads `ptr[tid + offset * tsize]`
|
||||
- **`.store(ptr, val, offset)`** — stores to `ptr[tid + offset * tsize]`
|
||||
- **`.in_bound(n, offset)`** — boundary check
|
||||
|
||||
For a **2D tile**, either flatten `(row, col)` into a linear tile index first, or compute the address manually with `ptr[row * stride + col]` using your thread/block coordinates.
|
||||
|
||||
### `math.cuh` — Device math (`device::math::`)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/math.cuh>
|
||||
```
|
||||
|
||||
- `device::math::max/min<T>(a, b)` — type-dispatched binary math via `dtype_trait`
|
||||
- `device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x)` — type-dispatched unary math via `dtype_trait`
|
||||
|
||||
### `warp.cuh` — Warp-level primitives
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
```
|
||||
|
||||
- `device::warp::reduce_sum<T>(value)` — warp-level sum reduction via `__shfl_xor_sync`
|
||||
- `device::warp::reduce_max<T>(value)` — warp-level max reduction
|
||||
|
||||
### `cta.cuh` — CTA-level primitives
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/cta.cuh>
|
||||
```
|
||||
|
||||
- `device::cta::reduce_max<T>(value, smem, min_value)` — CTA-wide max using shared memory + warp reduction. Caller is responsible for a `__syncthreads()` after if the result in `smem[0]` is needed.
|
||||
|
||||
### `atomic.cuh` — Atomic operations
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/atomic.cuh>
|
||||
```
|
||||
|
||||
- `device::atomic::max(float* addr, float value)` — float atomic max (handles negative values correctly via bit tricks).
|
||||
|
||||
### `runtime.cuh` — Occupancy and device info
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
```
|
||||
|
||||
- `host::runtime::get_blocks_per_sm(kernel, block_dim)` — max active blocks per SM (occupancy)
|
||||
- `host::runtime::get_sm_count(device_id)` — number of SMs on the device
|
||||
- `host::runtime::get_cc_major(device_id)` — compute capability major version
|
||||
|
||||
**Persistent kernel pattern** (cap blocks to SM count × occupancy):
|
||||
```cpp
|
||||
static const uint32_t max_occ = runtime::get_blocks_per_sm(kernel, kBlockSize);
|
||||
static const uint32_t num_sm = runtime::get_sm_count(device.unwrap().device_id);
|
||||
const auto num_blocks = std::min(num_sm * max_occ, div_ceil(n, kBlockSize));
|
||||
LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 0 (optional): Generate a `.clangd` config for better IDE support
|
||||
|
||||
```bash
|
||||
python -m sglang.jit_kernel
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 1: Implement the CUDA kernel in `jit_kernel/csrc/`
|
||||
|
||||
Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`.
|
||||
|
||||
The implementation fully uses the project abstractions described above:
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/type.cuh> // For dtype_trait, fp16_t, bf16_t, fp32_t
|
||||
#include <sgl_kernel/utils.h> // For RuntimeCheck, div_ceil
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel, SGL_DEVICE
|
||||
#include <sgl_kernel/vec.cuh> // For AlignedVector
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// ----------------------------------------------------------------
|
||||
// Kernel: element-wise scale using vectorized 128-bit loads/stores
|
||||
// T = fp16_t | bf16_t | fp32_t
|
||||
// kVecN = number of elements per vector load (e.g. 8 for fp16)
|
||||
// factor = runtime scale factor
|
||||
// ----------------------------------------------------------------
|
||||
template <typename T, int kVecN>
|
||||
__global__ void scale_kernel(T* __restrict__ dst,
|
||||
const T* __restrict__ src,
|
||||
float factor,
|
||||
uint32_t n_total) {
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
const uint32_t n_vecs = n_total / kVecN;
|
||||
|
||||
// --- vectorised body ---
|
||||
const uint32_t vec_stride = blockDim.x * gridDim.x;
|
||||
for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
vi < n_vecs;
|
||||
vi += vec_stride) {
|
||||
vec_t v;
|
||||
v.load(src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
v[i] = static_cast<T>(static_cast<float>(v[i]) * factor);
|
||||
}
|
||||
v.store(dst, vi);
|
||||
}
|
||||
|
||||
// --- scalar tail ---
|
||||
const uint32_t base = n_vecs * kVecN;
|
||||
const uint32_t scalar_stride = blockDim.x * gridDim.x;
|
||||
for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
base + i < n_total;
|
||||
i += scalar_stride) {
|
||||
dst[base + i] = static_cast<T>(static_cast<float>(src[base + i]) * factor);
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------
|
||||
// Launcher: validates tensors, selects vector width, launches kernel
|
||||
// ----------------------------------------------------------------
|
||||
template <typename T>
|
||||
void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
|
||||
using namespace host;
|
||||
|
||||
// 1. Validate input tensors with TensorMatcher
|
||||
SymbolicSize N = {"num_elements"};
|
||||
SymbolicDevice device_;
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N}) //
|
||||
.with_dtype<T>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.verify(dst)
|
||||
.verify(src); // same shape / dtype / device as dst
|
||||
|
||||
const uint32_t n = static_cast<uint32_t>(N.unwrap());
|
||||
const DLDevice device = device_.unwrap();
|
||||
|
||||
RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n);
|
||||
|
||||
// 2. Choose vector width for 128-bit loads (16 bytes)
|
||||
// fp16/bf16: 8 elements × 2 bytes = 16 bytes
|
||||
// fp32: 4 elements × 4 bytes = 16 bytes
|
||||
constexpr int kVecN = 16 / sizeof(T);
|
||||
const uint32_t n_work_items = div_ceil(n, static_cast<uint32_t>(kVecN));
|
||||
|
||||
// 3. Launch
|
||||
constexpr uint32_t kBlockSize = 256;
|
||||
const uint32_t grid = div_ceil(n_work_items, kBlockSize);
|
||||
|
||||
LaunchKernel(grid, kBlockSize, device)(
|
||||
scale_kernel<T, kVecN>,
|
||||
static_cast<T*>(dst.data_ptr()),
|
||||
static_cast<const T*>(src.data_ptr()),
|
||||
factor,
|
||||
n);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
```
|
||||
|
||||
**Key points:**
|
||||
|
||||
- Include headers from `sgl_kernel/` — **not** raw CUDA headers for anything already covered
|
||||
- Add a short trailing `// For ...` explanation to every `#include <sgl_kernel/...>` line
|
||||
- Use `TensorMatcher` for all tensor validation; never manually check shape/dtype/device
|
||||
- Use `AlignedVector` for vectorised 128-bit loads/stores — significant bandwidth win
|
||||
- Use `LaunchKernel` — it resolves the stream and checks errors automatically
|
||||
- Use `RuntimeCheck` for runtime assertions with useful error messages
|
||||
- Prefer passing runtime scalars like `factor` directly unless compile-time specialisation is genuinely required
|
||||
- `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`)
|
||||
- `device::cast<To, From>` or `dtype_trait<T>::from(val)` for cross-type conversions
|
||||
- `device::math::` functions for device math instead of bare `__` intrinsics
|
||||
|
||||
---
|
||||
|
||||
## Step 2: Add the Python wrapper in `jit_kernel/`
|
||||
|
||||
Create `python/sglang/jit_kernel/scale.py`:
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_scale_module(dtype: torch.dtype) -> Module:
|
||||
"""Compile and cache the JIT scale module for a given dtype."""
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"scale",
|
||||
*args,
|
||||
cuda_files=["elementwise/scale.cuh"],
|
||||
cuda_wrappers=[("scale", f"scale<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> torch.Tensor:
|
||||
"""
|
||||
Element-wise scale: dst = src * factor.
|
||||
|
||||
Supported dtypes: torch.float16, torch.bfloat16, torch.float32.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
src : CUDA tensor (FP16 / BF16 / FP32)
|
||||
factor : scale factor
|
||||
out : optional pre-allocated output tensor (same shape/dtype as src)
|
||||
|
||||
Returns
|
||||
-------
|
||||
Scaled tensor (dst = src * factor).
|
||||
"""
|
||||
if not src.is_cuda:
|
||||
raise RuntimeError("src must be a CUDA tensor")
|
||||
if src.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise RuntimeError(
|
||||
f"Unsupported dtype {src.dtype}. Supported: float16, bfloat16, float32"
|
||||
)
|
||||
if out is None:
|
||||
out = torch.empty_like(src)
|
||||
else:
|
||||
if out.shape != src.shape:
|
||||
raise RuntimeError("out shape must match src")
|
||||
if out.dtype != src.dtype:
|
||||
raise RuntimeError("out dtype must match src")
|
||||
if out.device != src.device:
|
||||
raise RuntimeError("out device must match src")
|
||||
|
||||
# Keep the Python wrapper thin, but still enforce the basic preconditions
|
||||
# that the current JIT/FFI path does not reject safely on its own.
|
||||
module = _jit_scale_module(src.dtype)
|
||||
module.scale(out, src, factor)
|
||||
return out
|
||||
```
|
||||
|
||||
**Key points:**
|
||||
|
||||
- Use `cache_once` — **not** `functools.lru_cache` (incompatible with `torch.compile`)
|
||||
- `load_jit` first arg(s) form the unique build marker; same marker = same cached binary
|
||||
- Only include compile-time specialisation knobs in the build marker; runtime values like `factor` should stay runtime unless the kernel truly needs templating
|
||||
- `cuda_wrappers`: `(export_name, kernel_symbol)` — `export_name` is called from Python
|
||||
- `make_cpp_args(dtype, ...)` converts `torch.dtype` to C++ type alias:
|
||||
- Keep Python launchers thin, but still validate the basic invariants (`is_cuda`, supported dtype, `out` metadata). In the current JIT/FFI path, invalid tensors are not always rejected safely before launch
|
||||
|
||||
| `torch.dtype` | C++ type |
|
||||
|--------------------|------------|
|
||||
| `torch.float16` | `fp16_t` |
|
||||
| `torch.bfloat16` | `bf16_t` |
|
||||
| `torch.float32` | `fp32_t` |
|
||||
|
||||
---
|
||||
|
||||
## Step 3 (optional): Tune JIT build flags
|
||||
|
||||
```python
|
||||
return load_jit(
|
||||
"scale",
|
||||
*args,
|
||||
cuda_files=["elementwise/scale.cuh"],
|
||||
cuda_wrappers=[("scale", f"scale<{args}>")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
)
|
||||
```
|
||||
|
||||
If your kernel requires SM90+, raise a clear Python error before calling `load_jit`:
|
||||
|
||||
```python
|
||||
if torch.cuda.get_device_capability()[0] < 9:
|
||||
raise RuntimeError("This kernel requires SM90 (Hopper) or later")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 4: Write tests (required)
|
||||
|
||||
JIT kernel tests live under `python/sglang/jit_kernel/tests/`. **CI does not run `pytest` in that directory directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` there (and every `bench_*.py` under `benchmark/`), collects `register_*_ci(...)` calls by **statically parsing each file’s AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check.
|
||||
|
||||
- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `stage-b-kernel-unit-1-gpu-large` (see `.github/workflows/pr-test-jit-kernel.yml`: `python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large`).
|
||||
- **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/jit_kernel/utils.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`).
|
||||
|
||||
Registration pattern (module level, **literal** `est_time` and `suite` strings — required for AST parsing):
|
||||
|
||||
```python
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
# Optional second registration: same file also listed under the nightly kernel suite
|
||||
# register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
```
|
||||
|
||||
Keep `est_time` and `suite` as literal values. `run_suite.py` collects them from the file AST, so computed values and helper wrappers can break CI discovery.
|
||||
|
||||
Use `register_cuda_ci(..., disabled="reason")` if the file must stay in-tree but should be skipped in CI (e.g. multi-GPU only).
|
||||
|
||||
**Run like CI** (from repo root):
|
||||
|
||||
```bash
|
||||
cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large
|
||||
```
|
||||
|
||||
For fast iteration you can still run `pytest` on a single file locally; CI coverage is via `run_suite.py`.
|
||||
|
||||
Create `python/sglang/jit_kernel/tests/test_scale.py`:
|
||||
|
||||
```python
|
||||
import pytest
|
||||
import torch
|
||||
from sglang.jit_kernel.scale import scale
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("size", [1, 127, 128, 1024, 4097]) # cover tail remainder
|
||||
@pytest.mark.parametrize("factor", [0.5, 1.0, 2.0, 3.0])
|
||||
def test_scale_correctness(dtype, size, factor):
|
||||
src = torch.randn(size, dtype=dtype, device="cuda")
|
||||
out = scale(src, factor)
|
||||
expected = src * 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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_scale_out_param(dtype):
|
||||
src = torch.randn(1024, dtype=dtype, device="cuda")
|
||||
out = torch.empty_like(src)
|
||||
result = scale(src, 2.0, out=out)
|
||||
assert result is out
|
||||
torch.testing.assert_close(out, src * 2.0, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
def test_scale_cpu_error():
|
||||
src = torch.randn(128, dtype=torch.float16) # CPU tensor
|
||||
with pytest.raises(RuntimeError, match="CUDA"):
|
||||
scale(src, 2.0)
|
||||
|
||||
|
||||
def test_scale_unsupported_dtype():
|
||||
src = torch.randint(0, 10, (128,), dtype=torch.int32, device="cuda")
|
||||
with pytest.raises(RuntimeError, match="dtype"):
|
||||
scale(src, 2.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 5: Add a benchmark (required)
|
||||
|
||||
Benchmarks are `bench_*.py` files under `python/sglang/jit_kernel/benchmark/`. They are picked up by the same `run_suite.py` machinery as unit tests. Register them for **`stage-b-kernel-benchmark-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite stage-b-kernel-benchmark-1-gpu-large`).
|
||||
|
||||
Create `python/sglang/jit_kernel/benchmark/bench_scale.py`:
|
||||
|
||||
```python
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
from sglang.jit_kernel.scale import scale as jit_scale
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
SIZE_LIST = get_benchmark_range(
|
||||
full_range=[2**n for n in range(10, 20)], # 1K … 512K elements
|
||||
ci_range=[4096, 65536],
|
||||
)
|
||||
|
||||
configs = list(itertools.product(SIZE_LIST))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["size"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=["jit", "torch"],
|
||||
line_names=["SGL JIT Kernel", "PyTorch"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="scale-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(size: int, provider: str):
|
||||
src = torch.randn(size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
|
||||
factor = 2.0
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_scale(src, factor)
|
||||
else:
|
||||
fn = lambda: src * factor
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
```
|
||||
|
||||
Run locally:
|
||||
|
||||
```bash
|
||||
python python/sglang/jit_kernel/benchmark/bench_scale.py
|
||||
```
|
||||
|
||||
Run the benchmark suite the way CI does:
|
||||
|
||||
```bash
|
||||
cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-benchmark-1-gpu-large
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **`No CI registry found in ...` from `run_suite.py`**: add a module-level `register_cuda_ci(...)` with literal `est_time` and `suite` (and optional `nightly=True`); starred args and non-literal values break AST collection
|
||||
- **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/jit_kernel/csrc/`; reduce template argument combinations
|
||||
- **CUDA crash / illegal memory access**: `CUDA_LAUNCH_BLOCKING=1`; `compute-sanitizer --tool memcheck python ...`
|
||||
- **Unstable benchmark results**: `run_benchmark` uses CUDA-graph-based timing by default
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
- `docs/developer_guide/development_jit_kernel_guide.md`
|
||||
- `test/run_suite.py` — suite names, discovery of `jit_kernel/tests/` and `jit_kernel/benchmark/`, execution entrypoint for CI
|
||||
- `python/sglang/test/ci/ci_register.py` — `register_cuda_ci` and AST registration rules
|
||||
- `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args`, `should_run_full_tests`, `get_ci_test_range`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `dtype_trait`, `packed_t`, `device::cast`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce_sum/max`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers
|
||||
- `python/sglang/jit_kernel/csrc/add_constant.cuh` — minimal runnable reference
|
||||
- `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh` — real example using `TensorMatcher` + `LaunchKernel` + `tile::Memory`
|
||||
- `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern
|
||||
- `python/sglang/jit_kernel/benchmark/utils.py` — benchmark helpers
|
||||
|
||||
## Summary of Files Created
|
||||
|
||||
```
|
||||
python/sglang/jit_kernel/csrc/elementwise/scale.cuh # NEW: CUDA kernel
|
||||
python/sglang/jit_kernel/scale.py # NEW: Python wrapper
|
||||
python/sglang/jit_kernel/tests/test_scale.py # NEW: Tests
|
||||
python/sglang/jit_kernel/benchmark/bench_scale.py # NEW: Benchmark
|
||||
```
|
||||
363
third_party/sglang/.claude/skills/add-sgl-kernel/SKILL.md
vendored
Normal file
363
third_party/sglang/.claude/skills/add-sgl-kernel/SKILL.md
vendored
Normal file
@@ -0,0 +1,363 @@
|
||||
---
|
||||
name: add-sgl-kernel
|
||||
description: Step-by-step tutorial for adding a heavyweight AOT CUDA/C++ kernel to sgl-kernel (including tests & benchmarks)
|
||||
---
|
||||
|
||||
# Tutorial: Adding a New Kernel to `sgl-kernel` (AOT / Heavyweight)
|
||||
|
||||
This tutorial walks through adding a simple element-wise scale operation as an AOT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow.
|
||||
|
||||
## Goal
|
||||
|
||||
Add a new operation that scales each element of a tensor by a scalar factor:
|
||||
|
||||
- Input: tensor `x` (CUDA) and scalar `factor` (float)
|
||||
- Output: `x * factor` (element-wise, in-place or into pre-allocated `out`)
|
||||
- Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)**
|
||||
- Dispatched via `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro (defined in `sgl-kernel/include/utils.h`)
|
||||
|
||||
## Two rules of thumb (must follow)
|
||||
|
||||
1. **Prefer `python/sglang/jit_kernel` first** when the kernel does **not** depend on CUTLASS or another large C++ project. This is the default path for lightweight kernels that benefit from rapid iteration.
|
||||
2. **Prefer `sgl-kernel`** when the kernel **does** depend on CUTLASS or another large C++ project, or when it should be part of the AOT wheel / torch op registration flow.
|
||||
3. **Exception**: if the dependency is `flashinfer`, or CUTLASS that is already provided through `flashinfer`, the kernel can still be implemented as `jit_kernel`.
|
||||
|
||||
In addition, every new kernel must ship with:
|
||||
|
||||
- **Tests** (pytest)
|
||||
- **A benchmark script** (triton.testing)
|
||||
|
||||
---
|
||||
|
||||
## Repository integration map
|
||||
|
||||
You will typically touch these files/areas:
|
||||
|
||||
- Implementation: `sgl-kernel/csrc/elementwise/scale.cu` (pick the right subdirectory)
|
||||
- Public declarations: `sgl-kernel/include/sgl_kernel_ops.h`
|
||||
- Torch extension registration: `sgl-kernel/csrc/common_extension.cc`
|
||||
- Build: `sgl-kernel/CMakeLists.txt` (`set(SOURCES ...)`)
|
||||
- Python API: `sgl-kernel/python/sgl_kernel/` and `sgl-kernel/python/sgl_kernel/__init__.py`
|
||||
- Tests: `sgl-kernel/tests/test_scale.py`
|
||||
- Benchmarks: `sgl-kernel/benchmark/bench_scale.py`
|
||||
|
||||
---
|
||||
|
||||
## Step 1: Implement the kernel in `csrc/`
|
||||
|
||||
Pick the right subdirectory:
|
||||
|
||||
- `csrc/elementwise/` — for element-wise ops (our example)
|
||||
- `csrc/gemm/`, `csrc/attention/`, `csrc/moe/` — for other categories
|
||||
|
||||
Create `sgl-kernel/csrc/elementwise/scale.cu`:
|
||||
|
||||
```cpp
|
||||
#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;
|
||||
});
|
||||
}
|
||||
```
|
||||
|
||||
**Key points:**
|
||||
|
||||
- Use `at::Tensor` (PyTorch tensors), `TORCH_CHECK` for validation, `at::cuda::getCurrentCUDAStream()` for stream
|
||||
- Keep Python wrappers thin; do shape/dtype/device validation in C++ right around the launch path
|
||||
- `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` covers `float`, `half` (FP16), `__nv_bfloat16` (BF16)
|
||||
- Add device error checking after every kernel launch
|
||||
- If a kernel only works on certain architectures, enforce that with `TORCH_CHECK` and skip logic in tests
|
||||
|
||||
---
|
||||
|
||||
## Step 2: Add a C++ declaration in `include/sgl_kernel_ops.h`
|
||||
|
||||
Edit `sgl-kernel/include/sgl_kernel_ops.h`, add to the elementwise section:
|
||||
|
||||
```cpp
|
||||
void scale(at::Tensor& out, const at::Tensor& input, double factor);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 3: Register the op in `csrc/common_extension.cc`
|
||||
|
||||
Edit `sgl-kernel/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`:
|
||||
|
||||
```cpp
|
||||
// From csrc/elementwise
|
||||
m.def("scale(Tensor! out, Tensor input, float factor) -> ()");
|
||||
m.impl("scale", torch::kCUDA, &scale);
|
||||
```
|
||||
|
||||
**Key points:**
|
||||
|
||||
- `Tensor!` means in-place / mutable output argument
|
||||
- The schema is important for `torch.compile` and for consistent call signatures
|
||||
- Keep the torch schema in PyTorch scalar types (`float` here), but note that the C++ launcher signature still needs `double` for scalar arguments accepted by `torch::Library`
|
||||
|
||||
---
|
||||
|
||||
## Step 4: Add the new source file to `CMakeLists.txt`
|
||||
|
||||
Edit `sgl-kernel/CMakeLists.txt`, add to `set(SOURCES ...)`:
|
||||
|
||||
```cmake
|
||||
csrc/elementwise/scale.cu
|
||||
```
|
||||
|
||||
**Key points:**
|
||||
|
||||
- Keep the list **alphabetically sorted** (the file explicitly requires this)
|
||||
- If the kernel has arch constraints, reflect that in tests/benchmarks via skip logic
|
||||
|
||||
---
|
||||
|
||||
## Step 5: Expose a Python API under `sgl-kernel/python/sgl_kernel/`
|
||||
|
||||
Prefer following the existing module organization first. For elementwise kernels, the usual pattern is:
|
||||
|
||||
- implement the Python wrapper in `sgl-kernel/python/sgl_kernel/elementwise.py`
|
||||
- then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py`
|
||||
|
||||
For example, in `sgl-kernel/python/sgl_kernel/elementwise.py`, add:
|
||||
|
||||
```python
|
||||
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
|
||||
```
|
||||
|
||||
Then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` following the existing import style used by other kernels.
|
||||
|
||||
---
|
||||
|
||||
## Step 6: Write tests (required)
|
||||
|
||||
Create `sgl-kernel/tests/test_scale.py`:
|
||||
```python
|
||||
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"]))
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 7: Add a benchmark (required)
|
||||
|
||||
Create `sgl-kernel/benchmark/bench_scale.py`:
|
||||
|
||||
```python
|
||||
import 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_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=["sglang", "torch"],
|
||||
line_names=["SGL Kernel", "PyTorch"],
|
||||
styles=[("green", "-"), ("red", "--")],
|
||||
ylabel="µs (median)",
|
||||
plot_name="scale-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(dtype, size, provider):
|
||||
input = torch.randn(size, dtype=dtype, device="cuda")
|
||||
out = torch.empty_like(input)
|
||||
factor = 2.0
|
||||
|
||||
if provider == "sglang":
|
||||
fn = lambda: sgl_kernel.scale(input, factor, out=out)
|
||||
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_data=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 8: Build
|
||||
|
||||
Build:
|
||||
|
||||
```bash
|
||||
cd sgl-kernel
|
||||
make build -j16
|
||||
```
|
||||
|
||||
If you need to limit host resource usage:
|
||||
|
||||
```bash
|
||||
cd sgl-kernel
|
||||
make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 9: Validate
|
||||
|
||||
After building successfully, run the test and benchmark:
|
||||
|
||||
```bash
|
||||
pytest sgl-kernel/tests/test_scale.py -q
|
||||
python sgl-kernel/benchmark/bench_scale.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **Async CUDA errors**: `CUDA_LAUNCH_BLOCKING=1`
|
||||
- **Memory errors**: `compute-sanitizer --tool memcheck python ...`
|
||||
- **Build is too slow / OOM**: reduce `MAX_JOBS` and `SGL_KERNEL_COMPILE_THREADS`
|
||||
- **Binary bloat**: use `sgl-kernel/analyze_whl_kernel_sizes.py`
|
||||
- **CMake sources list**: if your `.cu` file is missing from `SOURCES`, the symbol will be undefined at link time
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
- `sgl-kernel/README.md`
|
||||
- `sgl-kernel/include/sgl_kernel_ops.h`
|
||||
- `sgl-kernel/csrc/common_extension.cc`
|
||||
- `sgl-kernel/CMakeLists.txt`
|
||||
- `sgl-kernel/include/utils.h` — `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro and friends
|
||||
- `sgl-kernel/csrc/elementwise/activation.cu` — reference for the FP16/BF16/FP32 dispatch pattern
|
||||
|
||||
## Summary of Files Created/Modified
|
||||
|
||||
```
|
||||
sgl-kernel/csrc/elementwise/scale.cu # NEW: CUDA kernel + launcher
|
||||
sgl-kernel/include/sgl_kernel_ops.h # MODIFIED: C++ declaration
|
||||
sgl-kernel/csrc/common_extension.cc # MODIFIED: schema + dispatch registration
|
||||
sgl-kernel/CMakeLists.txt # MODIFIED: add source file (alphabetical)
|
||||
sgl-kernel/python/sgl_kernel/elementwise.py # MODIFIED: Python wrapper
|
||||
sgl-kernel/python/sgl_kernel/__init__.py # MODIFIED: re-export Python API
|
||||
sgl-kernel/tests/test_scale.py # NEW: tests
|
||||
sgl-kernel/benchmark/bench_scale.py # NEW: benchmark
|
||||
```
|
||||
386
third_party/sglang/.claude/skills/ci-workflow-guide/SKILL.md
vendored
Normal file
386
third_party/sglang/.claude/skills/ci-workflow-guide/SKILL.md
vendored
Normal file
@@ -0,0 +1,386 @@
|
||||
---
|
||||
name: ci-workflow-guide
|
||||
description: Guide to SGLang CI workflow orchestration — stage ordering, fast-fail, gating, partitioning, execution modes, and debugging CI failures. Use when modifying CI workflows, adding stages, debugging CI pipeline issues, or understanding how tests are dispatched and gated across stages.
|
||||
---
|
||||
|
||||
# SGLang CI Workflow Orchestration Guide
|
||||
|
||||
This skill covers the CI **infrastructure** layer — how tests are dispatched, gated, and fast-failed across stages. For test authoring (templates, fixtures, registration, model selection), see the [write-sglang-test skill](../write-sglang-test/SKILL.md).
|
||||
|
||||
---
|
||||
|
||||
## Naming Conventions
|
||||
|
||||
- **Suite**: `stage-{a,b,c}-test-{gpu_count}-gpu-{hardware}` (e.g., `stage-b-test-1-gpu-small`)
|
||||
- **CI runner**: `{gpu_count}-gpu-{hardware}` (e.g., `1-gpu-5090`, `4-gpu-h100`, `8-gpu-h200`)
|
||||
|
||||
---
|
||||
|
||||
## Key Files
|
||||
|
||||
| File | Role |
|
||||
|------|------|
|
||||
| `.github/workflows/pr-test.yml` | Main workflow — all stages, jobs, conditions, matrix definitions |
|
||||
| `.github/workflows/pr-gate.yml` | PR gating: draft check, `run-ci` label, per-user rate limiting |
|
||||
| `.github/actions/check-stage-health/action.yml` | Cross-job fast-fail: queries API for any failed job |
|
||||
| `.github/actions/wait-for-jobs/action.yml` | Stage gating: polls API until stage jobs complete |
|
||||
| `.github/actions/check-maintenance/action.yml` | Maintenance mode check |
|
||||
| `test/run_suite.py` | Suite runner: collects, filters, partitions, executes tests |
|
||||
| `python/sglang/test/ci/ci_register.py` | Test registration (AST-parsed markers), LPT auto-partition |
|
||||
| `python/sglang/test/ci/ci_utils.py` | `run_unittest_files()`: execution, retry, continue-on-error |
|
||||
| `scripts/ci/utils/slash_command_handler.py` | Handles slash commands from PR comments |
|
||||
|
||||
---
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
```
|
||||
┌──────────────┐
|
||||
│ build kernel │
|
||||
└──────┬───────┘
|
||||
│
|
||||
├─ check-changes ──── detects which packages changed
|
||||
│ (main_package, sgl_kernel, jit_kernel, multimodal_gen)
|
||||
│
|
||||
├─ call-gate ──────── pr-gate.yml (draft? label? rate limit?)
|
||||
│
|
||||
├─────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
▼ │
|
||||
┌─────────────────────────────────────┐ │
|
||||
│ Stage A (~3 min) │ │
|
||||
│ pre-flight check │ │
|
||||
│ │ │
|
||||
│ ┌─────────────────────────────┐ │ │
|
||||
│ │ stage-a-test-1-gpu-small │ │ │
|
||||
│ │ (small GPUs) │ │ │
|
||||
│ └─────────────────────────────┘ │ │
|
||||
│ ┌─────────────────────────────┐ │ │
|
||||
│ │ stage-a-test-cpu │ │ │
|
||||
│ │ (CPU) │ │ │
|
||||
│ └─────────────────────────────┘ │ │
|
||||
└──────┬──────────────────────────────┘ │
|
||||
│ │
|
||||
▼ ▼
|
||||
┌─────────────────────────────────────┐ ┌──────────────────────────┐
|
||||
│ Stage B (~30 min) │ │ kernel test │
|
||||
│ basic tests │ └──────────────────────────┘
|
||||
│ │ ┌──────────────────────────┐
|
||||
│ ┌─────────────────────────────┐ │ │ multimodal gen test │
|
||||
│ │ stage-b-test-1-gpu-small │ │ └──────────────────────────┘
|
||||
│ │ (small GPUs, e.g. 5090) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ stage-b-test-1-gpu-large │ │
|
||||
│ │ (large GPUs, e.g. H100) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ stage-b-test-2-gpu-large │ │
|
||||
│ │ (large GPUs, e.g. H100) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
└──────┬──────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────┐
|
||||
│ Stage C (~30 min) │
|
||||
│ advanced tests │
|
||||
│ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ stage-c-test-4-gpu-h100 │ │
|
||||
│ │ (H100 GPUs) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ stage-c-test-8-gpu-h200 │ │
|
||||
│ │ (8 x H200 GPUs) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ stage-c-test-4-gpu-b200 │ │
|
||||
│ │ (4 x B200 GPUs) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
│ ┌─────────────────────────────┐ │
|
||||
│ │ Other advanced tests │ │
|
||||
│ │ (DeepEP, PD Disagg, GB300) │ │
|
||||
│ └─────────────────────────────┘ │
|
||||
└──────┬──────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────┐
|
||||
│ pr-test-finish │
|
||||
│ aggregates all results, fails if │
|
||||
│ any job failed/cancelled │
|
||||
└─────────────────────────────────────┘
|
||||
```
|
||||
|
||||
**Every stage test job** includes a `check-stage-health` step after checkout — if any job in the run has already failed, the job fast-fails (red X) with a root cause annotation.
|
||||
|
||||
**Scheduled runs** skip `wait-for-stage-*` jobs, running all stages in parallel. Fast-fail is also disabled.
|
||||
|
||||
---
|
||||
|
||||
## Fast-Fail Layers
|
||||
|
||||
4 layers of fast-fail, from fine to coarse:
|
||||
|
||||
| Layer | Mechanism | Granularity | Disabled on schedule? |
|
||||
|-------|-----------|-------------|----------------------|
|
||||
| **1. Test method → file** | `unittest -f` (failfast) | One test method fails → entire test file stops immediately | Yes |
|
||||
| **2. File → suite** | `run_unittest_files()` default | One test file fails → entire suite stops (`--continue-on-error` off) | Yes |
|
||||
| **3. Job → job (same stage)** | `check-stage-health` action | One job fails → other waiting jobs in same stage fast-fail (red X) | Yes |
|
||||
| **4. Stage → stage (cross-stage)** | `wait-for-stage` + `needs` | Stage A fails → stage B/C jobs skip entirely (never get a runner) | Yes (wait jobs skipped) |
|
||||
|
||||
- **Layer 1**: `-f` flag appended to all `python3 -m pytest` / `unittest` invocations in `ci_utils.py`
|
||||
- **Layer 2**: `--continue-on-error` flag in `run_suite.py` — off for PRs, on for scheduled runs
|
||||
- **Layer 3**: `check-stage-health` auto-detects `schedule` event and skips; filters out cascade failures to show only root cause jobs
|
||||
- **Layer 4**: `wait-for-stage-*` jobs are conditioned on `github.event_name == 'pull_request'` — skipped for scheduled runs
|
||||
|
||||
---
|
||||
|
||||
## Execution Modes
|
||||
|
||||
| Aspect | PR (`pull_request`) | Scheduled (`cron`, every 6h) | `/rerun-stage` (`workflow_dispatch`) |
|
||||
|--------|---------------------|------------------------------|--------------------------------------|
|
||||
| **Stage ordering** | Sequential: A → B → C via `wait-for-stage-*` | Parallel (all at once) | Single target stage only |
|
||||
| **Cross-job fast-fail** | Yes (`check-stage-health`) | Yes | Yes |
|
||||
| **continue-on-error** | No (stop at first failure within suite) | Yes (run all tests) | No |
|
||||
| **Retry** | Enabled | Enabled | Enabled |
|
||||
| **max_parallel** | 3 (default), 14 if `high priority` label | 14 | 3 (default), 14 if `high priority` |
|
||||
| **PR gate** | Yes (draft, label, rate limit) | Skipped | Skipped |
|
||||
| **Concurrency** | `cancel-in-progress: true` per branch | Queue (no cancel) | Isolated per stage+SHA |
|
||||
|
||||
---
|
||||
|
||||
## Stage Gating (`wait-for-jobs` action)
|
||||
|
||||
`wait-for-stage-a` and `wait-for-stage-b` are lightweight `ubuntu-latest` jobs that poll the GitHub Actions API.
|
||||
|
||||
**How it works:**
|
||||
1. Calls `listJobsForWorkflowRun` to list all jobs in the current run
|
||||
2. Matches jobs by exact name or prefix (for matrix jobs, e.g., `stage-b-test-1-gpu-small (3)`)
|
||||
3. If any matched job has `conclusion === 'failure'` → fail immediately (fast-fail)
|
||||
4. If all matched jobs are completed and count matches `expected_count` → success
|
||||
5. Otherwise → sleep `poll-interval-seconds` (default: 60s) and retry
|
||||
6. Timeout after `max-wait-minutes` (240 min for stage-a, 480 min for stage-b)
|
||||
|
||||
**Job specs example** (stage-b):
|
||||
```json
|
||||
[
|
||||
{"prefix": "stage-b-test-1-gpu-small", "expected_count": 8},
|
||||
{"prefix": "stage-b-test-1-gpu-large", "expected_count": 14},
|
||||
{"prefix": "stage-b-test-2-gpu-large", "expected_count": 4},
|
||||
{"prefix": "stage-b-test-4-gpu-b200", "expected_count": 1}
|
||||
]
|
||||
```
|
||||
|
||||
> **Critical**: `expected_count` must match the matrix size. If you add/remove matrix entries, update the wait job's spec accordingly.
|
||||
|
||||
**PR only**: Condition `github.event_name == 'pull_request' && !inputs.target_stage` — scheduled runs and `/rerun-stage` skip these entirely, allowing parallel execution.
|
||||
|
||||
---
|
||||
|
||||
## Cross-Job Fast-Fail (`check-stage-health` action)
|
||||
|
||||
Composite action called after checkout in every stage test job (21 jobs total across `pr-test.yml`, `pr-test-multimodal-gen.yml`, `pr-test-sgl-kernel.yml`, `pr-test-jit-kernel.yml`).
|
||||
|
||||
**How it works:**
|
||||
1. Queries `listJobsForWorkflowRun` for the current workflow run
|
||||
2. Filters for **root cause failures only** — jobs with `conclusion === 'failure'` whose failing step is NOT `check-stage-health` (excludes cascade failures)
|
||||
3. If root cause failures found → calls `core.setFailed()` with the list of root cause job names
|
||||
4. If none → does nothing (step succeeds)
|
||||
|
||||
**Cascade filtering**: When job A fast-fails due to health check, it also has `conclusion: failure`. Without filtering, job B would list both the original failure AND job A's fast-fail. The filter checks each failed job's `steps` array — if the failing step name contains `check-stage-health` or `Check stage health`, it's excluded from the root cause list.
|
||||
|
||||
**Usage pattern:**
|
||||
```yaml
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
...
|
||||
|
||||
- uses: ./.github/actions/check-stage-health
|
||||
id: stage-health
|
||||
|
||||
- name: Install dependencies # skipped automatically if health check failed
|
||||
... # (default if: success() is false)
|
||||
|
||||
- name: Run test # also skipped
|
||||
...
|
||||
```
|
||||
|
||||
**Visual effect**: Job shows **red X** (failure) with error annotation showing root cause job names. Subsequent steps are naturally skipped (default `if: success()` is false after a failed step). No per-step `if` guards needed.
|
||||
|
||||
**No stage filtering**: Checks ALL jobs in the run, not just the current stage. Any failure anywhere triggers fast-fail.
|
||||
|
||||
**Error message example:**
|
||||
```
|
||||
Fast-fail: skipping — root cause job(s): stage-b-test-1-gpu-small (0), stage-b-test-1-gpu-small (1)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Within-Suite Failure Handling
|
||||
|
||||
Controlled by `run_unittest_files()` in `python/sglang/test/ci/ci_utils.py`.
|
||||
|
||||
### Flags
|
||||
|
||||
| Flag | PR default | Scheduled default | Effect |
|
||||
|------|------------|-------------------|--------|
|
||||
| `--continue-on-error` | Off | On | Off: stop at first failure. On: run all files, report all failures at end |
|
||||
| `--enable-retry` | On | On | Retry retriable failures (accuracy/perf assertions) |
|
||||
| `--max-attempts` | 2 | 2 | Max attempts per file including initial run |
|
||||
|
||||
### Retry Classification
|
||||
|
||||
When a test fails and retry is enabled, the output is classified:
|
||||
|
||||
**Non-retriable** (checked first — real code errors):
|
||||
`SyntaxError`, `ImportError`, `ModuleNotFoundError`, `NameError`, `TypeError`, `AttributeError`, `RuntimeError`, `CUDA out of memory`, `OOM`, `Segmentation fault`, `core dumped`, `ConnectionRefusedError`, `FileNotFoundError`
|
||||
|
||||
**Retriable** (accuracy/performance):
|
||||
`AssertionError` with comparison patterns (`not greater than`, `not less than`, `not equal to`), `accuracy`, `score`, `latency`, `throughput`, `timeout`
|
||||
|
||||
**Default**: Unknown `AssertionError` → retriable. Other unknown failures → not retriable.
|
||||
|
||||
### How `continue_on_error` is set
|
||||
|
||||
In `pr-test.yml`'s `check-changes` job:
|
||||
- `schedule` runs or `run_all_tests` flag → `continue_on_error = 'true'`
|
||||
- PR runs → `continue_on_error = 'false'`
|
||||
|
||||
Each test job propagates via:
|
||||
```yaml
|
||||
env:
|
||||
CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
|
||||
run: |
|
||||
python3 run_suite.py --hw cuda --suite <name> $CONTINUE_ON_ERROR_FLAG
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Test Partitioning
|
||||
|
||||
Large suites are split across matrix jobs using the **LPT (Longest Processing Time) heuristic** in `ci_register.py:auto_partition()`:
|
||||
|
||||
1. Sort tests by `est_time` descending, filename as tie-breaker (deterministic)
|
||||
2. Greedily assign each test to the partition with smallest cumulative time
|
||||
3. Result: roughly equal total time per partition
|
||||
|
||||
**Partition table** (CUDA per-commit suites):
|
||||
|
||||
| Suite | Partitions | Runner | max_parallel |
|
||||
|-------|-----------|--------|-------------|
|
||||
| `stage-a-test-1-gpu-small` | 1 (no matrix) | `1-gpu-5090` | — |
|
||||
| `stage-a-test-cpu` | 1 (no matrix) | `ubuntu-latest` | — |
|
||||
| `stage-b-test-1-gpu-small` | 8 | `1-gpu-5090` | 8 |
|
||||
| `stage-b-test-1-gpu-large` | 14 | `1-gpu-h100` | dynamic (3 or 14) |
|
||||
| `stage-b-test-2-gpu-large` | 4 | `2-gpu-h100` | — |
|
||||
| `stage-b-test-4-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — |
|
||||
| `stage-b-kernel-unit-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
|
||||
| `stage-b-kernel-unit-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — |
|
||||
| `stage-b-kernel-benchmark-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
|
||||
| `stage-c-test-4-gpu-h100` | 3 | `4-gpu-h100` | — |
|
||||
| `stage-c-test-8-gpu-h200` | 4 | `8-gpu-h200` | — |
|
||||
| `stage-c-test-8-gpu-h20` | 2 | `8-gpu-h20` | — |
|
||||
| `stage-c-test-deepep-4-gpu-h100` | 1 (no matrix) | `4-gpu-h100` | — |
|
||||
| `stage-c-test-deepep-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — |
|
||||
| `stage-c-test-4-gpu-b200` | 4 | `4-gpu-b200` | — |
|
||||
| `stage-c-test-4-gpu-gb200` | 1 (no matrix) | `4-gpu-gb200` | — |
|
||||
|
||||
> **Note**: Kernel suites (`stage-b-kernel-*`) run via `pr-test-jit-kernel.yml` and `pr-test-sgl-kernel.yml`, not the main `pr-test.yml`. Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`.
|
||||
|
||||
**Workflow usage:**
|
||||
```yaml
|
||||
strategy:
|
||||
matrix:
|
||||
partition: [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
steps:
|
||||
- run: python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-small \
|
||||
--auto-partition-id ${{ matrix.partition }} --auto-partition-size 8
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## check-changes Job
|
||||
|
||||
Determines which test suites to run based on file changes.
|
||||
|
||||
### Detection Methods
|
||||
|
||||
| Trigger | Method | Details |
|
||||
|---------|--------|---------|
|
||||
| `pull_request` | `dorny/paths-filter` | Detects changes via GitHub diff |
|
||||
| `workflow_dispatch` (with `pr_head_sha`) | GitHub API | `repos/{repo}/compare/main...{sha}` |
|
||||
| `schedule` / `run_all_tests` | Force all true | Runs everything |
|
||||
|
||||
### Output Flags
|
||||
|
||||
| Output | Triggers |
|
||||
|--------|----------|
|
||||
| `main_package` | Stage A/B/C test suites |
|
||||
| `sgl_kernel` | Kernel wheel builds + kernel test suites |
|
||||
| `jit_kernel` | JIT kernel test workflow |
|
||||
| `multimodal_gen` | Multimodal-gen test workflow |
|
||||
|
||||
> **Note**: `sgl_kernel` is forced to `false` when `target_stage` is set, because `sgl-kernel-build-wheels` won't run and wheel artifacts won't be available.
|
||||
|
||||
---
|
||||
|
||||
## Concurrency Control
|
||||
|
||||
```
|
||||
group: pr-test-{event_name}-{branch}-{pr_sha}-{stage}
|
||||
```
|
||||
|
||||
| Segment | Source | Purpose |
|
||||
|---------|--------|---------|
|
||||
| `event_name` | `github.event_name` | Prevents scheduled runs colliding with fork PRs named `main` |
|
||||
| `branch` | `github.head_ref \|\| github.ref_name` | Per-branch isolation |
|
||||
| `pr_sha` | `inputs.pr_head_sha \|\| 'current'` | Isolates `/rerun-stage` from main runs |
|
||||
| `stage` | `inputs.target_stage \|\| 'all'` | Allows parallel stage dispatches |
|
||||
|
||||
`cancel-in-progress: true` for `pull_request` events (new push cancels old run), `false` for `workflow_call`.
|
||||
|
||||
---
|
||||
|
||||
## How To: Add a New Stage Job
|
||||
|
||||
1. Define the job in `pr-test.yml` with `needs: [check-changes, call-gate, wait-for-stage-X, ...]`
|
||||
2. Copy the `if:` condition pattern from an existing same-stage job (handles `target_stage`, `schedule`, `main_package`)
|
||||
3. Add `checkout` step
|
||||
4. Add `check-stage-health` step (after checkout) — if any prior job failed, `core.setFailed()` fires and all subsequent steps auto-skip via default `if: success()`
|
||||
5. Add `check-maintenance` step
|
||||
6. Add `download-artifact` step if `sgl_kernel` changed
|
||||
7. Add `install dependencies` step
|
||||
8. Add `run test` step with `$CONTINUE_ON_ERROR_FLAG`
|
||||
9. Add `upload-cuda-coredumps` step with `if: always()`
|
||||
10. Register the suite name in `PER_COMMIT_SUITES` in `test/run_suite.py`
|
||||
11. If using matrix, add `--auto-partition-id` and `--auto-partition-size` to the run command
|
||||
12. **Update `wait-for-stage-X`** job spec with the new job name and `expected_count` (if matrix)
|
||||
13. **Add the job to `pr-test-finish.needs`** list
|
||||
|
||||
---
|
||||
|
||||
## How To: Debug CI Failures
|
||||
|
||||
| Symptom | Likely cause | What to check |
|
||||
|---------|-------------|---------------|
|
||||
| All stage-B/C jobs green but steps skipped | Earlier job failed, `check-stage-health` triggered | Find the actual failed job (red X) |
|
||||
| `wait-for-stage-b` timeout | `expected_count` doesn't match matrix size | Verify job spec counts match `matrix:` array length |
|
||||
| `pr-test-finish` fails but all jobs green | A job was `cancelled` (counts as failure in finish) | Check concurrency cancellation |
|
||||
| Tests pass locally but fail in CI | Partition assignment, runner GPU type, or `est_time` inaccuracy | Check which partition the test lands in; verify runner label |
|
||||
| Flaky test retried and passed | Retriable failure (accuracy/perf) | Check `[CI Retry]` markers in job logs |
|
||||
| Flaky test NOT retried | Matched non-retriable pattern | Check if error matches `NON_RETRIABLE_PATTERNS` in `ci_utils.py` |
|
||||
|
||||
---
|
||||
|
||||
## Slash Commands
|
||||
|
||||
| Command | Effect |
|
||||
|---------|--------|
|
||||
| `/tag-run-ci-label` | Adds `run-ci` label to PR |
|
||||
| `/rerun-failed-ci` | Reruns failed jobs in the latest workflow run |
|
||||
| `/tag-and-rerun-ci` | Adds label + reruns |
|
||||
| `/rerun-stage <stage>` | Dispatches `pr-test.yml` with `target_stage=<stage>` |
|
||||
| `/rerun-test <test-file>` | Reruns a specific test file via `rerun-test.yml` |
|
||||
|
||||
Handled by `scripts/ci/utils/slash_command_handler.py` → `.github/workflows/slash-command-handler.yml`.
|
||||
657
third_party/sglang/.claude/skills/debug-cuda-crash/SKILL.md
vendored
Normal file
657
third_party/sglang/.claude/skills/debug-cuda-crash/SKILL.md
vendored
Normal file
@@ -0,0 +1,657 @@
|
||||
---
|
||||
name: debug-cuda-crash
|
||||
description: Call this skill when you need to debug CUDA crashes in SGLang using kernel API logging
|
||||
---
|
||||
|
||||
# Tutorial: Debugging CUDA Crashes with Kernel API Logging
|
||||
|
||||
This tutorial shows you how to debug CUDA crashes and errors in SGLang using the `@debug_kernel_api` logging decorator.
|
||||
|
||||
## Goal
|
||||
|
||||
When your code crashes with CUDA errors such as illegal memory access, device-side assert, out-of-bounds, or NaN/Inf, use kernel API logging to:
|
||||
- Capture input tensors BEFORE the crash occurs
|
||||
- Understand what data caused the problem
|
||||
- Track tensor shapes, dtypes, and values through the call boundary that triggered the crash
|
||||
- Detect numerical issues such as NaN, Inf, or obviously wrong shapes
|
||||
|
||||
## Why Use Kernel API Logging?
|
||||
|
||||
**Problem**: CUDA errors often crash the program before normal debugging output is flushed.
|
||||
|
||||
**Solution**: SGLang's `@debug_kernel_api` decorator logs inputs before execution, so you can still see what caused the crash even after the program aborts.
|
||||
|
||||
## What Is Covered?
|
||||
|
||||
The current logging coverage focuses on the highest-value kernel boundaries in SGLang:
|
||||
- Custom ops registered through `register_custom_op(...)`
|
||||
- External custom ops registered through `register_custom_op_from_extern(...)`
|
||||
- LLM attention, linear, quantization, and multi-platform wrapper entry points
|
||||
- Diffusion attention impl, linear, rotary, and custom-op wrapper entry points
|
||||
- Selected direct `torch.ops.sglang.*` hotspots and model-specific bypasses
|
||||
|
||||
This means the logging is useful for both LLM and diffusion kernel debugging, but it does not automatically cover every pure PyTorch call in the repository.
|
||||
|
||||
## Step 1: Enable Kernel API Logging
|
||||
|
||||
### Basic Logging (Function Names Only)
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=1
|
||||
export SGLANG_KERNEL_API_LOGDEST=stdout
|
||||
|
||||
python my_script.py
|
||||
```
|
||||
|
||||
Output:
|
||||
```
|
||||
================================================================================
|
||||
[2026-03-19 00:47:06] SGLang Kernel API Call: RMSNorm.forward
|
||||
================================================================================
|
||||
[2026-03-19 00:47:06] SGLang Kernel API Call: sglang.quant_method.UnquantizedLinearMethod.apply
|
||||
================================================================================
|
||||
[2026-03-19 00:47:06] SGLang Kernel API Call: sglang.custom_op.fused_inplace_qknorm
|
||||
```
|
||||
|
||||
This is a real level-1 excerpt captured from `Qwen/Qwen3-0.6B`.
|
||||
|
||||
### Detailed Logging (Inputs with Metadata)
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
|
||||
python my_script.py
|
||||
```
|
||||
|
||||
Output in `debug.log`:
|
||||
```
|
||||
================================================================================
|
||||
[2026-03-19 00:47:30] SGLang Kernel API Call: sglang.quant_method.UnquantizedLinearMethod.apply
|
||||
Positional input arguments:
|
||||
arg[0]=QKVParallelLinear(
|
||||
repr=QKVParallelLinear(in_features=1024, output_features=4096, bias=False, tp_size=1, gather_output=False)
|
||||
)
|
||||
arg[1]=Tensor(
|
||||
shape=(1, 1024)
|
||||
dtype=torch.bfloat16
|
||||
device=cuda:0
|
||||
requires_grad=False
|
||||
is_contiguous=True
|
||||
)
|
||||
arg[2]=None
|
||||
Output:
|
||||
return=Tensor(
|
||||
shape=(1, 4096)
|
||||
dtype=torch.bfloat16
|
||||
device=cuda:0
|
||||
requires_grad=False
|
||||
is_contiguous=True
|
||||
)
|
||||
```
|
||||
|
||||
This is a real level-3 excerpt captured from `Qwen/Qwen3-0.6B`.
|
||||
|
||||
### Full Logging (With Tensor Statistics)
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=5
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
|
||||
python my_script.py
|
||||
```
|
||||
|
||||
Additional output:
|
||||
```
|
||||
================================================================================
|
||||
[2026-03-19 01:00:42] SGLang Kernel API Call: diffusion.quant_method.UnquantizedLinearMethod.apply
|
||||
Positional input arguments:
|
||||
arg[1]=Tensor(
|
||||
shape=(1, 77, 768)
|
||||
dtype=torch.bfloat16
|
||||
device=cuda:0
|
||||
requires_grad=False
|
||||
is_contiguous=True
|
||||
min=-27.250000
|
||||
max=28.500000
|
||||
mean=0.011723
|
||||
nan_count=0
|
||||
inf_count=0
|
||||
)
|
||||
Output:
|
||||
return=Tensor(
|
||||
shape=(1, 77, 2304)
|
||||
dtype=torch.bfloat16
|
||||
device=cuda:0
|
||||
requires_grad=False
|
||||
is_contiguous=True
|
||||
min=-8.937500
|
||||
max=9.375000
|
||||
mean=0.009460
|
||||
nan_count=0
|
||||
inf_count=0
|
||||
)
|
||||
```
|
||||
|
||||
This is a real level-5 excerpt captured from `black-forest-labs/FLUX.1-dev`.
|
||||
|
||||
### Crash-Safe Dumps (Inputs Saved Before Execution)
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=10
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps
|
||||
|
||||
python my_script.py
|
||||
```
|
||||
|
||||
At level 10, SGLang saves the inputs before execution. If the kernel crashes, the dump directory still contains the inputs and exception metadata.
|
||||
|
||||
If CUDA graph capture is active, tensor dumps are skipped automatically to avoid capture-time CUDA errors. In that case, you still get the kernel API call log, but not `inputs.pt` / `outputs.pt`.
|
||||
|
||||
Level-10 dumps are best understood as crash-safe call snapshots. They always preserve the observed call boundary. They do not guarantee one-click replay for every method, because some methods depend on module state that is not serialized into the dump.
|
||||
|
||||
Real level-10 dump layout from `Qwen/Qwen3-0.6B`:
|
||||
|
||||
```text
|
||||
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps
|
||||
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001
|
||||
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/inputs.pt
|
||||
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/metadata.json
|
||||
/tmp/sglang_kernel_api_validation/qwen_qwen3_0_6b_level10_dumps/20260319_004821_182_pid919286_RotaryEmbedding.forward_call0001/outputs.pt
|
||||
```
|
||||
|
||||
Real `metadata.json` excerpt:
|
||||
|
||||
```json
|
||||
{
|
||||
"function_name": "RotaryEmbedding.forward",
|
||||
"timestamp": "20260319_004821_182",
|
||||
"process_id": 919286,
|
||||
"execution_status": "completed",
|
||||
"input_tensor_keys": ["arg_0", "arg_1", "arg_2"],
|
||||
"output_tensor_keys": ["result_0", "result_1"]
|
||||
}
|
||||
```
|
||||
|
||||
## Step 2: Reproduce an LLM CUDA Crash
|
||||
|
||||
Create a temporary reproducer:
|
||||
|
||||
```bash
|
||||
python3 - <<'PY'
|
||||
from pathlib import Path
|
||||
Path("/tmp/sglang_llm_crash.py").write_text(
|
||||
"import torch\\n"
|
||||
"import torch.nn.functional as F\\n"
|
||||
"from sglang.srt.utils.custom_op import register_custom_op\\n\\n"
|
||||
"def _fake_embedding(indices, table):\\n"
|
||||
" return torch.empty((*indices.shape, table.shape[-1]), device=table.device, dtype=table.dtype)\\n\\n"
|
||||
"@register_custom_op(op_name='mock_llm_cuda_crash', fake_impl=_fake_embedding)\\n"
|
||||
"def mock_llm_cuda_crash(indices, table):\\n"
|
||||
" out = F.embedding(indices, table)\\n"
|
||||
" torch.cuda.synchronize()\\n"
|
||||
" return out\\n\\n"
|
||||
"table = torch.randn(4, 8, device='cuda', dtype=torch.float16)\\n"
|
||||
"indices = torch.tensor([0, 7], device='cuda', dtype=torch.long)\\n"
|
||||
"mock_llm_cuda_crash(indices, table)\\n"
|
||||
)
|
||||
PY
|
||||
|
||||
SGLANG_KERNEL_API_LOGLEVEL=1 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level1.log \
|
||||
python3 /tmp/sglang_llm_crash.py
|
||||
```
|
||||
|
||||
What to expect:
|
||||
- The script exits with a CUDA `device-side assert`
|
||||
- The log still contains the last API boundary before the crash
|
||||
|
||||
Try the same example at level 3:
|
||||
|
||||
```bash
|
||||
SGLANG_KERNEL_API_LOGLEVEL=3 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level3.log \
|
||||
python3 /tmp/sglang_llm_crash.py
|
||||
```
|
||||
|
||||
Now the log shows tensor metadata before the crash.
|
||||
|
||||
Try level 10:
|
||||
|
||||
```bash
|
||||
SGLANG_KERNEL_API_LOGLEVEL=10 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_llm_level10.log \
|
||||
SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_llm_level10_dumps \
|
||||
python3 /tmp/sglang_llm_crash.py
|
||||
```
|
||||
|
||||
Now you should see:
|
||||
- A log entry for `sglang.custom_op.mock_llm_cuda_crash`
|
||||
- A dump directory with `inputs.pt`
|
||||
- `metadata.json` showing `execution_status: "exception"`
|
||||
- No `outputs.pt`, because the kernel crashed before producing output
|
||||
|
||||
For real-model success-path level-10 dumps, it is often easier to temporarily disable CUDA graph and piecewise CUDA graph for the debug run.
|
||||
|
||||
## Step 3: Reproduce a Diffusion CUDA Crash
|
||||
|
||||
Create a temporary diffusion-side reproducer:
|
||||
|
||||
```bash
|
||||
python3 - <<'PY'
|
||||
from pathlib import Path
|
||||
Path("/tmp/sglang_diffusion_crash.py").write_text(
|
||||
"import torch\\n"
|
||||
"import torch.nn.functional as F\\n"
|
||||
"from sglang.multimodal_gen.runtime.layers.utils import register_custom_op\\n\\n"
|
||||
"def _fake_embedding(positions, cache):\\n"
|
||||
" return torch.empty((*positions.shape, cache.shape[-1]), device=cache.device, dtype=cache.dtype)\\n\\n"
|
||||
"@register_custom_op(op_name='mock_diffusion_cuda_crash', fake_impl=_fake_embedding)\\n"
|
||||
"def mock_diffusion_cuda_crash(positions, cache):\\n"
|
||||
" out = F.embedding(positions, cache)\\n"
|
||||
" torch.cuda.synchronize()\\n"
|
||||
" return out\\n\\n"
|
||||
"cache = torch.randn(4, 64, device='cuda', dtype=torch.float16)\\n"
|
||||
"positions = torch.tensor([0, 9], device='cuda', dtype=torch.long)\\n"
|
||||
"mock_diffusion_cuda_crash(positions, cache)\\n"
|
||||
)
|
||||
PY
|
||||
|
||||
SGLANG_KERNEL_API_LOGLEVEL=1 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level1.log \
|
||||
python3 /tmp/sglang_diffusion_crash.py
|
||||
```
|
||||
|
||||
Try level 3:
|
||||
|
||||
```bash
|
||||
SGLANG_KERNEL_API_LOGLEVEL=3 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level3.log \
|
||||
python3 /tmp/sglang_diffusion_crash.py
|
||||
```
|
||||
|
||||
Try level 10:
|
||||
|
||||
```bash
|
||||
SGLANG_KERNEL_API_LOGLEVEL=10 \
|
||||
SGLANG_KERNEL_API_LOGDEST=/tmp/sglang_diffusion_level10.log \
|
||||
SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_diffusion_level10_dumps \
|
||||
python3 /tmp/sglang_diffusion_crash.py
|
||||
```
|
||||
|
||||
If your local environment has unrelated FlashInfer import issues, resolve them in the shell before running the example. The example itself does not set any `FLASHINFER_*` environment variable.
|
||||
|
||||
## Step 4: Multi-Process Debugging
|
||||
|
||||
When running with multiple GPUs or worker processes, use `%i` in the log path:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug_rank_%i.log
|
||||
|
||||
torchrun --nproc_per_node=4 my_script.py
|
||||
```
|
||||
|
||||
This creates separate logs such as:
|
||||
- `debug_rank_12345.log`
|
||||
- `debug_rank_12346.log`
|
||||
- `debug_rank_12347.log`
|
||||
- `debug_rank_12348.log`
|
||||
|
||||
Real multi-process example from a 2-GPU `Qwen/Qwen2.5-0.5B-Instruct` run:
|
||||
|
||||
```text
|
||||
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950201.log
|
||||
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950349.log
|
||||
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950350.log
|
||||
/tmp/sglang_kernel_api_validation_multi/qwen_qwen2_5_0_5b_instruct_level3_950351.log
|
||||
```
|
||||
|
||||
You should usually do the same for level-10 dump directories:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=10
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug_rank_%i.log
|
||||
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps_%i
|
||||
```
|
||||
|
||||
This avoids multiple ranks writing into the same dump directory tree.
|
||||
|
||||
## Step 5: Filter Level-10 Dumps
|
||||
|
||||
If level 10 is too noisy, restrict dumps to specific APIs:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=10
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
export SGLANG_KERNEL_API_DUMP_DIR=/tmp/sglang_kernel_api_dumps
|
||||
export SGLANG_KERNEL_API_DUMP_INCLUDE='sglang.custom_op.*'
|
||||
export SGLANG_KERNEL_API_DUMP_EXCLUDE='*.fake_impl'
|
||||
```
|
||||
|
||||
`SGLANG_KERNEL_API_DUMP_INCLUDE` and `SGLANG_KERNEL_API_DUMP_EXCLUDE` use shell-style wildcard matching.
|
||||
|
||||
## Step 6: Common CUDA Errors and What to Check
|
||||
|
||||
### Illegal Memory Access or Device-Side Assert
|
||||
|
||||
**Typical errors**:
|
||||
```
|
||||
RuntimeError: CUDA error: an illegal memory access was encountered
|
||||
torch.AcceleratorError: CUDA error: device-side assert triggered
|
||||
```
|
||||
|
||||
Use:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
```
|
||||
|
||||
Check in the logs:
|
||||
- ✅ Tensor shapes
|
||||
- ✅ Tensor dtypes
|
||||
- ✅ CUDA vs CPU device placement
|
||||
- ✅ Tensor stride / contiguity
|
||||
- ✅ Whether the failing call has inputs logged but no outputs logged
|
||||
|
||||
Typical shape-mismatch pattern:
|
||||
|
||||
```text
|
||||
SGLang Kernel API Call: ...
|
||||
arg[0]=Tensor(shape=(..., 128), ...) # ✅ expected dimension
|
||||
arg[1]=Tensor(shape=(..., 64), ...) # ❌ mismatch
|
||||
```
|
||||
|
||||
This often points to head-dim, hidden-dim, or cache-layout mismatch rather than a random CUDA failure.
|
||||
|
||||
### NaN or Inf
|
||||
|
||||
Use:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=5
|
||||
```
|
||||
|
||||
Check:
|
||||
- `min`
|
||||
- `max`
|
||||
- `mean`
|
||||
- `nan_count`
|
||||
- `inf_count`
|
||||
|
||||
Typical bad pattern:
|
||||
|
||||
```text
|
||||
Tensor(
|
||||
...
|
||||
min=-1234567.000000 # ❌ suspiciously large
|
||||
max=9876543.000000 # ❌ suspiciously large
|
||||
mean=nan # ❌ bad
|
||||
nan_count=128 # ❌ found NaNs
|
||||
inf_count=0 # ✅ no Infs here
|
||||
)
|
||||
```
|
||||
|
||||
This usually means the bad values were already present before the crashing kernel.
|
||||
|
||||
### Out of Memory
|
||||
|
||||
Use:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
```
|
||||
|
||||
Check:
|
||||
- Unexpectedly large tensor shapes
|
||||
- Batch size
|
||||
- Sequence length
|
||||
- Frame count or image resolution in diffusion workloads
|
||||
|
||||
Also check whether a supposedly per-token or per-frame tensor accidentally became full-sequence or full-image sized.
|
||||
|
||||
Typical bad pattern:
|
||||
|
||||
```text
|
||||
Tensor(
|
||||
shape=(1024, 8192, 128, 128) # ❌ way too large
|
||||
...
|
||||
)
|
||||
```
|
||||
|
||||
### Example: Spot a Shape Bug from the Log
|
||||
|
||||
Suppose the failing API log looks like this:
|
||||
|
||||
```text
|
||||
[2026-03-19 00:47:30] SGLang Kernel API Call: RotaryEmbedding.forward
|
||||
Positional input arguments:
|
||||
arg[0]=Tensor(shape=(1, 8), dtype=torch.int64, ...)
|
||||
arg[1]=Tensor(shape=(1, 8, 8, 256), dtype=torch.bfloat16, ...) # ✅ query
|
||||
arg[2]=Tensor(shape=(1, 8, 4, 64), dtype=torch.bfloat16, ...) # ❌ key head_dim mismatch
|
||||
```
|
||||
|
||||
What this tells you:
|
||||
- ✅ positions look reasonable
|
||||
- ✅ query looks plausible
|
||||
- ❌ key last dimension is inconsistent with the expected rotary/head dimension
|
||||
|
||||
That usually means the bug is in projection layout, head packing, or cache format rather than in the rotary kernel itself.
|
||||
|
||||
## Step 7: Combine with compute-sanitizer
|
||||
|
||||
For harder bugs, combine kernel API logging with CUDA memory checking:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
|
||||
compute-sanitizer --tool memcheck python3 /tmp/sglang_llm_crash.py
|
||||
```
|
||||
|
||||
Use `debug.log` to see the exact inputs that reached the crashing API boundary.
|
||||
|
||||
Typical `compute-sanitizer` output:
|
||||
|
||||
```text
|
||||
========= COMPUTE-SANITIZER
|
||||
========= Invalid __global__ write of size 4 bytes
|
||||
========= at 0x1234 in SomeKernel
|
||||
========= by thread (256,0,0) in block (10,0,0)
|
||||
========= Address 0x... is out of bounds
|
||||
```
|
||||
|
||||
Use the sanitizer output to identify the failing kernel and use `debug.log` to identify the exact tensors that reached the API boundary right before it.
|
||||
|
||||
If you need more synchronous host-side error reporting, you can try `CUDA_LAUNCH_BLOCKING=1` as a separate follow-up experiment. It is not part of the default workflow because it changes execution timing and can hide concurrency-related behavior.
|
||||
|
||||
## Step 8: Combine with cuda-gdb
|
||||
|
||||
For crashes that need a stack trace instead of only memory diagnostics:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
export SGLANG_KERNEL_API_LOGDEST=debug.log
|
||||
|
||||
cuda-gdb --args python3 /tmp/sglang_llm_crash.py
|
||||
```
|
||||
|
||||
Inside `cuda-gdb`:
|
||||
|
||||
```text
|
||||
(cuda-gdb) run
|
||||
(cuda-gdb) where
|
||||
```
|
||||
|
||||
Then correlate the backtrace with `debug.log`.
|
||||
|
||||
## Step 9: Kernel-Level Debugging with printf()
|
||||
|
||||
When you own the CUDA kernel, `printf()` is still useful for narrowing down bad indices, bad launch geometry, or broken state propagation.
|
||||
|
||||
Basic pattern:
|
||||
|
||||
```cpp
|
||||
__global__ void MyKernel(const float* input, float* output, int n) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
printf("n=%d input0=%f\n", n, input[0]);
|
||||
}
|
||||
|
||||
if (idx < n) {
|
||||
output[idx] = input[idx] * 2.0f;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
After launch, force the output to flush:
|
||||
|
||||
```python
|
||||
my_kernel(...)
|
||||
torch.cuda.synchronize()
|
||||
```
|
||||
|
||||
For warp-specialized kernels, do not blindly print only on `threadIdx.x == 0`. Pick one representative thread per warp or per specialization group instead.
|
||||
|
||||
### Warp-Specialized Kernels: Choosing the Right Print Thread
|
||||
|
||||
Problem:
|
||||
- `threadIdx.x == 0` only prints from the first warp in the block
|
||||
- for warp-specialized kernels, that often misses the warp or group that is actually wrong
|
||||
|
||||
Better pattern:
|
||||
|
||||
```cpp
|
||||
__global__ void WarpSpecializedKernel(...) {
|
||||
// Example: first lane of each warp
|
||||
if ((threadIdx.x % 32) == 0) {
|
||||
printf("warp=%d\n", threadIdx.x / 32);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Or, if the kernel is organized in larger specialization groups, print once per group instead of once per block.
|
||||
|
||||
Common mistake:
|
||||
|
||||
```cpp
|
||||
// Only warp 0 prints
|
||||
if (threadIdx.x == 0) {
|
||||
printf("warp=%d\n", threadIdx.x / 32);
|
||||
}
|
||||
```
|
||||
|
||||
### Quick Reference
|
||||
|
||||
| Kernel Type | Print Condition | Notes |
|
||||
|----------|----------|-------------|
|
||||
| Simple kernel | `threadIdx.x == 0` | One thread per block is usually enough |
|
||||
| Warp-specialized kernel | one representative lane per warp | e.g. `threadIdx.x % 32 == 0` |
|
||||
| Group-specialized kernel | one representative lane per group | choose based on the kernel's scheduling layout |
|
||||
|
||||
### Other Kernel Debugging Tools
|
||||
|
||||
```cpp
|
||||
assert(value >= 0.0f && "value must be non-negative");
|
||||
static_assert(BLOCK_SIZE % 32 == 0, "BLOCK_SIZE must be warp aligned");
|
||||
```
|
||||
|
||||
## Environment Variables Reference
|
||||
|
||||
| Variable | Values | Description |
|
||||
|----------|--------|-------------|
|
||||
| `SGLANG_KERNEL_API_LOGLEVEL` | `0` | No logging (default) |
|
||||
| | `1` | Function names only |
|
||||
| | `3` | Inputs and outputs with metadata |
|
||||
| | `5` | Level 3 plus tensor statistics |
|
||||
| | `10` | Level 5 plus crash-safe tensor dumps |
|
||||
| `SGLANG_KERNEL_API_LOGDEST` | `stdout` | Log to stdout |
|
||||
| | `stderr` | Log to stderr |
|
||||
| | `<path>` | Log to file |
|
||||
| | `log_%i.txt` | `%i` expands to process ID |
|
||||
| `SGLANG_KERNEL_API_DUMP_DIR` | `<path>` | Directory for level-10 dumps |
|
||||
| `SGLANG_KERNEL_API_DUMP_INCLUDE` | wildcard list | Only dump matching API names |
|
||||
| `SGLANG_KERNEL_API_DUMP_EXCLUDE` | wildcard list | Skip matching API names |
|
||||
|
||||
## Best Practices
|
||||
|
||||
### 1. Start with Level 3
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
```
|
||||
|
||||
Level 3 is usually enough to catch wrong shapes, wrong dtypes, and wrong devices.
|
||||
|
||||
### 2. Use Level 5 for Numerical Issues
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=5
|
||||
```
|
||||
|
||||
Use it when you suspect NaN or Inf values.
|
||||
|
||||
### 3. Use Level 10 for Crash Reproduction
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=10
|
||||
```
|
||||
|
||||
This is the most useful mode when the process crashes before you can inspect live tensors.
|
||||
|
||||
If you need successful input/output dumps from a real model run, temporarily disable CUDA graph for that debug session.
|
||||
|
||||
When level 10 is too noisy, pair it with `SGLANG_KERNEL_API_DUMP_INCLUDE` / `SGLANG_KERNEL_API_DUMP_EXCLUDE` instead of dumping every covered API.
|
||||
|
||||
### 4. Log to File for Crashes
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGDEST=crash.log
|
||||
```
|
||||
|
||||
File logs are safer than stdout when the process aborts.
|
||||
|
||||
### 5. Disable Logging in Production
|
||||
|
||||
```bash
|
||||
unset SGLANG_KERNEL_API_LOGLEVEL
|
||||
```
|
||||
|
||||
When disabled, the decorator returns the original callable and adds no runtime logging overhead.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### No Logs Appear
|
||||
|
||||
Check:
|
||||
1. `echo $SGLANG_KERNEL_API_LOGLEVEL`
|
||||
2. `echo $SGLANG_KERNEL_API_LOGDEST`
|
||||
3. Whether the failing path goes through a covered API boundary
|
||||
|
||||
### Too Much Output
|
||||
|
||||
Reduce the level:
|
||||
|
||||
```bash
|
||||
export SGLANG_KERNEL_API_LOGLEVEL=3
|
||||
```
|
||||
|
||||
### Statistics Are Skipped During CUDA Graph Capture
|
||||
|
||||
If you see:
|
||||
```text
|
||||
statistics=[skipped: CUDA graph capture in progress]
|
||||
```
|
||||
|
||||
That is expected. Level-5 statistics are intentionally skipped during CUDA graph capture to avoid synchronization side effects.
|
||||
|
||||
### Tensor Dumps Are Skipped During CUDA Graph Capture
|
||||
|
||||
If you see:
|
||||
```text
|
||||
Tensor dump skipped: CUDA graph capture in progress
|
||||
```
|
||||
|
||||
That is also expected. Level-10 dumps require copying tensors to CPU, which is not allowed during CUDA graph capture.
|
||||
141
third_party/sglang/.claude/skills/generate-profile/SKILL.md
vendored
Normal file
141
third_party/sglang/.claude/skills/generate-profile/SKILL.md
vendored
Normal file
@@ -0,0 +1,141 @@
|
||||
---
|
||||
name: generate-profile
|
||||
description: Generate an e2e profiling trace of an SGLang server run. Launches a server, validates accuracy, captures a Chrome-compatible trace, and returns the profile path.
|
||||
---
|
||||
|
||||
# Generate an E2E Profile of an SGLang Server Run
|
||||
|
||||
This skill launches an SGLang server, validates it with a quick accuracy test, generates a profiling trace, and returns the profile file path.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- A working SGLang installation (`pip install -e .` or equivalent)
|
||||
- At least one available CUDA GPU
|
||||
|
||||
## Step-by-step Workflow
|
||||
|
||||
### Step 1: Launch the server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=<gpu_id> sglang serve --model-path <model> --port <port> &
|
||||
```
|
||||
|
||||
- Default model: `Qwen/Qwen3-8B` (good balance of speed and quality)
|
||||
- Default port: `30000`
|
||||
- The server runs in the background. Save the PID for cleanup.
|
||||
- Use the GPU specified by the user's preferences (check memory files for GPU preferences).
|
||||
|
||||
### Step 2: Wait for server readiness
|
||||
|
||||
Poll the health endpoint until the server is ready:
|
||||
|
||||
```bash
|
||||
for i in $(seq 1 120); do
|
||||
if curl -s http://127.0.0.1:<port>/health 2>/dev/null | grep -q "ok\|healthy"; then
|
||||
echo "Server ready"
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
```
|
||||
|
||||
The server prints **"The server is fired up and ready to roll!"** to stdout when ready. The health endpoint returns 200 once the server can accept requests.
|
||||
|
||||
Typical startup time: 30-90 seconds depending on model size and whether CUDA graphs are being compiled.
|
||||
|
||||
### Step 3: Validate accuracy (sanity check)
|
||||
|
||||
```bash
|
||||
python3 -m sglang.test.few_shot_gsm8k --num-q 20
|
||||
```
|
||||
|
||||
- Expected accuracy: **> 0.8** for capable models (Qwen3-8B, Llama-3.1-8B-Instruct, etc.)
|
||||
- This is a quick sanity check, not a rigorous benchmark.
|
||||
- If accuracy is unexpectedly low, something is wrong — do not proceed to profiling.
|
||||
|
||||
### Step 4: Generate the profile
|
||||
|
||||
```bash
|
||||
python3 -m sglang.test.send_one --profile
|
||||
```
|
||||
|
||||
This command:
|
||||
1. Sends a request to the server
|
||||
2. Triggers the profiler for 5 steps (default)
|
||||
3. Generates a trace file under `/tmp/<timestamp>/`
|
||||
4. The trace directory contains:
|
||||
- `<timestamp>-TP-0.trace.json.gz` — Chrome trace format (open in `chrome://tracing` or Perfetto)
|
||||
- `server_args.json` — the server configuration used
|
||||
|
||||
**Output format:**
|
||||
```
|
||||
Dump profiling traces to /tmp/<timestamp>
|
||||
```
|
||||
|
||||
The profile path is printed to stdout. Parse it from the output.
|
||||
|
||||
**Optional flags:**
|
||||
- `--profile-steps N` — number of profiling steps (default: 5)
|
||||
- `--profile-by-stage` — profile by stage (prefill/decode separately)
|
||||
- `--profile-prefix <path>` — custom output prefix
|
||||
|
||||
### Step 5: Kill the server
|
||||
|
||||
```bash
|
||||
pkill -9 -f "sglang.launch_server\|sglang serve\|sglang.srt"
|
||||
```
|
||||
|
||||
Wait a moment and verify no sglang processes remain:
|
||||
```bash
|
||||
sleep 2 && pgrep -af "sglang serve" || echo "Server killed"
|
||||
```
|
||||
|
||||
### Step 6: Report the profile path
|
||||
|
||||
Return the profile directory path (e.g., `/tmp/1773999986.4769795`) and list its contents so the user knows what files were generated.
|
||||
|
||||
## Example Full Run
|
||||
|
||||
```bash
|
||||
# 1. Launch server
|
||||
source cleanup/bin/activate
|
||||
CUDA_VISIBLE_DEVICES=1 sglang serve --model-path Qwen/Qwen3-8B --port 30000 &
|
||||
|
||||
# 2. Wait for ready
|
||||
for i in $(seq 1 120); do
|
||||
curl -s http://127.0.0.1:30000/health | grep -q "ok" && break
|
||||
sleep 5
|
||||
done
|
||||
|
||||
# 3. Accuracy check
|
||||
python3 -m sglang.test.few_shot_gsm8k --num-q 20
|
||||
# Expected: Accuracy > 0.8
|
||||
|
||||
# 4. Profile
|
||||
python3 -m sglang.test.send_one --profile
|
||||
# Output: "Dump profiling traces to /tmp/1773999986.4769795"
|
||||
|
||||
# 5. Cleanup
|
||||
pkill -9 -f "sglang.launch_server\|sglang serve\|sglang.srt"
|
||||
sleep 2
|
||||
|
||||
# 6. Check output
|
||||
ls -la /tmp/1773999986.4769795/
|
||||
# 1773999986.4851577-TP-0.trace.json.gz (Chrome trace)
|
||||
# server_args.json (server config)
|
||||
```
|
||||
|
||||
## Customization
|
||||
|
||||
- **Different port**: Pass `--port <port>` and use `--host 127.0.0.1 --port <port>` for test commands
|
||||
- **Multi-GPU**: Use `--tp <N>` for tensor parallelism; trace files will be generated per TP rank
|
||||
- **Longer profile**: Use `--profile-steps 10` for more steps in the trace
|
||||
- **Stage profiling**: Use `--profile-by-stage` to separate prefill and decode phases
|
||||
|
||||
## Viewing the Profile
|
||||
|
||||
Open the `.trace.json.gz` file in:
|
||||
- **Perfetto UI**: https://ui.perfetto.dev/ (drag and drop the file)
|
||||
- **Chrome tracing**: `chrome://tracing` (load the file)
|
||||
|
||||
Both support the gzipped Chrome trace format natively.
|
||||
219
third_party/sglang/.claude/skills/sglang-bisect-ci-regression/SKILL.md
vendored
Normal file
219
third_party/sglang/.claude/skills/sglang-bisect-ci-regression/SKILL.md
vendored
Normal file
@@ -0,0 +1,219 @@
|
||||
# SGLang Bisect CI Regression
|
||||
|
||||
Investigate a consistently failing CI test to find the root cause - whether it's a code regression from a specific PR, a hardware/runner-specific issue, or an environment change. Optionally reproduce the failure on a remote GPU server.
|
||||
|
||||
## Slash Command
|
||||
|
||||
`/sglang-bisect-ci-regression <test_name_or_ci_url> [ssh_target] [docker_container]`
|
||||
|
||||
## When to Use This Skill
|
||||
|
||||
- A CI test is failing consistently on main (scheduled runs)
|
||||
- You need to find which PR introduced a regression
|
||||
- You suspect a runner-specific or GPU-specific issue
|
||||
- You want to reproduce a CI failure on a remote server
|
||||
|
||||
## Arguments
|
||||
|
||||
- **First argument (required)**: Test file name (e.g. `test_lora_tp.py`) or a GitHub Actions job URL
|
||||
- **Second argument (optional)**: SSH target for remote reproduction (e.g. `user@host`)
|
||||
- **Third argument (optional)**: Docker container name on the SSH target (e.g. `sglang_dev`)
|
||||
|
||||
If SSH target and docker container are not provided, the skill will only perform the CI log analysis and bisection, without remote reproduction. **Ask the user** for these if reproduction is needed and they weren't provided.
|
||||
|
||||
## Background: Scheduled CI Runs
|
||||
|
||||
SGLang uses the `pr-test.yml` workflow with **scheduled runs** (cron-triggered) to periodically test the `main` branch. These runs are the primary data source for detecting regressions:
|
||||
|
||||
- **Workflow**: `pr-test.yml` with `event: schedule`
|
||||
- **Branch**: `main`
|
||||
- **Dashboard**: https://github.com/sgl-project/sglang/actions/workflows/pr-test.yml?query=event%3Aschedule
|
||||
- **Frequency**: Runs multiple times daily, each pinned to the HEAD of `main` at trigger time
|
||||
- **Purpose**: Catches regressions that slip through PR-level CI (e.g., interaction bugs between merged PRs, hardware-specific issues)
|
||||
|
||||
Always use these scheduled runs (not PR-triggered runs) when bisecting regressions on `main`. The `--event schedule` filter in `gh run list` ensures you only see these periodic main-branch runs.
|
||||
|
||||
## Workflow
|
||||
|
||||
### Phase 1: Extract the Failure Signature
|
||||
|
||||
1. **Get the failing test details from CI logs.** If given a URL, fetch logs directly. If given a test name, find recent scheduled runs of `pr-test.yml` on `main` that failed:
|
||||
|
||||
```bash
|
||||
# List recent scheduled runs targeting main (the primary source of truth for regressions)
|
||||
# These are cron-triggered runs visible at:
|
||||
# https://github.com/sgl-project/sglang/actions/workflows/pr-test.yml?query=event%3Aschedule
|
||||
gh run list --repo sgl-project/sglang --workflow="pr-test.yml" --event schedule --branch main --limit 20 --json databaseId,conclusion,createdAt,headSha
|
||||
|
||||
# Find the job containing the test
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --json jobs --jq '.jobs[] | select(.conclusion == "failure") | {name, conclusion, databaseId}'
|
||||
|
||||
# Get the failure details
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E -B 5 -A 30 "AssertionError|FAIL|Error|{TEST_NAME}"
|
||||
```
|
||||
|
||||
2. **Record the failure signature:**
|
||||
- Exact error message and assertion
|
||||
- Affected test method name
|
||||
- Model/config involved
|
||||
- Numeric values (e.g., tolerance diffs, scores)
|
||||
- Whether the failure is deterministic (same values across runs)
|
||||
|
||||
### Phase 2: Temporal Bisection
|
||||
|
||||
3. **Find the boundary between passing and failing runs.** Walk through the scheduled run history (from the `pr-test.yml` schedule runs on `main`) to identify:
|
||||
- Last known PASSING run (sha + date)
|
||||
- First known FAILING run (sha + date)
|
||||
|
||||
```bash
|
||||
# For each scheduled run, check the specific partition/job status
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --json jobs --jq '.jobs[] | select(.name == "{JOB_NAME}") | {conclusion, databaseId}'
|
||||
|
||||
# Verify a specific test passed or failed in a run
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "{TEST_NAME}|PASSED|FAILED|logprobs mismatch" | head -10
|
||||
```
|
||||
|
||||
4. **List commits between the boundary:**
|
||||
|
||||
```bash
|
||||
git log --oneline {LAST_PASS_SHA}..{FIRST_FAIL_SHA}
|
||||
```
|
||||
|
||||
5. **Filter for relevant commits** that touch files related to the failing test (model layers, kernels, test utilities, etc.):
|
||||
|
||||
```bash
|
||||
git log --oneline {LAST_PASS_SHA}..{FIRST_FAIL_SHA} -- {relevant_paths}
|
||||
```
|
||||
|
||||
### Phase 3: Runner/Hardware Analysis
|
||||
|
||||
6. **Check if the failure is runner-specific.** Extract the runner identity from each failing and passing run:
|
||||
|
||||
```bash
|
||||
# Get runner name and machine
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "Runner name|Machine name" | head -5
|
||||
|
||||
# Get GPU/driver info
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -i -E "NVIDIA-SMI|Driver Version|CUDA Version" | head -5
|
||||
|
||||
# Get package versions
|
||||
gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "sgl.kernel.*==|flashinfer.*==" | head -5
|
||||
```
|
||||
|
||||
7. **Correlate runners with pass/fail outcomes.** Build a table:
|
||||
|
||||
| Run ID | Date | Runner | GPU Type | Driver | Result |
|
||||
|--------|------|--------|----------|--------|--------|
|
||||
|
||||
If all failures map to a specific runner type/GPU and all passes map to another, the issue is **hardware-specific**, not a code regression.
|
||||
|
||||
### Phase 4: Code Analysis
|
||||
|
||||
8. **If a code regression is suspected** (failures not runner-specific), examine the candidate commits:
|
||||
- Read the changed files
|
||||
- Understand how the changes could affect the failing test
|
||||
- Look for prefill-vs-decode differences, TP-specific paths, kernel changes
|
||||
|
||||
9. **If a hardware issue is suspected**, analyze:
|
||||
- Kernel compatibility (CUDA compute capability)
|
||||
- Driver version differences
|
||||
- All-reduce / NCCL behavior differences
|
||||
- CUDA graph capture differences across GPU architectures
|
||||
|
||||
### Phase 5: Remote Reproduction (Optional)
|
||||
|
||||
Only if SSH target and docker container were provided.
|
||||
|
||||
10. **Verify the remote environment:**
|
||||
|
||||
```bash
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} nvidia-smi --query-gpu=name,driver_version --format=csv"
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} pip show sgl-kernel sglang flashinfer-python 2>&1 | grep -E 'Name:|Version:'"
|
||||
```
|
||||
|
||||
11. **Ensure latest code is installed.** If the container is stale, update:
|
||||
|
||||
```bash
|
||||
# Try fetching latest main
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /path/to/sglang && git fetch origin main && git checkout origin/main'"
|
||||
# Or download and install from tarball if git auth fails
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /tmp && curl -L https://github.com/sgl-project/sglang/archive/refs/heads/main.tar.gz | tar xz && cd sglang-main && pip install -e \"python[all]\"'"
|
||||
# Reinstall (after git fetch)
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /path/to/sglang && pip install -e \"python[all]\"'"
|
||||
# Install test dependencies if needed
|
||||
ssh {SSH_TARGET} "docker exec {CONTAINER} pip install peft rouge-score"
|
||||
```
|
||||
|
||||
12. **Create a minimal reproduction script** that:
|
||||
- Uses `if __name__ == '__main__'` with `mp.set_start_method("spawn")`
|
||||
- Runs the specific failing test configuration
|
||||
- Prints key metrics (diffs, scores, outputs)
|
||||
- Exits with code 1 on failure
|
||||
|
||||
13. **Copy and run the reproduction script:**
|
||||
|
||||
```bash
|
||||
scp /tmp/repro_script.py {SSH_TARGET}:/tmp/
|
||||
ssh {SSH_TARGET} "docker cp /tmp/repro_script.py {CONTAINER}:/tmp/"
|
||||
ssh {SSH_TARGET} "docker exec -e CUDA_VISIBLE_DEVICES=0,1 {CONTAINER} python3 /tmp/repro_script.py"
|
||||
```
|
||||
|
||||
14. **Run control experiments** to isolate the variable:
|
||||
- If suspecting TP issue: run with TP=1 as control
|
||||
- If suspecting GPU issue: compare same code on different GPU
|
||||
- If suspecting a specific commit: test before/after that commit
|
||||
|
||||
### Phase 6: Report
|
||||
|
||||
15. **Produce a structured report:**
|
||||
|
||||
```markdown
|
||||
## CI Regression Bisection Report
|
||||
|
||||
### Failure Signature
|
||||
- **Test**: {test_file}::{test_method}
|
||||
- **Error**: {exact error message}
|
||||
- **Key metrics**: {numeric values}
|
||||
- **Deterministic**: Yes/No
|
||||
|
||||
### Root Cause Classification
|
||||
One of:
|
||||
- **Code Regression**: PR #{number} introduced the bug
|
||||
- **Hardware-Specific**: Fails on {GPU_TYPE}, passes on others
|
||||
- **Environment Change**: New runner/driver/package version
|
||||
- **Pre-existing Flakiness**: Intermittent, not a new regression
|
||||
|
||||
### Evidence
|
||||
| Condition | Result |
|
||||
|-----------|--------|
|
||||
| {condition1} | PASS/FAIL |
|
||||
| {condition2} | PASS/FAIL |
|
||||
|
||||
### Timeline
|
||||
- {date}: Last known pass ({sha}, {runner})
|
||||
- {date}: First known fail ({sha}, {runner})
|
||||
- {date}: Confirmed reproduction on {server}
|
||||
|
||||
### Recommended Fix
|
||||
- **Short-term**: {workaround}
|
||||
- **Long-term**: {proper fix}
|
||||
```
|
||||
|
||||
## Key Patterns to Recognize
|
||||
|
||||
| Pattern | Diagnosis |
|
||||
|---------|-----------|
|
||||
| Same SHA passes on runner A, fails on runner B | Hardware/runner-specific |
|
||||
| All runners fail after commit X | Code regression from commit X |
|
||||
| Intermittent - same runner sometimes passes/fails | Flaky test or race condition |
|
||||
| Prefill OK but decode fails | TP/all-reduce issue in decode path |
|
||||
| Works with TP=1, fails with TP>1 | Tensor parallelism bug |
|
||||
| Exact same numeric diff every time | Deterministic bug, not flakiness |
|
||||
|
||||
## Important Notes
|
||||
|
||||
- **Always check runner identity** before concluding it's a code regression. Many "consistent" failures are actually runner-specific.
|
||||
- **Test partition assignments change over time** as tests are added/removed. A test may move between partitions, landing on different runner types.
|
||||
- **H200 runners** use `/root/actions-runner/` path and machine names like `gpu-h200-worker-*`. Non-H200 runners use `/public_sglang_ci/runner-*` paths.
|
||||
- When running remote reproduction, use `run_in_background` for long-running tests and check output with `TaskOutput`.
|
||||
- Container environments may be stale - always verify package versions match CI before drawing conclusions.
|
||||
444
third_party/sglang/.claude/skills/write-sglang-test/SKILL.md
vendored
Normal file
444
third_party/sglang/.claude/skills/write-sglang-test/SKILL.md
vendored
Normal file
@@ -0,0 +1,444 @@
|
||||
---
|
||||
name: write-sglang-test
|
||||
description: Guide for writing SGLang CI/UT tests. Covers CustomTestCase, CI registration, server fixtures, model selection, mock testing, and test placement. Always read test/README.md for the full CI layout, how to run tests, and extra tips. Use when creating new tests, adding CI test cases, writing unit tests, or when the user asks to add tests for SGLang features.
|
||||
---
|
||||
|
||||
# Writing SGLang CI / UT Tests
|
||||
|
||||
This skill covers **how to write and register tests**. For CI pipeline internals (stage ordering, fast-fail, gating, partitioning, debugging CI failures), see the [CI workflow guide](../ci-workflow-guide/SKILL.md).
|
||||
|
||||
## Core Rules
|
||||
|
||||
1. **Always use `CustomTestCase`** — never raw `unittest.TestCase`. It ensures `tearDownClass` runs even when `setUpClass` fails, preventing resource leaks in CI.
|
||||
2. **`tearDownClass` must be defensive** — use `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
|
||||
3. **Place tests in `test/registered/<category>/`** — except JIT kernel tests and benchmarks, which live in `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/` (nested subfolders are allowed)
|
||||
4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
|
||||
5. **Prefer mock over real server** — when testing logic that doesn't need a server / engine launch (middleware, request routing, config validation, argument parsing), use `unittest.mock.patch` / `MagicMock` and place tests in `test/registered/unit/`. Only launch a real server when the test genuinely needs inference results or server lifecycle behavior.
|
||||
|
||||
JIT kernel exception:
|
||||
- If the task is adding or updating code under `python/sglang/jit_kernel/`, prefer the `add-jit-kernel` skill first.
|
||||
- JIT kernel correctness tests use `python/sglang/jit_kernel/tests/**/test_*.py`.
|
||||
- JIT kernel benchmarks use `python/sglang/jit_kernel/benchmark/**/bench_*.py`.
|
||||
- Those files are still executed by `test/run_suite.py`, but through dedicated kernel suites rather than `test/registered/`.
|
||||
|
||||
---
|
||||
|
||||
## Model & Backend Selection
|
||||
|
||||
| Scenario | Model | CI Registration | Suite |
|
||||
|----------|-------|-----------------|-------|
|
||||
| **Unit tests** (no server / engine launch) | None | `register_cpu_ci` (prefer) or `register_cuda_ci` | `stage-a-test-cpu` or `stage-b-test-1-gpu-small` |
|
||||
| **Common / backend-independent** (middleware, abort, routing, config, arg parsing) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` only | `stage-b-test-1-gpu-small` |
|
||||
| **Model-agnostic functionality** (sampling, session, OpenAI API features) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` (+ AMD if relevant) | `stage-b-test-1-gpu-small` |
|
||||
| **General performance** (single node, no spec/DP/parallelism) | `DEFAULT_MODEL_NAME_FOR_TEST` (8B) | `register_cuda_ci` | `stage-b-test-1-gpu-large` |
|
||||
| **Bigger features** (spec, DP, TP, disaggregation) | Case by case | Case by case | See suite table below |
|
||||
|
||||
**Key principle for E2E tests**: Do NOT add `register_amd_ci` unless the test specifically exercises AMD/ROCm code paths. Common E2E tests just need any GPU to run — duplicating across backends wastes CI time with no extra coverage.
|
||||
|
||||
### All model constants
|
||||
|
||||
Defined in `python/sglang/test/test_utils.py`:
|
||||
|
||||
| Constant | Model | When to use |
|
||||
|----------|-------|-------------|
|
||||
| `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` | Llama-3.2-1B-Instruct | Common features, model-agnostic tests |
|
||||
| `DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE` | Llama-3.2-1B | Base (non-instruct) model tests |
|
||||
| `DEFAULT_MODEL_NAME_FOR_TEST` | Llama-3.1-8B-Instruct | General performance (single node) |
|
||||
| `DEFAULT_MOE_MODEL_NAME_FOR_TEST` | Mixtral-8x7B-Instruct | MoE-specific tests |
|
||||
| `DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST` | — | Embedding tests |
|
||||
| `DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST` | — | Vision-language tests |
|
||||
|
||||
### Naming Conventions
|
||||
|
||||
- **Suite**: `stage-{a,b,c}-test-{gpu_count}-gpu-{hardware}` (e.g., `stage-b-test-1-gpu-small`)
|
||||
- **CI runner**: `{gpu_count}-gpu-{hardware}` (e.g., `1-gpu-5090`, `4-gpu-h100`, `8-gpu-h200`)
|
||||
|
||||
### All CI Suites
|
||||
|
||||
#### Per-commit (CUDA)
|
||||
|
||||
| Suite | Runner (label) | Description |
|
||||
|-------|----------------|-------------|
|
||||
| `stage-a-test-1-gpu-small` | `1-gpu-5090` | Quick checks on a small NVIDIA GPU before heavier stages |
|
||||
| `stage-a-test-cpu` | `ubuntu-latest` | CPU-only unit tests |
|
||||
| `stage-b-test-1-gpu-small` | `1-gpu-5090` | Core engine tests that fit a 5090-class card |
|
||||
| `stage-b-test-1-gpu-large` | `1-gpu-h100` | Tests that need H100-class memory or kernels (e.g. FA3) |
|
||||
| `stage-b-test-2-gpu-large` | `2-gpu-h100` | Two-GPU correctness and parallelism (TP/PP) on H100 |
|
||||
| `stage-b-test-4-gpu-b200` | `4-gpu-b200` | Early Blackwell coverage (SM100+ paths) on four GPUs |
|
||||
| `stage-b-kernel-unit-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `python/sglang/jit_kernel/tests/` |
|
||||
| `stage-b-kernel-unit-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `python/sglang/jit_kernel/tests/` |
|
||||
| `stage-b-kernel-benchmark-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `python/sglang/jit_kernel/benchmark/` |
|
||||
| `stage-c-test-4-gpu-h100` | `4-gpu-h100` | Large 4-GPU H100 integration and scaling tests |
|
||||
| `stage-c-test-8-gpu-h200` | `8-gpu-h200` | Large 8-GPU H200 runs for big models and parallelism |
|
||||
| `stage-c-test-8-gpu-h20` | `8-gpu-h20` | Large 8-GPU H20 runs for big models |
|
||||
| `stage-c-test-deepep-4-gpu-h100` | `4-gpu-h100` | DeepEP expert-parallel and networking on four H100s |
|
||||
| `stage-c-test-deepep-8-gpu-h200` | `8-gpu-h200` | DeepEP at 8-GPU H200 scale |
|
||||
| `stage-c-test-8-gpu-b200` | `8-gpu-b200` | 8-GPU B200 suite (registered but not yet wired to a workflow) |
|
||||
| `stage-c-test-4-gpu-b200` | `4-gpu-b200` | 4-GPU B200 suite for large models on Blackwell |
|
||||
| `stage-c-test-4-gpu-gb200` | `4-gpu-gb200` | 4-GPU GB200 suite for large models on Grace Blackwell |
|
||||
|
||||
#### Per-commit (AMD)
|
||||
|
||||
| Suite | Runner (label) | Description |
|
||||
|-------|----------------|-------------|
|
||||
| `stage-a-test-1-gpu-small-amd` | `linux-mi325-1gpu-sglang` | Quick checks on one MI325-class GPU |
|
||||
| `stage-b-test-1-gpu-small-amd` | `linux-mi325-1gpu-sglang` | Core 1-GPU AMD tests (14 partitions) |
|
||||
| `stage-b-test-1-gpu-small-amd-nondeterministic` | `linux-mi325-1gpu-sglang` | Non-deterministic 1-GPU AMD tests |
|
||||
| `stage-b-test-1-gpu-small-amd-mi35x` | `linux-mi35x-gpu-1` | 1-GPU tests on MI35x hardware |
|
||||
| `stage-b-test-1-gpu-large-amd` | `linux-mi325-1gpu-sglang` | Large 1-GPU AMD tests (2 partitions) |
|
||||
| `stage-b-test-2-gpu-large-amd` | `linux-mi325-2gpu-sglang` | 2-GPU ROCm correctness and parallel setups |
|
||||
| `stage-b-test-large-8-gpu-35x-disaggregation-amd` | `linux-mi35x-gpu-8.fabric` | PD disaggregation and RDMA on 8×MI35x fabric |
|
||||
| `stage-c-test-4-gpu-amd` | `linux-mi325-4gpu-sglang` | 4-GPU AMD integration (2 partitions) |
|
||||
| `stage-c-test-large-8-gpu-amd` | `linux-mi325-8gpu-sglang` | 8-GPU MI325 scaling and integration |
|
||||
| `stage-c-test-large-8-gpu-amd-mi35x` | `linux-mi35x-gpu-8` | 8-GPU MI35x scaling (2 partitions) |
|
||||
|
||||
|
||||
### Per-commit (Ascend NPU)
|
||||
|
||||
| Suite | Runner (label) | Description |
|
||||
| --- | --- | --- |
|
||||
| `per-commit-1-npu-a2` | `linux-aarch64-a2-1` | 1-NPU LLM CI machine |
|
||||
| `per-commit-2-npu-a2` | `linux-aarch64-a2-2` | 2-NPU LLM CI machine |
|
||||
| `per-commit-4-npu-a3` | `linux-aarch64-a3-4` | 4-NPU LLM CI machine |
|
||||
| `per-commit-16-npu-a3` | `linux-aarch64-a3-16` | 16-NPU LLM CI machine |
|
||||
| `multimodal-gen-test-1-npu-a3` | `linux-aarch64-a3-2` | 1-NPU multimodal CI machine |
|
||||
| `multimodal-gen-test-2-npu-a3` | `linux-aarch64-a3-16` | 2-NPU multimodal CI machine |
|
||||
| `multimodal-gen-test-8-npu-a3` | `linux-aarch64-a3-16` | 8-NPU multimodal CI machine |
|
||||
|
||||
#### Nightly
|
||||
|
||||
Nightly suites are listed in `NIGHTLY_SUITES` in [`test/run_suite.py`](../../../test/run_suite.py). They run via `nightly-test-nvidia.yml`, `nightly-test-amd.yml` amd `nightly-test-npu.yml`, not `pr-test.yml`. Examples:
|
||||
|
||||
- `nightly-1-gpu` (CUDA)
|
||||
- `nightly-kernel-1-gpu` (CUDA, JIT kernel full grids)
|
||||
- `nightly-kernel-8-gpu-h200` (CUDA, multi-GPU JIT kernel nightly)
|
||||
- `nightly-8-gpu-h200` (CUDA)
|
||||
- `nightly-eval-vlm-2-gpu` (CUDA)
|
||||
- `nightly-amd` (AMD)
|
||||
- `nightly-amd-8-gpu-mi35x` (AMD)
|
||||
- `nightly-1-npu-a3` (NPU)
|
||||
- `nightly-2-npu-a3` (NPU)
|
||||
- `nightly-4-npu-a3` (NPU)
|
||||
- `nightly-8-npu-a3` (NPU)
|
||||
- `nightly-16-npu-a3` (NPU)
|
||||
|
||||
> **Note**: Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`.
|
||||
|
||||
### Choosing a Suite
|
||||
|
||||
Use the lightest suite that meets your test's needs:
|
||||
|
||||
- **No GPU required** → `stage-a-test-cpu`
|
||||
- **Most small GPU tests** → `stage-b-test-1-gpu-small` (default choice)
|
||||
- **Need H100 memory or Hopper features** → `stage-b-test-1-gpu-large`
|
||||
- **JIT kernel correctness** → `stage-b-kernel-unit-1-gpu-large`
|
||||
- **JIT kernel benchmarks** → `stage-b-kernel-benchmark-1-gpu-large`
|
||||
- **Multi-GPU** → only when the test actually needs multiple GPUs
|
||||
|
||||
---
|
||||
|
||||
## Test File Templates
|
||||
|
||||
### Unit Tests (no server / engine launch)
|
||||
|
||||
See `test/registered/unit/README.md` for quick-start and rules. Unit tests live in `test/registered/unit/`, mirroring `python/sglang/srt/`:
|
||||
|
||||
```python
|
||||
"""Unit tests for srt/<module>"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.<module> import TargetClass
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
|
||||
# Prefer CPU. Only use register_cuda_ci when the test truly needs a GPU.
|
||||
|
||||
class TestTargetClass(CustomTestCase):
|
||||
def test_basic_behavior(self):
|
||||
obj = TargetClass(...)
|
||||
self.assertEqual(obj.method(), expected)
|
||||
|
||||
@patch("sglang.srt.<module>.some_dependency")
|
||||
def test_with_mock(self, mock_dep):
|
||||
mock_dep.return_value = MagicMock()
|
||||
# test logic with dependency mocked
|
||||
...
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
```
|
||||
|
||||
Use `unittest.mock.patch` / `MagicMock` to mock dependencies and isolate the logic under test. If the module transitively imports GPU-only packages (e.g. `sgl_kernel`), they can be stubbed so the test runs on CPU CI. See `test/registered/unit/README.md` for details and examples.
|
||||
|
||||
**Quality bar** — test real logic (validation boundaries, state transitions, error paths, branching, etc.). Skip tests that just verify Python itself works (e.g., "does calling an abstract method raise `NotImplementedError`?", "does a dataclass store the field I assigned?"). Consolidate repetitive patterns into parameterized tests. No production code changes in test PRs.
|
||||
|
||||
### E2E test (small model, server needed)
|
||||
|
||||
```python
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=60, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
|
||||
class TestMyFeature(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--arg1", "value1"], # feature-specific args
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_basic_functionality(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={"text": "Hello", "sampling_params": {"max_new_tokens": 32}},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=3)
|
||||
```
|
||||
|
||||
### E2E test (8B model, server needed, performance)
|
||||
|
||||
```python
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=300, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
|
||||
class TestMyFeaturePerf(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_latency(self):
|
||||
start = time.perf_counter()
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={"text": "Hello", "sampling_params": {"max_new_tokens": 128}},
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertLess(elapsed, 5.0, "Latency exceeded threshold")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=3)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Server Fixture Reuse
|
||||
|
||||
For tests that only need a standard server, inherit from `DefaultServerBase` and override class attributes:
|
||||
|
||||
```python
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
|
||||
class TestMyFeature(DefaultServerBase):
|
||||
model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
other_args = ["--enable-my-feature"]
|
||||
|
||||
def test_something(self):
|
||||
...
|
||||
```
|
||||
|
||||
Available fixtures in `python/sglang/test/server_fixtures/`:
|
||||
|
||||
| Fixture | Use case |
|
||||
|---------|----------|
|
||||
| `DefaultServerBase` | Standard single-server tests |
|
||||
| `EagleServerBase` | EAGLE speculative decoding |
|
||||
| `PDDisaggregationServerBase` | Disaggregated prefill/decode |
|
||||
| `MMMUServerBase` | Multimodal VLM tests |
|
||||
|
||||
---
|
||||
|
||||
## CI Registration
|
||||
|
||||
Every CI-discovered test file must call a registration function at module level:
|
||||
|
||||
```python
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_cuda_ci,
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
register_npu_ci,
|
||||
)
|
||||
|
||||
# Per-commit test (small 1-gpu, runs on 5090)
|
||||
register_cuda_ci(est_time=80, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
# Per-commit test (large 1-gpu, runs on H100)
|
||||
register_cuda_ci(est_time=120, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
# Nightly-only test
|
||||
register_cuda_ci(est_time=200, suite="nightly-1-gpu", nightly=True)
|
||||
|
||||
# Multi-backend test (only when testing backend-specific code paths)
|
||||
register_cuda_ci(est_time=80, suite="stage-a-test-1-gpu-small")
|
||||
register_amd_ci(est_time=120, suite="stage-a-test-1-gpu-small-amd")
|
||||
register_npu_ci(est_time=400, suite="nightly-8-npu-a3", nightly=True)
|
||||
|
||||
# Temporarily disabled test
|
||||
register_cuda_ci(est_time=80, suite="stage-b-test-1-gpu-small", disabled="flaky - see #12345")
|
||||
```
|
||||
|
||||
Parameters:
|
||||
- `est_time`: estimated runtime in seconds (used for CI partitioning)
|
||||
- `suite`: which CI suite to run in (see suite tables above)
|
||||
- `nightly=True`: for nightly-only tests (default `False` = per-commit)
|
||||
- `disabled="reason"`: temporarily disable with explanation
|
||||
|
||||
**Key principle**: Only add `register_amd_ci` / `register_npu_ci` when the test exercises backend-specific code paths. Common E2E tests just need `register_cuda_ci` — duplicating across backends wastes CI time.
|
||||
|
||||
### JIT Kernel Registration
|
||||
|
||||
JIT kernel files live outside `test/registered/` but still use registration:
|
||||
|
||||
```python
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
# Correctness tests in python/sglang/jit_kernel/tests/
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="stage-b-kernel-unit-8-gpu-h200")
|
||||
|
||||
# Benchmarks in python/sglang/jit_kernel/benchmark/
|
||||
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
# Optional nightly registration
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-8-gpu-h200", nightly=True)
|
||||
```
|
||||
|
||||
Keep `est_time` and `suite` as **literal values** — `run_suite.py` collects them by AST parsing
|
||||
|
||||
---
|
||||
|
||||
## Test Placement
|
||||
|
||||
```
|
||||
test/
|
||||
├── registered/ # CI tests (auto-discovered by run_suite.py)
|
||||
│ ├── unit/ # No server / engine launch (see test/registered/unit/README.md)
|
||||
│ ├── kernels/ # CUDA kernel correctness (no server, GPU required)
|
||||
│ ├── sampling/ # test_penalty.py, test_sampling_params.py ...
|
||||
│ ├── sessions/ # test_session_control.py ...
|
||||
│ ├── openai_server/ # basic/, features/, validation/ ...
|
||||
│ ├── spec/ # eagle/, utils/ ...
|
||||
│ ├── models/ # model-specific accuracy tests
|
||||
│ ├── perf/ # performance benchmarks
|
||||
│ └── <category>/ # create new category if needed
|
||||
├── manual/ # Non-CI: debugging, one-off, manual verification
|
||||
└── run_suite.py # CI runner (scans registered/ plus jit_kernel test/benchmark files)
|
||||
|
||||
python/sglang/jit_kernel/
|
||||
├── tests/ # JIT kernel correctness tests (CI-discovered by test/run_suite.py)
|
||||
└── benchmark/ # JIT kernel benchmarks (CI-discovered by test/run_suite.py)
|
||||
```
|
||||
|
||||
**Decision rule** (see also `test/registered/README.md`):
|
||||
- Component logic, no server → `registered/unit/`
|
||||
- JIT kernel correctness / benchmarks → `python/sglang/jit_kernel/tests/` or `python/sglang/jit_kernel/benchmark/`
|
||||
- Other kernel correctness → `registered/kernels/`
|
||||
- Server needed → `registered/<category>/`
|
||||
- Local debugging → `manual/`
|
||||
|
||||
---
|
||||
|
||||
## Eval Accuracy Mixins
|
||||
|
||||
**Design philosophy**: Most test files don't care about eval logic — they only need a "does this feature break model output quality?" sanity check. The mixin pattern separates **what to test** (threshold) from **how to test** (run_eval, assertions, CI summary). Test classes declare thresholds as class attributes; the mixin provides the `test_*` method. Override when you need extra assertions (e.g. EAGLE accept length).
|
||||
|
||||
Available mixins in `python/sglang/test/kits/eval_accuracy_kit.py`: `MMLUMixin`, `HumanEvalMixin`, `MGSMEnMixin`, `GSM8KMixin`. Can be combined freely. Read the source for attrs and defaults.
|
||||
|
||||
```python
|
||||
class TestMyFeature(CustomTestCase, MMLUMixin):
|
||||
mmlu_score_threshold = 0.65
|
||||
mmlu_num_examples = 64
|
||||
mmlu_num_threads = 32
|
||||
# test_mmlu is inherited — no code needed
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Key Utilities
|
||||
|
||||
```python
|
||||
from sglang.test.test_utils import (
|
||||
CustomTestCase, # base class with retry logic
|
||||
popen_launch_server, # launch server subprocess
|
||||
DEFAULT_URL_FOR_TEST, # auto-configured base URL
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, # 600s default
|
||||
run_bench_serving, # benchmark helper (launch + bench)
|
||||
)
|
||||
from sglang.srt.utils import kill_process_tree # cleanup server
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Checklist
|
||||
|
||||
Before submitting a test:
|
||||
|
||||
- [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`)
|
||||
- [ ] Has `register_*_ci(...)` call at module level
|
||||
- [ ] Placed in `test/registered/<category>/`, unless this is a JIT kernel test/benchmark
|
||||
- [ ] JIT kernel work: files live in `python/sglang/jit_kernel/tests/` or `python/sglang/jit_kernel/benchmark/`
|
||||
- [ ] Backend-independent tests: `register_cuda_ci` only + smallest model
|
||||
- [ ] Logic that doesn't need a server / engine launch → unit test in `registered/unit/` (see Unit Tests section)
|
||||
- [ ] `setUpClass` launches server, `tearDownClass` kills it (if server-based)
|
||||
- [ ] `tearDownClass` is defensive — uses `hasattr`/null checks before accessing resources that may not have been allocated
|
||||
- [ ] Has `if __name__ == "__main__": unittest.main()`
|
||||
- [ ] `est_time` is reasonable (measure locally)
|
||||
Reference in New Issue
Block a user