[Skills] Refine jit_kernel and sgl-kernel skills (#20095)

.claude/skills/add-jit-kernel/SKILL.md

@@ -17,8 +17,9 @@ Add a new operation that scales each element of a tensor by a scalar factor:
 
 ## When to use JIT vs AOT (`sgl-kernel`)
 
-- **JIT (`jit_kernel`)**: lightweight, few dependencies, rapid iteration, compiled on first use
+- **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`)**: depends on CUTLASS / FlashInfer / DeepGEMM, needs pre-built wheel
+- **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`.
 
 ---
 
@@ -26,6 +27,8 @@ Add a new operation that scales each element of a tensor by a scalar factor:
 
 **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
@@ -80,7 +83,7 @@ auto device = SymbolicDevice{};
 device.set_options<kDLCUDA>();
 TensorMatcher({N})  //
     .with_dtype<fp16_t>()
-    .with_device(device)
+    .with_device<kDLCUDA>(device)
     .verify(dst)
     .verify(src);  // same shape, dtype, device as dst
 const size_t n = N.unwrap();
@@ -95,7 +98,8 @@ const DLDevice dev = device.unwrap();
 
 - **`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/max/min(x)` — type-dispatched math (for `fp32_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)`.
 
@@ -105,7 +109,7 @@ const DLDevice dev = device.unwrap();
 #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 128-bit vector loads/stores for bandwidth efficiency.
+- **`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
@@ -117,19 +121,22 @@ const DLDevice dev = device.unwrap();
 #include <sgl_kernel/tile.cuh>
 ```
 
-- **`device::tile::Memory<T>::cta(blockDim.x)`** — Creates a tile accessor where each thread handles `tid = threadIdx.x` with stride `blockDim.x`. Common for loops over a 1D array.
+- `tile::Memory<T>` is fundamentally a **1D cooperative accessor** over a contiguous region.
-- **`.load(ptr, offset)`** — loads `ptr[tid + offset * blockDim.x]`
+- **`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.
-- **`.store(ptr, val, offset)`** — stores to `ptr[tid + offset * blockDim.x]`
+- **`.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/abs/sqrt/rsqrt<T>(a, b)` — type-dispatched math via `dtype_trait`
+- `device::math::max/min<T>(a, b)` — type-dispatched binary math via `dtype_trait`
-- `device::math::exp/sin/cos(float)` — fast float math wrappers
+- `device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x)` — type-dispatched unary math via `dtype_trait`
 
 ### `warp.cuh` — Warp-level primitives
 
@@ -191,11 +198,11 @@ Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`.
 The implementation fully uses the project abstractions described above:
 
 ```cpp
-#include <sgl_kernel/tensor.h>   // TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
-#include <sgl_kernel/type.cuh>   // dtype_trait, fp16_t, bf16_t, fp32_t
+#include <sgl_kernel/type.cuh>   // For dtype_trait, fp16_t, bf16_t, fp32_t
-#include <sgl_kernel/utils.h>    // RuntimeCheck, div_ceil
+#include <sgl_kernel/utils.h>    // For RuntimeCheck, div_ceil
-#include <sgl_kernel/utils.cuh>  // LaunchKernel, SGL_DEVICE
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
-#include <sgl_kernel/vec.cuh>    // AlignedVector
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
 
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/container/tensor.h>
@@ -257,7 +264,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
 
   TensorMatcher({N})  //
       .with_dtype<T>()
-      .with_device(device_)
+      .with_device<kDLCUDA>(device_)
       .verify(dst)
       .verify(src);  // same shape / dtype / device as dst
 
@@ -292,6 +299,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
 **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

.claude/skills/add-sgl-kernel/SKILL.md

@@ -18,8 +18,9 @@ Add a new operation that scales each element of a tensor by a scalar factor:
 
 ## Two rules of thumb (must follow)
 
-1. **Heavyweight kernels go to `sgl-kernel`.** If it depends on CUTLASS / FlashInfer / DeepGEMM (or similarly heavy stacks), implement it in `sgl-kernel/`.
+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. **Lightweight kernels go to `python/sglang/jit_kernel`.** If it is small, has few dependencies, and benefits from rapid iteration, implement it as a JIT kernel instead.
+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:
 
@@ -156,40 +157,51 @@ csrc/elementwise/scale.cu
 
 ## Step 5: Expose a Python API under `sgl-kernel/python/sgl_kernel/`
 
-In `sgl-kernel/python/sgl_kernel/__init__.py`, add:
+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
-from torch.ops import sgl_kernel as _ops
+import torch
 
-def scale(out: torch.Tensor, input: torch.Tensor, factor: float) -> None:
+def scale(
+    input: torch.Tensor,
+    factor: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
     """
-    Element-wise scale: out = input * factor (in-place into out).
+    Element-wise scale: out = input * factor.
 
     Supported dtypes: torch.float16, torch.bfloat16, torch.float32.
 
     Parameters
     ----------
-    out    : pre-allocated CUDA output tensor (same shape/dtype as input)
     input  : CUDA input tensor
     factor : scale factor (float)
+    out    : optional pre-allocated CUDA output tensor (same shape/dtype as input)
     """
-    _ops.scale(out, input, factor)
+    if out is None:
+        out = torch.empty_like(input)
+    torch.ops.sgl_kernel.scale.default(out, input, factor)
+    return out
 ```
 
-Or export it from the existing module organisation — follow the pattern already used by similar ops in `__init__.py`.
+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])
@@ -197,7 +209,8 @@ def test_scale_correctness(dtype, size, factor):
     input = torch.randn(size, dtype=dtype, device="cuda")
     out   = torch.empty_like(input)
 
-    sgl_kernel.scale(out, input, factor)
+    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)
@@ -208,26 +221,20 @@ 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(out, input, 2.0)
+        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(out, input, 2.0)
+        sgl_kernel.scale(input, 2.0, out=out)
 
 
 if __name__ == "__main__":
     pytest.main([__file__, "-q"])
 ```
 
-Run:
-
-```bash
-pytest sgl-kernel/tests/test_scale.py -q
-```
-
 ---
 
 ## Step 7: Add a benchmark (required)
@@ -279,7 +286,7 @@ def benchmark(dtype, size, provider):
     factor = 2.0
 
     if provider == "sglang":
-        fn = lambda: sgl_kernel.scale(out, input, factor)
+        fn = lambda: sgl_kernel.scale(input, factor, out=out)
     else:
         fn = lambda: torch_scale(input, factor)
 
@@ -293,15 +300,9 @@ if __name__ == "__main__":
     benchmark.run(print_data=True)
 ```
 
-Run:
-
-```bash
-python sgl-kernel/benchmark/bench_scale.py
-```
-
 ---
 
-## Step 8: Build and validate
+## Step 8: Build
 
 Build:
 
@@ -317,7 +318,11 @@ cd sgl-kernel
 make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
 ```
 
-Validate:
+---
+
+## Step 9: Validate
+
+After building successfully, run the test and benchmark:
 
 ```bash
 pytest sgl-kernel/tests/test_scale.py -q
@@ -352,7 +357,8 @@ 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/__init__.py      # MODIFIED: export Python API
+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
 ```