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

docs/developer_guide/development_jit_kernel_guide.md

@@ -30,6 +30,9 @@ The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and
 To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
 The function can then be called in Python as `module.func`.
 
+For caching compiled modules, prefer `sglang.jit_kernel.utils.cache_once` over `functools.lru_cache`.
+`functools.lru_cache` is not compatible with `torch.compile`.
+
 ### C++ Utilities
 
 The following C++ utilities are available:
@@ -216,19 +219,17 @@ Create a new file at [jit_kernel/add_constant.py](../../python/sglang/jit_kernel
 
 ```python
 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(
@@ -255,4 +256,4 @@ Finally, import and use the kernel like a regular Python function:
 from sglang.jit_kernel.add_constant import add_constant
 ```
 
-For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/test_add_constant.py).
+For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/tests/test_add_constant.py).

docs/index.rst

@@ -100,6 +100,7 @@ Its core features include:
 
    developer_guide/contribution_guide.md
    developer_guide/development_guide_using_docker.md
+   developer_guide/development_jit_kernel_guide.md
    developer_guide/benchmark_and_profiling.md
    developer_guide/bench_serving.md
    developer_guide/evaluating_new_models.md

python/sglang/jit_kernel/add_constant.py

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

python/sglang/jit_kernel/cuda_wait_value.py

@@ -1,23 +1,22 @@
 from __future__ import annotations
 
-from functools import lru_cache
 from typing import TYPE_CHECKING
 
 import torch
 
-from sglang.jit_kernel.utils import load_jit
+from sglang.jit_kernel.utils import cache_once, load_jit
 
 if TYPE_CHECKING:
     import torch
     from tvm_ffi.module import Module
 
 
-@lru_cache(maxsize=1)
+@cache_once
 def _jit_stream_wait_value_module() -> Module:
     return load_jit(
         "cuda_wait_value",
         cuda_files=["cuda_wait_value.cuh"],
-        cuda_wrappers=[("stream_wait_value", "stream_wait_value")],
+        cuda_wrappers=[("stream_wait_value", "cuda_wait_value")],
     )
 
 

python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py

@@ -3,11 +3,13 @@
 import os
 import pathlib
 from typing import Tuple
-from functools import partial, lru_cache
+from functools import partial
 from dataclasses import dataclass, fields
 
 import torch
 
+from sglang.jit_kernel.utils import cache_once
+
 try:
     from triton.tools.disasm import extract
 except ImportError:
@@ -33,12 +35,12 @@ torch2cute_dtype_map = {
 }
 
 
-@lru_cache
+@cache_once
 def get_max_active_clusters(cluster_size):
     return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size)
 
 
-@lru_cache
+@cache_once
 def get_device_capacity(device: torch.device = None) -> Tuple[int, int]:
     return torch.cuda.get_device_capability(device)
 

python/sglang/jit_kernel/flash_attention/cute/interface.py

@@ -20,12 +20,14 @@
 # - bwd pass optimized for Hopper/Blackwell
 
 import math
-from functools import lru_cache
 from typing import Optional, Tuple, Callable
 
 import torch
 
 
+from sglang.jit_kernel.utils import cache_once
+
+
 import cuda.bindings.driver as cuda
 
 import cutlass
@@ -51,7 +53,7 @@ from .block_sparsity import (
     get_block_sparse_broadcast_pattern,
 )
 
-@lru_cache(maxsize=None)
+@cache_once
 def _get_device_capability():
     """Cached device capability check."""
     return torch.cuda.get_device_capability()[0]

python/sglang/jit_kernel/hicache.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import logging
-from functools import lru_cache
 from typing import TYPE_CHECKING
 
-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:
     import torch
@@ -13,7 +12,7 @@ if TYPE_CHECKING:
 DEFAULT_BLOCK_QUOTA = 2
 
 
-@lru_cache(maxsize=None)
+@cache_once
 def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module:
     num_threads, occupancy = 1024, 1
     args = make_cpp_args(

python/sglang/jit_kernel/timestep_embedding.py

@@ -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_timestep_embedding_module(dtype: torch.dtype) -> Module:
     args = make_cpp_args(dtype)
     return load_jit(

python/sglang/jit_kernel/utils.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import functools
 import pathlib
-from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union
 
 import torch
@@ -11,12 +10,32 @@ if TYPE_CHECKING:
     from tvm_ffi import Module
 
 
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def cache_once(fn: F) -> F:
+    """
+    NOTE: `functools.lru_cache` is not compatible with `torch.compile`
+    So we manually implement a simple cache_once decorator to replace it.
+    """
+    result_map = {}
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        key = (args, tuple(sorted(kwargs.items(), key=lambda x: x[0])))
+        if key not in result_map:
+            result_map[key] = fn(*args, **kwargs)
+        return result_map[key]
+
+    return wrapper  # type: ignore
+
+
 def _make_wrapper(tup: Tuple[str, str]) -> str:
     export_name, kernel_name = tup
     return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
 
 
-@lru_cache()
+@cache_once
 def _resolve_kernel_path() -> pathlib.Path:
     cur_dir = pathlib.Path(__file__).parent.resolve()
 
@@ -145,26 +164,6 @@ def load_jit(
     )
 
 
-F = TypeVar("F", bound=Callable[..., Any])
-
-
-def cache_once(fn: F) -> F:
-    """
-    NOTE: `functools.lru_cache` is not compatible with `torch.compile`
-    So we manually implement a simple cache_once decorator to replace it.
-    """
-    result_map = {}
-
-    @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
-        key = (args, tuple(sorted(kwargs.items(), key=lambda x: x[0])))
-        if key not in result_map:
-            result_map[key] = fn(*args, **kwargs)
-        return result_map[key]
-
-    return wrapper  # type: ignore
-
-
 @cache_once
 def is_arch_support_pdl() -> bool:
     import torch