add-jit-kernel

Installation
SKILL.md

Tutorial: Adding a New JIT Kernel to SGLang

Apply kernel-organization for the public operator namespace, logical grouping, lazy registry metadata, and test placement. The implementation tutorial below does not replace that API contract.

This tutorial walks through adding a simple element-wise scale operation as a JIT kernel. We'll implement scale(x, factor) = x * factor to demonstrate the complete workflow.

Goal

Add a new operation that scales each element of a tensor by a scalar factor:

  • 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)

When to use JIT vs AOT (sgl-kernel)

Installs
13
GitHub Stars
36.7K
First Seen
May 16, 2026