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

"""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)