[core] Consolidate compiled-kernel caches under SGLANG_CACHE_DIR (#32434)

scripts/ci/cuda/ci_install_dependency.sh

@@ -155,8 +155,20 @@ install_apt_packages() {
 }
 
 clean_site_packages() {
-    # Clear torch compilation cache
-    python3 -c 'import os, shutil, tempfile, getpass; cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") or os.path.join(tempfile.gettempdir(), "torchinductor_" + getpass.getuser()); shutil.rmtree(cache_dir, ignore_errors=True)'
+    # Clear torch compilation cache from every location it can be in; sglang
+    # is not installed yet, so it cannot be asked which one is in use.
+    python3 -c '
+import getpass, os, shutil, tempfile
+
+sglang_cache_dir = os.environ.get("SGLANG_CACHE_DIR") or "~/.cache/sglang"
+for cache_dir in (
+    os.environ.get("TORCHINDUCTOR_CACHE_DIR"),
+    os.path.join(tempfile.gettempdir(), "torchinductor_" + getpass.getuser()),
+    os.path.join(os.path.expanduser(sglang_cache_dir), "inductor"),
+):
+    if cache_dir:
+        shutil.rmtree(cache_dir, ignore_errors=True)
+'
 
     # Remove broken dist-info directories (missing METADATA per PEP 376)
     SITE_PACKAGES=$(python3 -c "import site; print(site.getsitepackages()[0])")

scripts/ci/cuda/warmup_deep_gemm.py

@@ -26,9 +26,15 @@ from math import ceil
 from pathlib import Path
 from typing import Dict, List
 
-# Shared with warmup_server.py. Wipe alongside /root/.cache/deep_gemm if you
-# clear the DeepGEMM JIT cache — a stale marker → in-test JIT compile.
-MARKER_DIR = os.path.join(os.path.expanduser("~"), ".cache", "sglang", "warmup_markers")
+from sglang.srt.environ import envs
+
+# Shared with warmup_server.py. Uses the same root as DG_JIT_CACHE_DIR below,
+# so overriding SGLANG_CACHE_DIR moves the markers and the cache together.
+# If only one moved, a marker could report a model as warmed while its cache
+# is empty, and the test would pay for the JIT compilation it should skip.
+MARKER_DIR = os.path.join(
+    os.path.expanduser(envs.SGLANG_CACHE_DIR.get()), "warmup_markers"
+)
 
 # Outer cap for stuck fallback subprocesses; CRASH_MARKERS abort sooner.
 FALLBACK_TIMEOUT_SEC = 600
@@ -67,11 +73,10 @@ CRASH_MARKERS = (
     "Received sigquit from a child",
 )
 
-# Configure DeepGEMM cache before importing deep_gemm
-os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
-    "SGLANG_DG_CACHE_DIR",
-    os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm"),
-)
+# Configure DeepGEMM cache before importing deep_gemm. Read through envs so
+# this warms the directory the server will actually compile into; duplicating
+# the default here is what let the two drift apart.
+os.environ["DG_JIT_CACHE_DIR"] = envs.SGLANG_DG_CACHE_DIR.get()
 os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0")
 
 BLOCK_SIZE = 128
@@ -510,9 +515,7 @@ def main():
     )
     print(f"=== DeepGEMM Lightweight Warmup ({len(model_tp_pairs)} model(s)) ===")
     print(f"    Fast warmup: {fast_warmup}")
-    print(
-        f"    Cache dir: {os.environ.get('DG_JIT_CACHE_DIR', '~/.cache/deep_gemm')}\n"
-    )
+    print(f"    Cache dir: {os.environ['DG_JIT_CACHE_DIR']}\n")
 
     # Load configs and deduplicate by architecture
     seen_keys = {}

scripts/ci/cuda/warmup_server.py

@@ -26,9 +26,8 @@ from pathlib import Path
 
 # Reuse helpers from warmup_deep_gemm (same directory)
 sys.path.insert(0, os.path.dirname(__file__))
-from warmup_deep_gemm import get_architecture_key, get_config_json
+from warmup_deep_gemm import MARKER_DIR, get_architecture_key, get_config_json
 
-MARKER_DIR = os.path.join(os.path.expanduser("~"), ".cache", "sglang", "warmup_markers")
 HEALTH_POLL_INTERVAL = 10  # seconds between health checks
 SERVER_STARTUP_TIMEOUT = 900  # 15 min max to wait for server ready
 DEFAULT_PORT = 39876