From cae022aa8817ac6e5d4c3648766d5971449e8b1f Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Mon, 8 Jun 2026 22:14:48 -0700
Subject: [PATCH] [JIT] Reuse JIT kernel build cache across CI runs (#27605)

---
 python/sglang/jit_kernel/utils.py | 55 ++++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py
index 4b04ce1c5..575e69de9 100644
--- a/python/sglang/jit_kernel/utils.py
+++ b/python/sglang/jit_kernel/utils.py
@@ -83,7 +83,13 @@ def _local_jit_source_hash(source_files: List[str]) -> str:
         seen.add(path)
 
         data = path.read_bytes()
-        digest.update(str(path).encode())
+        # Relative to kernel root, not absolute: the key must track source
+        # content, not install location (differs across runners / job dirs).
+        try:
+            ident = str(path.relative_to(KERNEL_PATH))
+        except ValueError:
+            ident = path.name
+        digest.update(ident.encode())
         digest.update(b"\0")
         digest.update(data)
         digest.update(b"\0")
@@ -166,6 +172,31 @@ def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
     return CPPArgList(_convert(arg) for arg in args)
 
 
+@cache_once
+def _tvm_ffi_version() -> str:
+    try:
+        import tvm_ffi
+
+        version = getattr(tvm_ffi, "__version__", None)
+        if version:
+            return str(version)
+    except Exception:
+        pass
+    try:
+        from importlib.metadata import version as dist_version
+
+        return dist_version("apache-tvm-ffi")
+    except Exception:
+        return "unknown"
+
+
+def _jit_build_dir_name(module_name: str) -> str:
+    # Key on arch + tvm-ffi ABI too (module_name only hashes sources), so a
+    # shared cache volume never reuses a cross-arch/ABI .so.
+    arch = get_jit_cuda_arch().target_name
+    return f"{module_name}__arch_{arch}__tvmffi_{_tvm_ffi_version()}"
+
+
 def load_jit(
     *args: str,
     cpp_files: List[str] | None = None,
@@ -235,6 +266,28 @@ def load_jit(
     module_name = "sgl_kernel_jit_" + "_".join(str(arg) for arg in args)
     if cpp_files or cuda_files:
         module_name += "_" + _local_jit_source_hash(cpp_files + cuda_files)
+
+    # A built .so under a deterministic dir is content-addressed: load it
+    # directly to skip ninja, whose mtime check rebuilds every CI run (pip
+    # install bumps dep header mtimes).
+    if build_directory is None:
+        cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")
+        build_directory = str(
+            pathlib.Path(cache_dir).expanduser() / _jit_build_dir_name(module_name)
+        )
+    prebuilt = pathlib.Path(build_directory) / f"{module_name}.so"
+    if prebuilt.is_file():
+        from tvm_ffi import load_module
+
+        try:
+            module = load_module(str(prebuilt))
+            logger.debug("Reused cached JIT module %s", module_name)
+            return module
+        except Exception:
+            logger.warning(
+                "Cached JIT module %s failed to load; rebuilding.", module_name
+            )
+
     if header_only:
         cpp_wrappers = cpp_wrappers or []
         cuda_wrappers = cuda_wrappers or []