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

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