diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 302aced46..64ef2a08a 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/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)) -# Separate, possibly-shrunk M grid for NON-grouped (dense/attention) shapes under -# CP prefill -- see _cp_dense_warmup_divisor / update_deep_gemm_config. -_BUILTIN_M_LIST_DENSE = _BUILTIN_M_LIST +_M_LIST_PER_RANK: List[int] = build_sparse_m_list(1024 * 16) +_M_LIST_GATHERED: List[int] = build_sparse_m_list(1024 * 16) +_M_LIST_DECODE: List[int] = build_sparse_m_list(1024) _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 -os.environ["DG_JIT_CACHE_DIR"] = os.getenv( - "SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm") +_DEEP_GEMM_CACHE_DIR = prepare_deep_gemm_jit_env( + preload_kernels=_ENABLE_JIT_DEEPGEMM_PRECOMPILE ) -# 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 _BUILTIN_M_LIST_DENSE + global _M_LIST_PER_RANK, _M_LIST_GATHERED, _M_LIST_DECODE global _DO_COMPILE_ALL global _IS_FIRST_RANK_ON_NODE - _BUILTIN_M_LIST = [] - - if _FAST_WARMUP: - # In fast warmup mode, only compile a small set of typical Ms - - # First cover all the small bs to ensure decode performance - _BUILTIN_M_LIST += list(range(1, 1025)) - - # Then cover larger batch sizes with gradually increasing steps - # 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 + m_lists = build_warmup_m_lists( + chunked_prefill_size=server_args.chunked_prefill_size, + attn_cp_size=getattr(server_args, "attn_cp_size", 1), + max_running_requests=server_args.max_running_requests, + speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, + ) + _M_LIST_PER_RANK = m_lists.per_rank + _M_LIST_GATHERED = m_lists.gathered + _M_LIST_DECODE = m_lists.decode _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 -# regardless of CP, so they always use the full M grid; non-grouped (dense) GEMMs -# are CP-per-rank and use the (possibly shrunk) dense grid. -_GROUPED_GEMM_KERNEL_TYPES = ( - DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED, - DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG, -) +def _m_list_for(kernel_type: DeepGemmKernelType) -> List[int]: + if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: + return _M_LIST_GATHERED + if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: + return _M_LIST_DECODE + 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); - # dense shapes are CP-per-rank and use the shrunk dense grid. - m_list = ( - _BUILTIN_M_LIST - if kernel_type in _GROUPED_GEMM_KERNEL_TYPES - else _BUILTIN_M_LIST_DENSE + m_list = _m_list_for(kernel_type) + manifest_key = make_warmup_manifest_key( + kernel_type=kernel_type.name, + n=n, + k=k, + 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 diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 34494f599..ba30fe465 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/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: diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/jit_cache.py b/python/sglang/srt/layers/deep_gemm_wrapper/jit_cache.py new file mode 100644 index 000000000..0557f8e0b --- /dev/null +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 02330c0e9..63b1ef54b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 - for module in self.model.modules(): + # Also compile Dense GEMM kernels for other layers (e.g. Attention) if + # 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") diff --git a/test/registered/unit/layers/test_deep_gemm_jit_cache.py b/test/registered/unit/layers/test_deep_gemm_jit_cache.py new file mode 100644 index 000000000..3de75cf40 --- /dev/null +++ b/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"]