minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Revert "[Feature] JIT activation and update skills (by codex)" (#22078)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-04-03 15:04:15 -07:00
committed by GitHub
parent 8cb337c8ea
commit ac1e437f6a
16 changed files with 64 additions and 489 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/add-jit-kernel/SKILL.md
+21 -40
View File
@@ -186,9 +186,7 @@ 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 -h # for verbose help info about clangd configuration
python -m sglang.jit_kernel
python -m sglang.jit_kernel --dep cutlass flashinfer # with cutlass/flashinfer dependency
```
---
@@ -200,8 +198,6 @@ Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`.
The implementation fully uses the project abstractions described above:
```cpp
// NOTE: Comments for headers are not common in practice.
// It is only shown here for tutorial purposes to highlight the key abstractions.
#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
@@ -219,7 +215,7 @@ namespace {
// kVecN = number of elements per vector load (e.g. 8 for fp16)
// factor = runtime scale factor
// ----------------------------------------------------------------
template <typename T, int kVecN, bool kUsePDL>
template <typename T, int kVecN>
__global__ void scale_kernel(T* __restrict__ dst,
const T* __restrict__ src,
float factor,
@@ -227,10 +223,6 @@ __global__ void scale_kernel(T* __restrict__ dst,
using vec_t = device::AlignedVector<T, kVecN>;
const uint32_t n_vecs = n_total / kVecN;
// If using PDL, wait for primary kernel before any global memory load.
// This is NOT a synchronization point, which means some threads can early exit before this.
device::PDLWaitPrimary<kUsePDL>();
// --- vectorised body ---
const uint32_t vec_stride = blockDim.x * gridDim.x;
for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x;
@@ -253,16 +245,12 @@ __global__ void scale_kernel(T* __restrict__ dst,
i += scalar_stride) {
dst[base + i] = static_cast<T>(static_cast<float>(src[base + i]) * factor);
}
// If using PDL, signal for the secondary kernel to start after all threads have finished
// This is NOT a synchronization point, which means some threads can early exit before this.
device::PDLTriggerSecondary<kUsePDL>();
}
// ----------------------------------------------------------------
// Launcher: validates tensors, selects vector width, launches kernel
// ----------------------------------------------------------------
template <typename T, bool kUsePDL>
template <typename T>
void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
using namespace host;
@@ -283,12 +271,8 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
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 x 2 bytes = 16 bytes
// fp32: 4 elements x 4 bytes = 16 bytes
// We encourage using `device::kMaxVecBytes`, which will change according to
// the target architecture and can enable 256-bit vectorization on SM100+ if desired.
// But 128-bit is more commonly adapted for better compatibility,
// so it's still ok to hardcode 16 here just for simplicity.
// 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));
@@ -296,10 +280,8 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
constexpr uint32_t kBlockSize = 256;
const uint32_t grid = div_ceil(n_work_items, kBlockSize);
// PDL feature is 100% optional. Without `enable_pdl`, the code should still be correct.
// Try to enable it if profiling shows that it can benefit the performance of this kernel.
LaunchKernel(grid, kBlockSize, device).enable_pdl(kUsePDL)(
scale_kernel<T, kVecN, kUsePDL>,
LaunchKernel(grid, kBlockSize, device)(
scale_kernel<T, kVecN>,
static_cast<T*>(dst.data_ptr()),
static_cast<const T*>(src.data_ptr()),
factor,
@@ -320,8 +302,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
- 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 if possible.
- Try to use `PDL` feature. In some cases, this will benefit the performance.
- `device::math::` functions for device math instead of bare `__` intrinsics
---
@@ -336,12 +317,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@@ -350,7 +326,7 @@ if TYPE_CHECKING:
@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, is_arch_support_pdl())
args = make_cpp_args(dtype)
return load_jit(
"scale",
*args,
@@ -375,16 +351,24 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) ->
-------
Scaled tensor (dst = src * factor).
"""
# DO NOT add too much proactive validation here.
# Keep the Python wrapper thin, only enforce the preconditions
# that the current JIT/FFI path (C++ side) does not reject on its own.
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
@@ -397,7 +381,6 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) ->
- 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:
- `is_arch_support_pdl()` checks if the current architecture supports PDL, which is typically passed as a template argument to the kernel.
- 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 |
@@ -410,8 +393,6 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) ->
## Step 3 (optional): Tune JIT build flags
If your kernel uses some math functions like `expf` or `sinf`, consider enabling `--use_fast_math` for better performance (with a potential precision tradeoff):
```python
return load_jit(
"scale",
@@ -433,7 +414,7 @@ if torch.cuda.get_device_capability()[0] < 9:
## 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.
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`).