Add SGLANG_CRASH_ON_JIT_COMPILE to forbid on-the-fly JIT compilation (#36615)

Signed-off-by: Shiyan Deng <dsy842974287@meta.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

python/sglang/kernels/jit/utils/compile/cache.py

@@ -298,10 +298,15 @@ def _normalize_text(text: str) -> str:
     return text
 
 
-def build_key_dir(*, module_name: str, build_key: str) -> pathlib.Path:
+def cache_root() -> pathlib.Path:
     configured = envs.SGLANG_JIT_CACHE_DIR.get() or "~/.cache/sglang/jit"
-    root = pathlib.Path(configured).expanduser()
-    return root / _target_tag() / module_name / f"{_BUILD_KEY_PREFIX}{build_key}"
+    return pathlib.Path(configured).expanduser()
+
+
+def build_key_dir(*, module_name: str, build_key: str) -> pathlib.Path:
+    return (
+        cache_root() / _target_tag() / module_name / f"{_BUILD_KEY_PREFIX}{build_key}"
+    )
 
 
 # ---------------------------------------------------------------------------

python/sglang/kernels/jit/utils/compile/loader.py

@@ -23,6 +23,7 @@ from sglang.kernels.jit.utils.compile.paths import (
 )
 from sglang.kernels.jit.utils.compile.spec import BuildSpec, resolve_sources
 from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES
+from sglang.srt.environ import envs
 
 if TYPE_CHECKING:
     from tvm_ffi import Module
@@ -125,6 +126,16 @@ def load_jit(
                 e,
             )
 
+    # Before the lock: with the flag set no process ever publishes a build,
+    # so waiting on the lock cannot turn this miss into a hit.
+    if envs.SGLANG_CRASH_ON_JIT_COMPILE.get():
+        raise RuntimeError(
+            f"JIT module '{spec.module_name}' has no usable cached build under "
+            f"{scope}, and SGLANG_CRASH_ON_JIT_COMPILE forbids compiling at "
+            "runtime. Seed the cache with prebuilt artifacts matching this "
+            "environment, or unset the flag."
+        )
+
     with _build_lock(scope):
         # Re-check: whoever held the lock before us has very likely just
         # published exactly what we were about to build. This is what turns N

python/sglang/srt/environ.py

@@ -1095,6 +1095,10 @@ class Envs:
     # is what makes reverting an edit an instant hit instead of a rebuild; set
     # it to trade that away for disk (1 keeps only the most recent build).
     SGLANG_JIT_CACHE_KEEP = EnvInt(None)
+    # Raise instead of compiling when a module misses the cache, so a
+    # deployment that expects a pre-seeded cache fails loudly at startup
+    # rather than silently eating a cold compile.
+    SGLANG_CRASH_ON_JIT_COMPILE = EnvBool(False)
 
     # ===================================================================
     # Expert-parallel dispatch and MoE execution