[JIT kernel] Update jit_kernel cache and develop doc (#17842)

This commit is contained in:
Xiaoyu Zhang
2026-01-28 15:09:47 +08:00
committed by GitHub
parent 2573a262af
commit c08b54a575
9 changed files with 46 additions and 45 deletions

View File

@@ -1,17 +1,16 @@
from __future__ import annotations
import functools
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import load_jit, make_cpp_args
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@functools.cache
@cache_once
def _jit_add_constant_module(constant: int) -> Module:
args = make_cpp_args(constant) # pass all the template argument
return load_jit(