From 20a23e3173b5b3bdb16f19451f55d448b3747fb0 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Wed, 18 Mar 2026 23:00:33 +0800 Subject: [PATCH] [SKILL] Refine kernel authoring docs and validate add-jit-kernel / add-sgl-kernel end to end with Codex (#20867) --- .claude/skills/add-jit-kernel/SKILL.md | 82 ++++++++++--------- .claude/skills/add-sgl-kernel/SKILL.md | 3 +- .../development_jit_kernel_guide.md | 6 ++ 3 files changed, 50 insertions(+), 41 deletions(-) diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index 232c9eb99..e43f37064 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -11,7 +11,7 @@ This tutorial walks through adding a simple element-wise scale operation as a JI Add a new operation that scales each element of a tensor by a scalar factor: -- Input: tensor `x` (CUDA) and scalar `factor` (float, passed as C++ template argument) +- 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`)** @@ -72,7 +72,7 @@ This is the **primary validation API** for all kernel launchers. Use it to valid - **`host::TensorMatcher({dims...})`** — Fluent builder for tensor validation: - `.with_dtype()` — require a specific C++ type (e.g. `fp16_t`) - `.with_dtype()` — allow a set of types - - `.with_device(device_sym)` — require CUDA, bind device to symbol + - `.with_device(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) @@ -213,18 +213,15 @@ 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) -// kFactor = scale factor encoded as kFactorNumer / kFactorDenom +// factor = runtime scale factor // ---------------------------------------------------------------- -template +template __global__ void scale_kernel(T* __restrict__ dst, const T* __restrict__ src, - uint32_t n_vecs, - uint32_t n_remainder, + float factor, uint32_t n_total) { - constexpr float kFactor = static_cast(kFactorNumer) - / static_cast(kFactorDenom); - using vec_t = device::AlignedVector; + const uint32_t n_vecs = n_total / kVecN; // --- vectorised body --- const uint32_t vec_stride = blockDim.x * gridDim.x; @@ -235,7 +232,7 @@ __global__ void scale_kernel(T* __restrict__ dst, v.load(src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { - v[i] = static_cast(static_cast(v[i]) * kFactor); + v[i] = static_cast(static_cast(v[i]) * factor); } v.store(dst, vi); } @@ -244,17 +241,17 @@ __global__ void scale_kernel(T* __restrict__ dst, 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; - i < n_remainder; + base + i < n_total; i += scalar_stride) { - dst[base + i] = static_cast(static_cast(src[base + i]) * kFactor); + dst[base + i] = static_cast(static_cast(src[base + i]) * factor); } } // ---------------------------------------------------------------- // Launcher: validates tensors, selects vector width, launches kernel // ---------------------------------------------------------------- -template -void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { +template +void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { using namespace host; // 1. Validate input tensors with TensorMatcher @@ -268,28 +265,26 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { .verify(dst) .verify(src); // same shape / dtype / device as dst - const uint32_t n = static_cast(N.unwrap()); - const DLDevice device = device_.unwrap(); + const uint32_t n = static_cast(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_vecs = n / kVecN; - const uint32_t n_remainder = n % kVecN; + constexpr int kVecN = 16 / sizeof(T); + const uint32_t n_work_items = div_ceil(n, static_cast(kVecN)); // 3. Launch constexpr uint32_t kBlockSize = 256; - const uint32_t grid = div_ceil(std::max(n_vecs, n_remainder), kBlockSize); + const uint32_t grid = div_ceil(n_work_items, kBlockSize); LaunchKernel(grid, kBlockSize, device)( - scale_kernel, + scale_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), - n_vecs, - n_remainder, + factor, n); } @@ -304,6 +299,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { - 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` or `dtype_trait::from(val)` for cross-type conversions - `device::math::` functions for device math instead of bare `__` intrinsics @@ -328,9 +324,9 @@ if TYPE_CHECKING: @cache_once -def _jit_scale_module(dtype: torch.dtype, factor_numer: int, factor_denom: int) -> Module: - """Compile and cache the JIT scale module for a given dtype and factor.""" - args = make_cpp_args(dtype, factor_numer, factor_denom) +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, @@ -355,22 +351,26 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> ------- Scaled tensor (dst = src * factor). """ - assert src.is_cuda, "src must be a CUDA tensor" - assert src.dtype in (torch.float16, torch.bfloat16, torch.float32), ( - f"Unsupported dtype {src.dtype}. Supported: float16, bfloat16, float32" - ) + 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: - assert out.shape == src.shape, "out shape must match src" - assert out.dtype == src.dtype, "out dtype must match src" + 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") - # Encode factor as integer ratio; denom=1000 gives 3 decimal places of precision - factor_denom = 1000 - factor_numer = round(factor * factor_denom) - - module = _jit_scale_module(src.dtype, factor_numer, factor_denom) - module.scale(out, 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 ``` @@ -378,8 +378,10 @@ def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> - 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 | |--------------------|------------| @@ -443,13 +445,13 @@ def test_scale_out_param(dtype): def test_scale_cpu_error(): src = torch.randn(128, dtype=torch.float16) # CPU tensor - with pytest.raises(AssertionError, match="CUDA"): + 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(AssertionError, match="Unsupported dtype"): + with pytest.raises(RuntimeError, match="dtype"): scale(src, 2.0) diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index 07767c16b..4c5dc7a52 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -106,6 +106,7 @@ void scale(at::Tensor& out, const at::Tensor& input, double factor) { **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 @@ -136,7 +137,7 @@ m.impl("scale", torch::kCUDA, &scale); - `Tensor!` means in-place / mutable output argument - The schema is important for `torch.compile` and for consistent call signatures -- If your underlying C++ API uses `float` but PyTorch bindings expect `double`, the implicit cast is fine for scalars; use shims if needed for other types +- 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` --- diff --git a/docs/developer_guide/development_jit_kernel_guide.md b/docs/developer_guide/development_jit_kernel_guide.md index ce0e9b8aa..b09476e48 100644 --- a/docs/developer_guide/development_jit_kernel_guide.md +++ b/docs/developer_guide/development_jit_kernel_guide.md @@ -241,6 +241,10 @@ def _jit_add_constant_module(constant: int) -> Module: def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor: + if not src.is_cuda: + raise RuntimeError("src must be a CUDA tensor") + if src.dtype != torch.int32: + raise RuntimeError(f"Unsupported dtype {src.dtype}. Supported: int32") dst = torch.empty_like(src) module = _jit_add_constant_module(constant) module.add_constant(dst, src) @@ -248,6 +252,8 @@ def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor: ``` +Keep the Python wrapper thin, but still validate the basic invariants such as device and dtype before dispatch. In the current JIT/FFI path, invalid tensors are not always rejected safely before launch. + ### STEP 3: Use your kernel Finally, import and use the kernel like a regular Python function: