Reuse DeepGEMM JIT cache across restarts

New sgl-deep-gemm wheels persist cubins, but SGLang was still replaying the full warmup envelope after every restart. The wrapper also translated SGLANG_DG_* env after importing deep_gemm, so the wheel could observe stale cache settings.\n\nMove DeepGEMM JIT env preparation before the first import, bridge the current SGLANG_DG_USE_NVRTC spelling, replace dense M walks with category-specific sparse M grids, and add a SGLang-side warmup manifest under the DeepGEMM cache dir. A manifest hit skips only the SGLang replay loop; missing or stale cache entries still force warmup and refresh the manifest.\n\nConstraint: sgl-deep-gemm 0.1.2 removed the compile-mode API, so dense 1..m_max loops launch real kernels.\nConstraint: DeepGEMM cache keys include compiler flags such as the deep_gemm include path; path-stable environments are still required for cross-container cache hits.\nRejected: SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 | bypasses warmup instead of making cache reuse correct.\nRejected: Enable per-process cache dirs by default | isolates corruption but defeats persistent cross-process cache reuse.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce dense DeepGEMM warmup for new wheels unless compile-mode support is verified on the target wheel.\nTested: Local py_compile for jit_cache.py configurer.py compile_utils.py model_runner.py\nTested: Local pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote g0034 cjy-glm5-new py_compile for same files\nTested: Remote g0034 cjy-glm5-new pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote env probe verified SGLANG_DG_CACHE_DIR and SGLANG_DG_USE_NVRTC bridge to DG_JIT_* before configurer completes\nNot-tested: Full decode cold-cache/hot-cache restart timing and manifest-hit ETE log validation

python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py

@@ -2,6 +2,7 @@ import logging
 import os
 from contextlib import contextmanager
 from enum import IntEnum, auto
+from importlib.metadata import PackageNotFoundError, version
 from typing import Dict, List, Tuple
 
 import torch
@@ -12,6 +13,15 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     restore_symmetric_memory_context,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers.deep_gemm_wrapper.jit_cache import (
+    build_warmup_m_lists,
+    build_sparse_m_list,
+    cache_entries,
+    make_warmup_manifest_key,
+    mark_warmup_complete,
+    prepare_deep_gemm_jit_env,
+    should_skip_warmup,
+)
 from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import ceil_div, get_available_gpu_memory
@@ -22,134 +32,51 @@ if ENABLE_JIT_DEEPGEMM:
     import deep_gemm
 
 
-_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1))
+_M_LIST_PER_RANK: List[int] = build_sparse_m_list(1024 * 16)
-# Separate, possibly-shrunk M grid for NON-grouped (dense/attention) shapes under
+_M_LIST_GATHERED: List[int] = build_sparse_m_list(1024 * 16)
-# CP prefill -- see _cp_dense_warmup_divisor / update_deep_gemm_config.
+_M_LIST_DECODE: List[int] = build_sparse_m_list(1024)
-_BUILTIN_M_LIST_DENSE = _BUILTIN_M_LIST
 _ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
 _DO_COMPILE_ALL = True
 _IS_FIRST_RANK_ON_NODE = envs.SGLANG_IS_FIRST_RANK_ON_NODE.get()
 _IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get()
-_FAST_WARMUP = envs.SGLANG_JIT_DEEPGEMM_FAST_WARMUP.get()
+# New DeepGEMM wheels dropped compile-mode, so dense 1..m_max enumeration would
+# launch real kernels. Sparse warmup is now always used.
 
-# Force redirect deep_gemm cache_dir
+_DEEP_GEMM_CACHE_DIR = prepare_deep_gemm_jit_env(
-os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
+    preload_kernels=_ENABLE_JIT_DEEPGEMM_PRECOMPILE
-    "SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
 )
 
-# Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f
-# NVRTC may have performance loss with some cases.
-# And NVCC JIT speed is also 9x faster in the ref commit
-os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0")
-
-# Enable DeepGEMM kernel preloading if precompile is enabled
-if _ENABLE_JIT_DEEPGEMM_PRECOMPILE:
-    os.environ["DG_PRELOAD_KERNELS"] = "1"
-
-
-def _cp_dense_warmup_divisor(server_args: ServerArgs) -> int:
-    """CP divisor for the NON-grouped (dense/attention) DeepGEMM warmup M grid.
-
-    Under NSA prefill context-parallel ``in-seq-split`` the sequence is split
-    across ``attn_cp_size`` ranks BEFORE the transformer layers
-    (deepseek_v2.py cp_split_and_rebuild_data), so every dense ``num_groups==1``
-    GEMM -- attention q/k/v/o projections, the dense/shared-expert MLP, the NSA
-    indexer ``weights_proj`` -- runs on only ~tokens/attn_cp_size per rank.
-    Warming those for the full non-CP M range wastes ~cp_size x.
-
-    The MoE GROUPED GEMM is deliberately NOT shrunk: deepep all-to-all re-gathers
-    every token, so its M = sum(num_recv_tokens_per_expert) ~= chunked * topk /
-    ep_size (== chunked for GLM-5.1 where topk==ep_size==8), independent of CP --
-    it keeps the full grid.
-
-    Shrink only when EVERY dense-extend path is CP-split: the main prefill extend
-    always is; the EAGLE draft extend is CP-split only when
-    ``SGLANG_CP_DRAFT_SHARED_KV`` is set (``_is_cp_shared_kv_draft_extend``,
-    ``include_v2=True``).  If a draft is configured without that env, the draft
-    extend runs non-CP at full tokens, so the dense shapes still need the full
-    grid and we return 1.
-    """
-
-    cp_size = int(getattr(server_args, "attn_cp_size", 1) or 1)
-    if cp_size <= 1:
-        return 1
-    prefill_cp_on = bool(
-        getattr(server_args, "enable_nsa_prefill_context_parallel", False)
-    ) and getattr(server_args, "nsa_prefill_cp_mode", None) == "in-seq-split"
-    if not prefill_cp_on:
-        return 1
-    has_draft = getattr(server_args, "speculative_algorithm", None) is not None
-    if has_draft and not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
-        return 1
-    return cp_size
-
 
 def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
-    global _BUILTIN_M_LIST
+    global _M_LIST_PER_RANK, _M_LIST_GATHERED, _M_LIST_DECODE
-    global _BUILTIN_M_LIST_DENSE
     global _DO_COMPILE_ALL
     global _IS_FIRST_RANK_ON_NODE
 
-    _BUILTIN_M_LIST = []
+    m_lists = build_warmup_m_lists(
-
+        chunked_prefill_size=server_args.chunked_prefill_size,
-    if _FAST_WARMUP:
+        attn_cp_size=getattr(server_args, "attn_cp_size", 1),
-        # In fast warmup mode, only compile a small set of typical Ms
+        max_running_requests=server_args.max_running_requests,
-
+        speculative_num_draft_tokens=server_args.speculative_num_draft_tokens,
-        # First cover all the small bs to ensure decode performance
+    )
-        _BUILTIN_M_LIST += list(range(1, 1025))
+    _M_LIST_PER_RANK = m_lists.per_rank
-
+    _M_LIST_GATHERED = m_lists.gathered
-        # Then cover larger batch sizes with gradually increasing steps
+    _M_LIST_DECODE = m_lists.decode
-        # For example, when chunekd prefill size is 16384
-        # The sampled Ms would be:
-        #   1024, 1026, ... 2046 (step 2)
-        #   2048, 2052, ... 4092 (step 4)
-        #   4096, 5004, ... 8184 (step 8)
-        #   8192, 9008, ... 16384 (step 16)
-        # Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels
-        next_m, sample_step = 1024, 2
-        max_prefill_bs = (
-            min(server_args.chunked_prefill_size, 32 * 1024)
-            if server_args.chunked_prefill_size >= 1
-            else 16 * 1024
-        )
-        while next_m < max_prefill_bs:
-            _BUILTIN_M_LIST += list(range(next_m, 2 * next_m, sample_step))
-            next_m = next_m * 2
-            sample_step = sample_step * 2
-        _BUILTIN_M_LIST.append(max_prefill_bs)
-        _BUILTIN_M_LIST = sorted(list(set(_BUILTIN_M_LIST)))
-    else:
-        # When fast warmup isn't enabled, generate m_max and compile all the covered Ms.
-        m_max = 1024 * 16
-        if server_args.chunked_prefill_size < 1:
-            m_max = 1024 * 64
-        elif server_args.chunked_prefill_size > 8192:
-            m_max = server_args.chunked_prefill_size * 2
-        m_max = min(1024 * 128, m_max)
-        _BUILTIN_M_LIST += list(range(1, m_max + 1))
-
-    # Dense (non-grouped) shapes run at ~tokens/cp_size under CP prefill in-seq-split;
-    # shrink their M grid by the (gated) CP divisor while the MoE grouped grid stays
-    # full.  Keep an absolute floor so degenerate cp_size/chunk combos still cover a
-    # reasonable dense M range.
-    cp_dense_div = _cp_dense_warmup_divisor(server_args)
-    if cp_dense_div > 1 and _BUILTIN_M_LIST:
-        dense_m_max = max(ceil_div(max(_BUILTIN_M_LIST), cp_dense_div), 2048)
-        _BUILTIN_M_LIST_DENSE = [m for m in _BUILTIN_M_LIST if m <= dense_m_max]
-        logger.info(
-            "DeepGEMM warmup: CP in-seq-split active (attn_cp_size divisor=%s); "
-            "dense (non-grouped) shapes warmed up to M=%s (%s Ms) vs full grouped "
-            "grid M=%s (%s Ms).",
-            cp_dense_div,
-            dense_m_max,
-            len(_BUILTIN_M_LIST_DENSE),
-            max(_BUILTIN_M_LIST),
-            len(_BUILTIN_M_LIST),
-        )
-    else:
-        _BUILTIN_M_LIST_DENSE = _BUILTIN_M_LIST
 
     _IS_FIRST_RANK_ON_NODE = server_args.base_gpu_id == gpu_id
+    if _IS_FIRST_RANK_ON_NODE:
+        logger.info(
+            "[DeepGEMM precompile] chunked=%s attn_cp=%s -> "
+            "per_rank max_m=%s (%s entries), gathered max_m=%s (%s entries), "
+            "decode max_m=%s (%s entries)",
+            server_args.chunked_prefill_size,
+            getattr(server_args, "attn_cp_size", 1),
+            max(_M_LIST_PER_RANK) if _M_LIST_PER_RANK else 0,
+            len(_M_LIST_PER_RANK),
+            max(_M_LIST_GATHERED) if _M_LIST_GATHERED else 0,
+            len(_M_LIST_GATHERED),
+            max(_M_LIST_DECODE) if _M_LIST_DECODE else 0,
+            len(_M_LIST_DECODE),
+        )
 
     # Check if is the first rank on node.
     # Default each rank will try compile all Ms to
@@ -167,13 +94,37 @@ class DeepGemmKernelType(IntEnum):
 
 _INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict()
 
-# Grouped (MoE) GEMMs are fed by the deepep all-to-all and run at M ~= chunked
+def _m_list_for(kernel_type: DeepGemmKernelType) -> List[int]:
-# regardless of CP, so they always use the full M grid; non-grouped (dense) GEMMs
+    if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
-# are CP-per-rank and use the (possibly shrunk) dense grid.
+        return _M_LIST_GATHERED
-_GROUPED_GEMM_KERNEL_TYPES = (
+    if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED:
-    DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED,
+        return _M_LIST_DECODE
-    DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG,
+    return _M_LIST_PER_RANK
-)
+
+
+def _deep_gemm_runtime_fingerprint() -> Dict[str, str]:
+    try:
+        deep_gemm_version = version("sgl-deep-gemm")
+    except PackageNotFoundError:
+        try:
+            deep_gemm_version = version("deep-gemm")
+        except PackageNotFoundError:
+            deep_gemm_version = "unknown"
+
+    try:
+        capability = torch.cuda.get_device_capability(0)
+        sm = f"{capability[0]}{capability[1]}"
+    except Exception:
+        sm = "unknown"
+
+    return {
+        "deep_gemm_version": deep_gemm_version,
+        "deep_gemm_file": str(getattr(deep_gemm, "__file__", "")),
+        "torch_version": str(torch.__version__),
+        "cuda_version": str(getattr(torch.version, "cuda", "")),
+        "sm": sm,
+        "dg_jit_use_nvrtc": os.environ.get("DG_JIT_USE_NVRTC", "0"),
+    }
 
 
 # TODO improve code
@@ -184,8 +135,6 @@ def _maybe_compile_deep_gemm_one_type_all(
     num_groups: int,
 ) -> None:
     global _INITIALIZATION_DICT
-    global _BUILTIN_M_LIST
-    global _BUILTIN_M_LIST_DENSE
 
     query_key = (kernel_type, n, k, num_groups)
     if (
@@ -207,13 +156,28 @@ def _maybe_compile_deep_gemm_one_type_all(
                 "`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`"
             )
 
-        # Grouped (MoE) shapes need the full M grid (deepep re-gathers all tokens);
+        m_list = _m_list_for(kernel_type)
-        # dense shapes are CP-per-rank and use the shrunk dense grid.
+        manifest_key = make_warmup_manifest_key(
-        m_list = (
+            kernel_type=kernel_type.name,
-            _BUILTIN_M_LIST
+            n=n,
-            if kernel_type in _GROUPED_GEMM_KERNEL_TYPES
+            k=k,
-            else _BUILTIN_M_LIST_DENSE
+            num_groups=num_groups,
+            m_list=m_list,
+            runtime_fingerprint=_deep_gemm_runtime_fingerprint(),
         )
+        if should_skip_warmup(_DEEP_GEMM_CACHE_DIR, manifest_key):
+            logger.info(
+                "[DeepGEMM-cache] warmup manifest hit; skip replay for "
+                "<%s> N=%s K=%s num_groups=%s m_count=%s max_m=%s cache_dir=%s",
+                kernel_type.name,
+                n,
+                k,
+                num_groups,
+                len(m_list),
+                max(m_list) if m_list else 0,
+                _DEEP_GEMM_CACHE_DIR,
+            )
+            return
 
         logger.info(
             f"Try DeepGEMM JIT Compiling for "
@@ -229,6 +193,11 @@ def _maybe_compile_deep_gemm_one_type_all(
             num_groups=num_groups,
             m_list=m_list,
         )
+        mark_warmup_complete(
+            _DEEP_GEMM_CACHE_DIR,
+            manifest_key,
+            cache_entries(_DEEP_GEMM_CACHE_DIR),
+        )
 
 
 # NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced

python/sglang/srt/layers/deep_gemm_wrapper/configurer.py

@@ -1,6 +1,7 @@
 import logging
 
 from sglang.srt.environ import envs
+from sglang.srt.layers.deep_gemm_wrapper.jit_cache import prepare_deep_gemm_jit_env
 from sglang.srt.utils import get_device_sm, is_blackwell_supported
 
 logger = logging.getLogger(__name__)
@@ -11,6 +12,9 @@ def _compute_enable_deep_gemm():
     if sm_version < 90:
         return False
 
+    prepare_deep_gemm_jit_env(
+        preload_kernels=envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
+    )
     try:
         import deep_gemm  # noqa: F401
     except ImportError:

python/sglang/srt/layers/deep_gemm_wrapper/jit_cache.py

@@ -0,0 +1,216 @@
+"""DeepGEMM JIT cache helpers.
+
+This module is deliberately stdlib-only.  It is imported before ``deep_gemm`` in
+the wrapper startup path, so it must not import torch, SGLang utility modules, or
+anything that can transitively import ``deep_gemm``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Mapping
+
+MANIFEST_FILENAME = "sglang_deep_gemm_warmup_manifest.v1.json"
+MANIFEST_VERSION = 1
+
+
+def _env_true(value: str | None) -> bool:
+    return value is not None and value.lower() in ("1", "true", "yes")
+
+
+def default_deep_gemm_cache_dir() -> str:
+    return os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
+
+
+def prepare_deep_gemm_jit_env(*, preload_kernels: bool = True) -> str:
+    """Set DeepGEMM JIT env vars before the first ``import deep_gemm``.
+
+    DeepGEMM's JIT cache key depends on its own environment variables and
+    compiler flags.  SGLang exposes ``SGLANG_*`` knobs, but the wheel consumes
+    ``DG_*`` knobs.  This bridge is intentionally idempotent:
+
+    * explicit ``DG_JIT_CACHE_DIR`` / ``DG_JIT_USE_NVRTC`` wins;
+    * otherwise SGLang env aliases are translated once;
+    * cache directory is created eagerly so startup fails at the real cause.
+    """
+
+    cache_dir = os.environ.get("DG_JIT_CACHE_DIR")
+    if not cache_dir:
+        cache_dir = os.environ.get("SGLANG_DG_CACHE_DIR", default_deep_gemm_cache_dir())
+        if _env_true(os.environ.get("SGLANG_DG_CACHE_DIR_PER_PROCESS")):
+            cache_dir = os.path.join(cache_dir, f"pid_{os.getpid()}")
+        os.environ["DG_JIT_CACHE_DIR"] = cache_dir
+
+    os.makedirs(cache_dir, exist_ok=True)
+
+    if "DG_JIT_USE_NVRTC" not in os.environ:
+        os.environ["DG_JIT_USE_NVRTC"] = os.environ.get(
+            "SGLANG_DG_USE_NVRTC",
+            os.environ.get("SGL_DG_USE_NVRTC", "0"),
+        )
+
+    if preload_kernels and "DG_PRELOAD_KERNELS" not in os.environ:
+        os.environ["DG_PRELOAD_KERNELS"] = "1"
+
+    return cache_dir
+
+
+def build_sparse_m_list(m_max: int) -> list[int]:
+    """Sparse M grid matching DeepGEMM v0.1+ block-size breakpoints.
+
+    The old dense 1..m_max walk was only cheap while DeepGEMM exposed
+    compile-mode.  New wheels launch real kernels for each M, while cubin
+    selection only changes at coarse breakpoints.  Keep 1..1024 dense for
+    decode/small prefill, then double both range and stride.
+    """
+
+    m_max = max(1, int(m_max))
+    out = list(range(1, min(1024, m_max) + 1))
+    next_m, step = 1024, 2
+    while next_m < m_max:
+        out.extend(range(next_m, min(2 * next_m, m_max), step))
+        next_m *= 2
+        step *= 2
+    out.append(m_max)
+    return sorted(set(out))
+
+
+@dataclass(frozen=True)
+class WarmupMLists:
+    per_rank: list[int]
+    gathered: list[int]
+    decode: list[int]
+
+
+def build_warmup_m_lists(
+    *,
+    chunked_prefill_size: int,
+    attn_cp_size: int,
+    max_running_requests: int,
+    speculative_num_draft_tokens: int | None,
+) -> WarmupMLists:
+    """Build category-specific DeepGEMM warmup M lists.
+
+    * per_rank: dense/attention/indexer GEMMs after NSA CP split.
+    * gathered: MoE contiguous GEMMs after DeepEP all-to-all.
+    * decode: masked GEMMs, bounded by running speculative decode batch.
+    """
+
+    chunked = int(chunked_prefill_size)
+    if chunked < 1:
+        chunked = 64 * 1024
+    chunked = min(chunked, 128 * 1024)
+
+    cp_size = max(1, int(attn_cp_size or 1))
+    m_max_per_rank = max(1024, chunked // cp_size)
+    m_max_gathered = chunked
+
+    spec_tokens = max(1, int(speculative_num_draft_tokens or 1))
+    m_max_decode = min(1024, max(256, int(max_running_requests) * spec_tokens))
+
+    return WarmupMLists(
+        per_rank=build_sparse_m_list(m_max_per_rank),
+        gathered=build_sparse_m_list(m_max_gathered),
+        decode=build_sparse_m_list(m_max_decode),
+    )
+
+
+def cache_entries(cache_root: str | os.PathLike[str]) -> set[str]:
+    """Return DeepGEMM cubin cache entry directory names under ``cache/``."""
+
+    cache_dir = Path(cache_root) / "cache"
+    if not cache_dir.is_dir():
+        return set()
+    return {p.name for p in cache_dir.iterdir() if p.is_dir()}
+
+
+def make_warmup_manifest_key(
+    *,
+    kernel_type: str,
+    n: int,
+    k: int,
+    num_groups: int,
+    m_list: Iterable[int],
+    runtime_fingerprint: Mapping[str, object],
+) -> str:
+    m_values = [int(m) for m in m_list]
+    payload = {
+        "version": MANIFEST_VERSION,
+        "kernel_type": kernel_type,
+        "n": int(n),
+        "k": int(k),
+        "num_groups": int(num_groups),
+        "m_list_count": len(m_values),
+        "m_list_max": max(m_values) if m_values else 0,
+        "m_list_hash": hashlib.sha256(
+            ",".join(str(m) for m in m_values).encode("utf-8")
+        ).hexdigest(),
+        "runtime": dict(sorted(runtime_fingerprint.items())),
+    }
+    return hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    ).hexdigest()
+
+
+def _manifest_path(cache_root: str | os.PathLike[str]) -> Path:
+    return Path(cache_root) / MANIFEST_FILENAME
+
+
+def _load_manifest(cache_root: str | os.PathLike[str]) -> dict:
+    path = _manifest_path(cache_root)
+    if not path.exists():
+        return {"version": MANIFEST_VERSION, "entries": {}}
+    try:
+        data = json.loads(path.read_text())
+    except Exception:
+        return {"version": MANIFEST_VERSION, "entries": {}}
+    if data.get("version") != MANIFEST_VERSION or not isinstance(
+        data.get("entries"), dict
+    ):
+        return {"version": MANIFEST_VERSION, "entries": {}}
+    return data
+
+
+def should_skip_warmup(cache_root: str | os.PathLike[str], key: str) -> bool:
+    manifest = _load_manifest(cache_root)
+    entry = manifest.get("entries", {}).get(key)
+    if not entry or entry.get("status") != "complete":
+        return False
+
+    required = set(entry.get("cache_entries") or [])
+    if not required:
+        return False
+    return required.issubset(cache_entries(cache_root))
+
+
+def mark_warmup_complete(
+    cache_root: str | os.PathLike[str],
+    key: str,
+    entries: Iterable[str],
+) -> None:
+    root = Path(cache_root)
+    root.mkdir(parents=True, exist_ok=True)
+    manifest = _load_manifest(root)
+    manifest.setdefault("entries", {})[key] = {
+        "status": "complete",
+        "cache_entries": sorted(set(entries)),
+        "updated_at": time.time(),
+    }
+
+    fd, tmp_name = tempfile.mkstemp(
+        prefix=f".{MANIFEST_FILENAME}.", suffix=".tmp", dir=str(root)
+    )
+    try:
+        with os.fdopen(fd, "w") as f:
+            json.dump(manifest, f, sort_keys=True)
+            f.write("\n")
+        os.replace(tmp_name, _manifest_path(root))
+    finally:
+        if os.path.exists(tmp_name):
+            os.unlink(tmp_name)

python/sglang/srt/model_executor/model_runner.py

@@ -2473,7 +2473,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
 
     def init_deep_gemm(self):
-        logger.info("[DeepGEMM Debug] Entering init_deep_gemm")
+        logger.info("Entering DeepGEMM init")
         if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
             return
         try:
@@ -2528,8 +2528,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                         k,
                         num_groups,
                     )
-        # Also compile Dense GEMM kernels for other layers (e.g. Attention) if they use FP8 and DeepGEMM
+        # Also compile Dense GEMM kernels for other layers (e.g. Attention) if
-        for module in self.model.modules():
+        # they use FP8 and DeepGEMM.  Skip embed_tokens/lm_head because they do
+        # not dispatch through DeepGEMM at runtime, and restrict BF16 warmup to
+        # NSA indexer weights_proj; warming every BF16 Linear can include
+        # vocab-sized shapes and turn cache-hit startup into minutes of replay.
+        for name, module in self.model.named_modules():
+            if "embed_tokens" in name or "lm_head" in name:
+                continue
             if hasattr(module, "weight") and module.weight is not None:
                 # Check if it is FP8
                 is_fp8 = module.weight.dtype in [
@@ -2544,7 +2550,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                         k,
                         1,  # num_groups=1 for standard GEMM
                     )
-                if module.weight.dim() == 2 and module.weight.dtype == torch.bfloat16:
+                if (
+                    module.weight.dim() == 2
+                    and module.weight.dtype == torch.bfloat16
+                    and name.endswith("weights_proj")
+                ):
                     n, k = module.weight.shape
                     _maybe_compile_deep_gemm_one_type_all(
                         DeepGemmKernelType.GEMM_NT_BF16BF16F32,
@@ -2552,7 +2562,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                         k,
                         1,
                     )
-        # deep_gemm.preload_kernels()
+        if hasattr(deep_gemm, "preload_all_cached_kernels"):
+            logger.info("Preloading cached DeepGEMM kernels...")
+            deep_gemm.preload_all_cached_kernels()
+        elif hasattr(deep_gemm, "preload_kernels"):
+            logger.info("Preloading cached DeepGEMM kernels...")
+            deep_gemm.preload_kernels()
 
     def init_threads_binding(self):
         omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")

test/registered/unit/layers/test_deep_gemm_jit_cache.py

@@ -0,0 +1,135 @@
+import json
+import os
+import importlib.util
+import sys
+from pathlib import Path
+
+
+def _load_jit_cache_module():
+    repo_root = Path(__file__).resolve().parents[4]
+    module_path = (
+        repo_root
+        / "python"
+        / "sglang"
+        / "srt"
+        / "layers"
+        / "deep_gemm_wrapper"
+        / "jit_cache.py"
+    )
+    spec = importlib.util.spec_from_file_location("deep_gemm_jit_cache_under_test", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_prepare_deep_gemm_jit_env_sets_cache_before_import(tmp_path, monkeypatch):
+    prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env
+
+    cache_dir = tmp_path / "dg-cache"
+    monkeypatch.delenv("DG_JIT_CACHE_DIR", raising=False)
+    monkeypatch.delenv("DG_JIT_USE_NVRTC", raising=False)
+    monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(cache_dir))
+    monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")
+
+    resolved = prepare_deep_gemm_jit_env()
+
+    assert resolved == str(cache_dir)
+    assert os.environ["DG_JIT_CACHE_DIR"] == str(cache_dir)
+    assert os.environ["DG_JIT_USE_NVRTC"] == "1"
+    assert cache_dir.is_dir()
+
+
+def test_prepare_deep_gemm_jit_env_respects_explicit_deepgemm_env(
+    tmp_path, monkeypatch
+):
+    prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env
+
+    explicit = tmp_path / "explicit"
+    ignored = tmp_path / "ignored"
+    monkeypatch.setenv("DG_JIT_CACHE_DIR", str(explicit))
+    monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(ignored))
+    monkeypatch.setenv("DG_JIT_USE_NVRTC", "0")
+    monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")
+
+    resolved = prepare_deep_gemm_jit_env()
+
+    assert resolved == str(explicit)
+    assert os.environ["DG_JIT_CACHE_DIR"] == str(explicit)
+    assert os.environ["DG_JIT_USE_NVRTC"] == "0"
+    assert explicit.is_dir()
+    assert not ignored.exists()
+
+
+def test_sparse_m_list_is_sparse_but_covers_boundaries():
+    build_sparse_m_list = _load_jit_cache_module().build_sparse_m_list
+
+    m_list = build_sparse_m_list(65536)
+
+    assert m_list[0] == 1
+    assert m_list[-1] == 65536
+    assert 1024 in m_list
+    assert 2048 in m_list
+    assert 4096 in m_list
+    assert len(m_list) < 5000
+    assert len(m_list) < 65536 // 10
+
+
+def test_build_warmup_m_lists_is_category_and_cp_aware():
+    build_warmup_m_lists = _load_jit_cache_module().build_warmup_m_lists
+
+    lists = build_warmup_m_lists(
+        chunked_prefill_size=65536,
+        attn_cp_size=8,
+        max_running_requests=200,
+        speculative_num_draft_tokens=4,
+    )
+
+    assert lists.per_rank[-1] == 8192
+    assert lists.gathered[-1] == 65536
+    assert lists.decode[-1] == 800
+    assert len(lists.per_rank) < len(lists.gathered)
+    assert len(lists.gathered) < 5000
+
+
+def test_warmup_manifest_requires_cache_key_superset(tmp_path):
+    jit_cache = _load_jit_cache_module()
+    cache_entries = jit_cache.cache_entries
+    make_warmup_manifest_key = jit_cache.make_warmup_manifest_key
+    mark_warmup_complete = jit_cache.mark_warmup_complete
+    should_skip_warmup = jit_cache.should_skip_warmup
+
+    cache_root = tmp_path / "deep_gemm"
+    (cache_root / "cache" / "a").mkdir(parents=True)
+    (cache_root / "cache" / "b").mkdir()
+    key = make_warmup_manifest_key(
+        kernel_type="GEMM_NT_F8F8BF16",
+        n=4096,
+        k=8192,
+        num_groups=1,
+        m_list=[1, 2, 4, 8],
+        runtime_fingerprint={"deep_gemm": "0.1.2", "sm": 90},
+    )
+
+    mark_warmup_complete(cache_root, key, cache_entries(cache_root))
+    assert should_skip_warmup(cache_root, key)
+
+    # Deleting a cubin directory must invalidate the SGLang-side manifest.
+    (cache_root / "cache" / "b").rmdir()
+    assert not should_skip_warmup(cache_root, key)
+
+
+def test_warmup_manifest_file_is_json_and_atomic(tmp_path):
+    jit_cache = _load_jit_cache_module()
+    MANIFEST_FILENAME = jit_cache.MANIFEST_FILENAME
+    mark_warmup_complete = jit_cache.mark_warmup_complete
+
+    cache_root = tmp_path / "deep_gemm"
+    cache_root.mkdir()
+    mark_warmup_complete(cache_root, "shape-key", {"abc"})
+
+    manifest = json.loads((cache_root / MANIFEST_FILENAME).read_text())
+    assert manifest["version"] == 1
+    assert manifest["entries"]["shape-key"]["status"] == "complete"
+    assert manifest["entries"]["shape-key"]["cache_entries"] == ["abc"]