chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[SKILL] Refine kernel authoring docs and validate add-jit-kernel / add-sgl-kernel end to end with Codex (#20867)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-18 23:00:33 +08:00
committed by GitHub
parent 21c4fc6334
commit 20a23e3173
3 changed files with 50 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
docs/developer_guide/development_jit_kernel_guide.md
View File

@@ -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: