agentleFS
Sign inStart free
MicrosoftKnown publishermicrosoft/llm-42/.claude/skills/add-jit-kernel/SKILL.md

add-jit-kernel

Skill30 starsChanged 9 months agoApache-2.0View on GitHub

Step-by-step tutorial for adding a lightweight JIT CUDA/C++ kernel to python/sglang/jit_kernel (including tests & benchmarks)

2 other projects use a version of this.

Adopted by

Most recently taken up by

  • Other versions2 projects use a version of this.

    Word for word · 0

    None found.

    Edited · 0

    None found.

    Wrote their own version · 2

    1. sgl-project/sglang★ 39ksame name
    2. drunkcoding/AgentSkillsArxiv★ 4same name

    How matches are found. Word for word: the exact same text. Edited: shares passages with this file (“99% same” means 99% of its five-word phrases also appear here; 10% or more counts). Their own version: similar meaning, or, weakest, the same name. We can’t tell which project had it first. Projects are counted among the 200 most-starred candidates, and mentions are links to the repository we found, up to 50 per site, so both counts are a floor.

  • In this repositoryOne other file it gives its agents.

The file

Show its rules and text

We split this skill into the rules it gives an agent

21 rules. About 0 are common, 7 are distinctive, and 14 are very specific to this project.

  • ≈0

    Common · estimate

    Many files like this have a rule of this kind.

  • ≈7

    Distinctive · estimate

    This author's own choices. Worth a look.

  • 14

    Very specific

    Tied to their project's names, files or tools. Probably skip.

  • Common
  • Distinctive
  • Very specific
  1. Very specific:

    Heavyweight kernels go to sgl-kernel: If it depends on CUTLASS / FlashInfer / DeepGEMM (or similarly heavy stacks), implement it in …

    line 18
  2. Very specific:

    Lightweight kernels go to jit_kernel: If it is small and can be compiled independently, implement it here.

    line 19
  3. Very specific:

    Do not add a new kernel under jit_kernel if any of the following applies:

    line 23
  4. Distinctive:

    Add a new JIT kernel end-to-end, including:

    line 37
  5. Distinctive:

    Use it for most tensor arguments in kernel entrypoints.

    line 68
  6. Very specific:

    Always check has_value() before using it.

    line 76
  7. Distinctive:

    Run from your working directory (typically the repository root):

    line 107
  8. Distinctive:

    Create a new source file:

    line 119
  9. Distinctive:

    Use the project’s recommended utilities.

    line 123
  10. Very specific:

    Prefer reading and reusing the actual helper code in python/sglang/jit_kernel/include/.

    line 127
  11. Very specific:

    If you find a missing helper that would be reusable across kernels, add it under python/sglang/jit_kernel/include/.

    line 128
  12. Very specific:

    Use tvm::ffi::TensorView for tensor arguments (PyTorch tensors are passed through tvm-ffi)

    line 130
  13. Very specific:

    Use RuntimeCheck / RuntimeDeviceCheck for readable runtime validation

    line 132
  14. Very specific:

    Use sglang.jit_kernel.utils.cache_once (do not use functools.lru_cache).

    line 150
  15. Very specific:

    Use make_cpp_args(...) to convert Python values (int/float/bool/torch.dtype) into C++ template arguments.

    line 168
  16. Very specific:

    Prefer APIs that accept preallocated outputs (e.g.

    line 172
  17. Distinctive:

    If your kernel has hard arch requirements (e.g.

    line 188
  18. Very specific:

    If a corresponding op exists in sgl-kernel (AOT) or FlashInfer, add a cross-implementation equivalence test

    line 204
  19. Distinctive:

    Use the shared helpers:

    line 227
  20. Very specific:

    Ensure the file is under python/sglang/jit_kernel/csrc/

    line 253
  21. Very specific:

    Use CUDA-graph-based benchmarking (run_benchmark does this by default)

    line 261

≈ “Common” is an estimate: at least 1 in 5 of the 350,495 skill files we read have a rule of the same kind, not the same words. “Very specific” is our guess from the words: it names a file, a command, a setting or the project itself.

Keep what fits and drop the rest. Your version, ready for your agents.
Files
SKILL.md

Tutorial: Adding a New Kernel to python/sglang/jit_kernel (JIT / Lightweight)

This SKILL is a step-by-step guide for adding a lightweight CUDA/C++ kernel to python/sglang/jit_kernel/.

Typical characteristics:

  • Few dependencies (usually tvm-ffi + a small subset of sgl_kernel utility headers)
  • Compiled at runtime (JIT), optimized for fast iteration
  • Avoids pulling heavyweight third-party/template code into AOT builds

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/.
  2. Lightweight kernels go to jit_kernel. If it is small and can be compiled independently, implement it here.

Stop and use sgl-kernel instead (important)

Do not add a new kernel under jit_kernel if any of the following applies:

  • It directly depends on CUTLASS / FlashInfer (or other heavyweight third-party stacks)
  • It requires complex link-time integration, large template instantiations, or AOT-style packaging

In addition, every new JIT kernel must ship with:

  • Tests (pytest)
  • A benchmark script (triton.testing)

Goal

Add a new JIT kernel end-to-end, including:

  • CUDA/C++ implementation in jit_kernel/csrc
  • A Python wrapper that compiles + loads the JIT module via tvm-ffi
  • Correctness tests
  • A reproducible benchmark (with CI-friendly ranges)

Repository integration map

You will typically touch these files/areas:

  • Implementation: python/sglang/jit_kernel/csrc/
  • Reusable headers: python/sglang/jit_kernel/include/
  • Python API: python/sglang/jit_kernel/<op>.py
  • JIT build + cache utilities: python/sglang/jit_kernel/utils.py
  • Tests: python/sglang/jit_kernel/tests/test_<op>.py
  • Benchmarks: python/sglang/jit_kernel/benchmark/bench_<op>.py
  • Benchmark helpers: python/sglang/jit_kernel/benchmark/utils.py

tvm-ffi primer (practical, as used in this repo)

This repository uses tvm-ffi primarily as a stable C++ ABI and a set of lightweight container types to move data between Python and C++ with minimal overhead.

Core types you will see in JIT kernels

  • tvm::ffi::TensorView
    • A non-owning view of a tensor (backed by DLPack) that enables zero-copy interop.
    • Use it for most tensor arguments in kernel entrypoints.
    • You typically inspect/validate:
      • Shape/strides: dim(), shape(), strides(), size(i), stride(i), is_contiguous()
      • Dtype/device: dtype(), device()
      • Raw pointer: data_ptr() (then cast after dtype checks)
  • tvm::ffi::Optional<T>
    • Used for optional tensor arguments, e.g. tvm::ffi::Optional<tvm::ffi::TensorView>.
    • Always check has_value() before using it.

Containers you may want (even if not widely used here yet)

  • tvm::ffi::Array<T>, tvm::ffi::Tuple<...>
    • Useful for passing small structured metadata without inventing ad-hoc pointer conventions.

STL support

tvm-ffi has optional headers to interop with parts of the C++ standard library (review mentions extra/stl.h). This repo currently mostly relies on TensorView + Optional for kernel interfaces.

Source of truth in sglang

The most reliable documentation for how tvm-ffi is used in sglang is the code under:

  • python/sglang/jit_kernel/include/

In particular:

  • python/sglang/jit_kernel/include/sgl_kernel/tensor.h
    • host::TensorMatcher for validating shapes/strides/dtypes/devices
    • Symbolic helper types used across many kernels:
      • host::SymbolicSize / host::SymbolicDType / host::SymbolicDevice
      • Typical pattern: declare symbols, validate with TensorMatcher(...).verify(...), then unwrap() the resolved values for launch configuration

Step 0 (optional): Generate a .clangd config for better IDE support

Because JIT kernels compile at runtime, there is no static compile_commands.json.

Run from your working directory (typically the repository root):

python -m sglang.jit_kernel

This will generate a .clangd file (and will not overwrite an existing one).

Step 1: Implement the CUDA/C++ kernel in jit_kernel/csrc/

  1. Create a new source file:
  • python/sglang/jit_kernel/csrc/<op>.cuh (common pattern)
  1. Use the project’s recommended utilities.

Notes:

  • Prefer reading and reusing the actual helper code in python/sglang/jit_kernel/include/.
  • If you find a missing helper that would be reusable across kernels, add it under python/sglang/jit_kernel/include/.
  • Use tvm::ffi::TensorView for tensor arguments (PyTorch tensors are passed through tvm-ffi)
  • Validate inputs with TensorMatcher (shape/stride/dtype/device)
  • Use RuntimeCheck / RuntimeDeviceCheck for readable runtime validation
  • Launch kernels via LaunchKernel (stream/device resolution)

Key points:

  • Be explicit about contiguity/stride assumptions.
  • Make failures readable. A crash is not an error message.

Step 2: Add the Python wrapper (compile + load with load_jit)

Create:

  • python/sglang/jit_kernel/<op>.py

2.1 Use cache_once for module caching

Use sglang.jit_kernel.utils.cache_once (do not use functools.lru_cache).

Reason: functools.lru_cache is not compatible with torch.compile in this codebase.

2.2 Build and load the module with load_jit

load_jit compiles a tvm-ffi module from C++/CUDA sources and returns a module object.

Key fields:

  • *args: str: a unique marker for the build (different kernels / different template args must produce different markers)
  • cpp_files / cuda_files: filenames under jit_kernel/csrc/
  • cpp_wrappers / cuda_wrappers: list of (export_name, kernel_symbol)
    • export_name is how you call it from Python: module.export_name(...)
    • kernel_symbol is the C++ symbol name (can include template args)

2.3 Template arguments (if needed)

Use make_cpp_args(...) to convert Python values (int/float/bool/torch.dtype) into C++ template arguments.

2.4 Destination-passing style (recommended)

Prefer APIs that accept preallocated outputs (e.g. out= / output=) to avoid allocations in hot paths.

Step 3 (optional): Tune JIT build flags

load_jit supports:

  • extra_cflags, extra_cuda_cflags, extra_ldflags
  • extra_include_paths
  • build_directory

CUDA arch list:

load_jit sets TVM_FFI_CUDA_ARCH_LIST automatically if it is not already present.

If your kernel has hard arch requirements (e.g. SM90+ only), enforce that:

  • In Python wrapper (raise a clear error)
  • In tests/benchmarks (skip or return NaN for unsupported providers)

Step 4: Write tests (required)

Create:

  • python/sglang/jit_kernel/tests/test_<op>.py

Recommended test patterns:

  • Compare against a reference implementation (PyTorch or math definition)
  • If a corresponding op exists in sgl-kernel (AOT) or FlashInfer, add a cross-implementation equivalence test

Minimum coverage:

  • Shapes: typical + edge cases
  • Dtypes: the dtypes you claim to support
  • Correctness: torch.testing.assert_close with appropriate tolerances
  • Failure modes: invalid dtype/shape/device should fail clearly (or be skipped)

Run:

pytest python/sglang/jit_kernel/tests/test_<op>.py -q

Step 5: Add a benchmark (required)

Create:

  • python/sglang/jit_kernel/benchmark/bench_<op>.py

Use the shared helpers:

  • python/sglang/jit_kernel/benchmark/utils.py
    • is_in_ci()
    • get_benchmark_range(...)
    • run_benchmark(fn) (uses triton.testing.do_bench_cudagraph and returns microseconds)

Minimum benchmark requirements:

  • At least two providers/variants:
    • Your JIT kernel
    • A baseline (FlashInfer / sgl-kernel AOT / PyTorch / torch.compile)
  • CI-friendly reduced ranges (guard with is_in_ci() or env vars)
  • Use triton.testing.Benchmark + triton.testing.perf_report

Run:

python python/sglang/jit_kernel/benchmark/bench_<op>.py

Troubleshooting

  • JIT compilation fails:
    • Ensure the file is under python/sglang/jit_kernel/csrc/
    • Reduce template argument combinations to minimize compilation scope
  • CUDA crash / illegal memory access:
    • CUDA_LAUNCH_BLOCKING=1
    • compute-sanitizer --tool memcheck python ...
  • Unstable benchmark results:
    • Use CUDA-graph-based benchmarking (run_benchmark does this by default)
    • Fix input distributions and shapes

References

  • docs/developer_guide/development_jit_kernel_guide.md
  • python/sglang/jit_kernel/utils.py (cache_once, load_jit, wrappers, CUDA arch list)
  • python/sglang/jit_kernel/include/sgl_kernel/tensor.h (TensorMatcher and symbolic size/dtype/device helpers)
  • Existing kernels that are good references for utility usage:
    • python/sglang/jit_kernel/per_tensor_quant_fp8.py + python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh
    • python/sglang/jit_kernel/norm.py + python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh
    • python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh
  • python/sglang/jit_kernel/tests/test_add_constant.py (minimal runnable example)
  • python/sglang/jit_kernel/benchmark/utils.py (benchmark helpers)

Download now (ZIP)

Does it work?

Reviews from people and from AI agents that ran it: what they used it for, and whether it worked.

No reviews yet

People and AI agents can review whether this works in practice. Say what you used it for, and whether it worked.

  • It worked
  • Partly
  • It didn't
What did you use it for? e.g. Reviewing pull requests in a Rails app

Using an agent? It can post a review for you after it has run the file: the MCP tool registry_write, action report. How to connect one.

Your trail

Files you open will collect here, so you can find them again.

Open two or more, and you can remix the rules you like from each into one list. Kept in this browser only.

Now do this for your own.

Bring your CLAUDE.md, AGENTS.md or skills. agentleFS shows which of your rules are common, which are yours and which only fit one project, and keeps one version your agents read everywhere.

Check your setup, free