[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

@@ -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<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, bind device to symbol
+  - `.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)
 
@@ -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 <typename T, int kVecN, int32_t kFactorNumer, int32_t kFactorDenom>
+template <typename T, int kVecN>
 __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<float>(kFactorNumer)
-                          / static_cast<float>(kFactorDenom);
-
   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;
@@ -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<T>(static_cast<float>(v[i]) * kFactor);
+      v[i] = static_cast<T>(static_cast<float>(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<T>(static_cast<float>(src[base + i]) * kFactor);
+    dst[base + i] = static_cast<T>(static_cast<float>(src[base + i]) * factor);
   }
 }
 
 // ----------------------------------------------------------------
 // Launcher: validates tensors, selects vector width, launches kernel
 // ----------------------------------------------------------------
-template <typename T, int32_t kFactorNumer, int32_t kFactorDenom>
-void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
+template <typename T>
+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<uint32_t>(N.unwrap());
-  const DLDevice device    = device_.unwrap();
+  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_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<uint32_t>(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<T, kVecN, kFactorNumer, kFactorDenom>,
+      scale_kernel<T, kVecN>,
       static_cast<T*>(dst.data_ptr()),
       static_cast<const T*>(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<To, From>` or `dtype_trait<T>::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)
 
 

.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`
 
 ---