diff --git a/python/sglang/kernels/jit/utils/compile.py b/python/sglang/kernels/jit/utils/compile.py deleted file mode 100644 index 107dedb06..000000000 --- a/python/sglang/kernels/jit/utils/compile.py +++ /dev/null @@ -1,341 +0,0 @@ -"""JIT compilation: load_jit, the build cache, and C++ template arguments.""" - -from __future__ import annotations - -import hashlib -import importlib.util -import logging -import os -import pathlib -import re -from contextlib import contextmanager -from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union - -import torch - -from sglang.kernels.jit.utils.arch import get_default_target_flags, get_jit_cuda_arch -from sglang.kernels.jit.utils.common import cache_once, is_hip_runtime -from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES - -if TYPE_CHECKING: - from tvm_ffi import Module - -logger = logging.getLogger(__name__) - - -_QUOTED_INCLUDE_RE = re.compile(r'^\s*#\s*include\s*"([^"]+)"', re.MULTILINE) -_ANGLE_INCLUDE_RE = re.compile(r"^\s*#\s*include\s*<(sgl_kernel/[^>]+)>", re.MULTILINE) - - -def _local_jit_source_hash(source_files: List[str]) -> str: - """Hash JIT source contents so TVM-FFI cache keys track included headers.""" - digest = hashlib.sha256() - seen: set[pathlib.Path] = set() - stack = [pathlib.Path(path).resolve() for path in source_files] - include_dir = KERNEL_PATH / "include" - - while stack: - path = stack.pop() - if path in seen or not path.is_file(): - continue - seen.add(path) - - data = path.read_bytes() - # Relative to kernel root, not absolute: the key must track source - # content, not install location (differs across runners / job dirs). - try: - ident = str(path.relative_to(KERNEL_PATH)) - except ValueError: - ident = path.name - digest.update(ident.encode()) - digest.update(b"\0") - digest.update(data) - digest.update(b"\0") - - text = data.decode("utf-8", errors="ignore") - for include in _QUOTED_INCLUDE_RE.findall(text): - include_path = (path.parent / include).resolve() - if include_path.is_file(): - stack.append(include_path) - for include in _ANGLE_INCLUDE_RE.findall(text): - include_path = (include_dir / include).resolve() - if include_path.is_file(): - stack.append(include_path) - - return digest.hexdigest()[:16] - - -@cache_once -def _resolve_kernel_path() -> pathlib.Path: - # Resolve via the package spec so the lookup is location-independent. - # The C++/CUDA sources (``csrc/``) and headers (``include/``) live next to - # this package under ``sglang.kernels.jit`` (RFC #29630). - spec = importlib.util.find_spec("sglang.kernels.jit") - assert spec is not None and spec.origin is not None - cur_dir = pathlib.Path(spec.origin).parent.resolve() - - # first, try this directory structure - def _environment_install(): - candidate = cur_dir.resolve() - if (candidate / "include").exists() and (candidate / "csrc").exists(): - return candidate - return None - - def _package_install(): - # TODO: support find path by package - return None - - path = _environment_install() or _package_install() - if path is None: - raise RuntimeError("Cannot find sglang.kernels.jit path") - return path - - -KERNEL_PATH = _resolve_kernel_path() -DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] -DEFAULT_CFLAGS = ["-std=c++20", "-O3"] -DEFAULT_LDFLAGS = [] -CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, str, bool, torch.dtype] - - -class CPPArgList(list[str]): - def __str__(self) -> str: - return ", ".join(self) - - -CPP_DTYPE_MAP = { - torch.float64: "double", - torch.float32: "fp32_t", - torch.float16: "fp16_t", - torch.bfloat16: "bf16_t", - # The fnuz variants are the ROCm-side torch dtypes; fp8_*_t resolves to - # the matching HIP type there (see HIP_FP8_TYPE_* in utils.cuh). - torch.float8_e4m3fn: "fp8_e4m3_t", - torch.float8_e4m3fnuz: "fp8_e4m3_t", - torch.float8_e5m2: "fp8_e5m2_t", - torch.float8_e5m2fnuz: "fp8_e5m2_t", - torch.int8: "int8_t", - torch.int16: "int16_t", - torch.int32: "int32_t", - torch.int64: "int64_t", - torch.uint8: "uint8_t", - torch.uint16: "uint16_t", - torch.uint32: "uint32_t", - torch.uint64: "uint64_t", - torch.bool: "bool", -} - - -def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: - def _convert(arg: CPP_TEMPLATE_TYPE) -> str: - if isinstance(arg, bool): - return "true" if arg else "false" - if isinstance(arg, (int, str, float)): - return str(arg) - if isinstance(arg, torch.dtype): - return CPP_DTYPE_MAP[arg] - raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}") - - return CPPArgList(_convert(arg) for arg in args) - - -@cache_once -def _tvm_ffi_version() -> str: - try: - import tvm_ffi - - version = getattr(tvm_ffi, "__version__", None) - if version: - return str(version) - except Exception: - pass - try: - from importlib.metadata import version as dist_version - - return dist_version("apache-tvm-ffi") - except Exception: - return "unknown" - - -def _jit_build_dir_name(module_name: str) -> str: - # Key on arch + tvm-ffi ABI too (module_name only hashes sources), so a - # shared cache volume never reuses a cross-arch/ABI .so. - arch = get_jit_cuda_arch().target_name - return f"{module_name}__arch_{arch}__tvmffi_{_tvm_ffi_version()}" - - -def _make_wrapper(tup: Tuple[str, str]) -> str: - export_name, kernel_name = tup - return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" - - -def _make_sources(files: List[str], wrappers: List[Tuple[str, str]]) -> List[str]: - sources = [f'#include "{path}"' for path in files] - sources += ["namespace sglang {"] - sources += [_make_wrapper(tup) for tup in wrappers] - sources += ["} // namespace sglang"] - return sources - - -# JIT compilation is pure Python/filesystem plumbing (path `.resolve()` calls -# `os.lstat`, etc.) that Dynamo cannot trace. When a lazily-loaded kernel is -# first reached from inside a `@torch.compile`d region, tracing into it produces -# spurious "Dynamo does not know how to trace the builtin `posix.lstat`" graph -# breaks. The load happens once and is memoized, so keep it out of the graph. -@torch.compiler.disable -def load_jit( - *args: str, - cpp_files: List[str] | None = None, - cuda_files: List[str] | None = None, - external_cpp_files: List[str] | None = None, - external_cuda_files: List[str] | None = None, - cpp_wrappers: List[Tuple[str, str]] | None = None, - cuda_wrappers: List[Tuple[str, str]] | None = None, - extra_cflags: List[str] | None = None, - extra_cuda_cflags: List[str] | None = None, - extra_ldflags: List[str] | None = None, - extra_include_paths: List[str] | None = None, - extra_dependencies: List[str] | None = None, - build_directory: str | None = None, - header_only: bool = True, -) -> Module: - """ - Loading a JIT module from C++/CUDA source files. - We define a wrapper as a tuple of (export_name, kernel_name), - where `export_name` is the name used to called from Python, - and `kernel_name` is the name of the kernel class in C++/CUDA source. - - :param args: Unique marker of the JIT module. Must be distinct for different kernels. - :type args: str - :param cpp_files: A list of C++ source files. - :type cpp_files: List[str] | None - :param cuda_files: A list of CUDA source files. - :type cuda_files: List[str] | None - :param external_cpp_files: A list of caller-resolved C++ source paths outside - the in-tree JIT source directory. - :type external_cpp_files: List[str] | None - :param external_cuda_files: A list of caller-resolved CUDA source paths outside - the in-tree JIT source directory. - :type external_cuda_files: List[str] | None - :param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name. - :type cpp_wrappers: List[Tuple[str, str]] | None - :param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name. - :type cuda_wrappers: List[Tuple[str, str]] | None - :param extra_cflags: Extra C++ compiler flags. - :type extra_cflags: List[str] | None - :param extra_cuda_cflags: Extra CUDA compiler flags. - :type extra_cuda_cflags: List[str] | None - :param extra_ldflags: Extra linker flags. - :type extra_ldflags: List[str] | None - :param extra_include_paths: Extra include paths. - :type extra_include_paths: List[str] | None - :param extra_dependencies: Extra dependencies for the JIT module, e.g., cutlass. - :type extra_dependencies: List[str] | None - :param build_directory: The build directory for JIT compilation. - :type build_directory: str | None - :param header_only: Whether the module is header-only. - If true, apply the wrappers to export given class/functions. - Otherwise, we must export from C++/CUDA side. - :return: A just-in-time(JIT) compiled module. - :rtype: Module - """ - - from tvm_ffi.cpp import load, load_inline - - cpp_files = cpp_files or [] - cuda_files = cuda_files or [] - external_cpp_files = external_cpp_files or [] - external_cuda_files = external_cuda_files or [] - extra_cflags = extra_cflags or [] - extra_cuda_cflags = extra_cuda_cflags or [] - extra_ldflags = extra_ldflags or [] - extra_include_paths = extra_include_paths or [] - - if torch.version.hip is not None: - extra_cuda_cflags = [ - flag - for flag in extra_cuda_cflags - if flag not in ("--use_fast_math", "-use_fast_math") - ] - - cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] + [ - str(pathlib.Path(f).resolve()) for f in external_cpp_files - ] - cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] + [ - str(pathlib.Path(f).resolve()) for f in external_cuda_files - ] - - for dep in set(extra_dependencies or []): - if dep not in REGISTERED_DEPENDENCIES: - raise ValueError(f"Dependency {dep} is not registered.") - extra_include_paths += REGISTERED_DEPENDENCIES[dep]() - - module_name = "sgl_kernel_jit_" + "_".join(str(arg) for arg in args) - if cpp_files or cuda_files: - module_name += "_" + _local_jit_source_hash(cpp_files + cuda_files) - - # A built .so under a deterministic dir is content-addressed: load it - # directly to skip ninja, whose mtime check rebuilds every CI run (pip - # install bumps dep header mtimes). - if build_directory is None: - cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi") - build_directory = str( - pathlib.Path(cache_dir).expanduser() / _jit_build_dir_name(module_name) - ) - prebuilt = pathlib.Path(build_directory) / f"{module_name}.so" - if prebuilt.is_file(): - from tvm_ffi import load_module - - try: - module = load_module(str(prebuilt)) - logger.debug("Reused cached JIT module %s", module_name) - return module - except Exception: - logger.warning( - "Cached JIT module %s failed to load; rebuilding.", module_name - ) - - if header_only: - cpp_sources = _make_sources(cpp_files, cpp_wrappers or []) - cuda_sources = _make_sources(cuda_files, cuda_wrappers or []) - with _jit_compile_context(): - return load_inline( - module_name, - cpp_sources=cpp_sources, - cuda_sources=cuda_sources, - extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags, - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, - extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, - build_directory=build_directory, - ) - else: - assert cpp_wrappers is None and cuda_wrappers is None - with _jit_compile_context(): - return load( - module_name, - cpp_files=cpp_files, - cuda_files=cuda_files, - extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags, - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, - extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, - build_directory=build_directory, - ) - - -@contextmanager -def _jit_compile_context(): - if is_hip_runtime(): - yield # TODO: support ROCm `TVM_FFI_ROCM_ARCH_LIST` if needed - return - env_key = "TVM_FFI_CUDA_ARCH_LIST" - old_value = os.environ.get(env_key, None) - os.environ[env_key] = get_jit_cuda_arch().target_name - try: - yield - finally: - if old_value is None: - os.environ.pop(env_key, None) - else: - os.environ[env_key] = old_value diff --git a/python/sglang/kernels/jit/utils/compile/__init__.py b/python/sglang/kernels/jit/utils/compile/__init__.py new file mode 100644 index 000000000..f3cf99c12 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/__init__.py @@ -0,0 +1,44 @@ +"""JIT compilation: source layout, ninja generation, the build cache, load_jit. + +The package owns the whole path from a ``load_jit`` call to a loaded module, +including the ``build.ninja`` it compiles through — only tvm-ffi's headers, +its shared library, and ``tvm_ffi.load_module`` are consumed from outside. + +Modules, in dependency order (there are no cycles):: + + paths where the in-tree sources live, and the default flags + cpp_args rendering Python values as C++ template arguments + spec BuildSpec: one fully-resolved build + toolchain compilers, tvm-ffi locations, platform base flags + ninja generating and running build.ninja, reading its depfiles + cache build_key / deps_key, cache layout, publication + loader load_jit +""" + +from sglang.kernels.jit.utils.compile.cpp_args import ( + CPP_DTYPE_MAP, + CPP_TEMPLATE_TYPE, + CPPArgList, + make_cpp_args, +) +from sglang.kernels.jit.utils.compile.loader import load_jit +from sglang.kernels.jit.utils.compile.paths import ( + DEFAULT_CFLAGS, + DEFAULT_INCLUDE, + DEFAULT_LDFLAGS, + KERNEL_PATH, +) +from sglang.kernels.jit.utils.compile.spec import BuildSpec + +__all__ = [ + "BuildSpec", + "CPPArgList", + "CPP_DTYPE_MAP", + "CPP_TEMPLATE_TYPE", + "DEFAULT_CFLAGS", + "DEFAULT_INCLUDE", + "DEFAULT_LDFLAGS", + "KERNEL_PATH", + "load_jit", + "make_cpp_args", +] diff --git a/python/sglang/kernels/jit/utils/compile/cache.py b/python/sglang/kernels/jit/utils/compile/cache.py new file mode 100644 index 000000000..78b664343 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/cache.py @@ -0,0 +1,522 @@ +"""Content-addressed JIT build cache: key derivation, layout, publication. + +The cache answers one question on every ``load_jit``: *is there an already-built +``.so`` that is guaranteed to be identical to what a build right now would +produce?* It does so with two keys, because the full answer is not computable +before the first build: + +``build_key`` — everything known *before* compiling: the module args, the + caller's flags, the wrapper exports, the compile target, and the contents + of the direct source files. It selects a directory. + +``deps_key`` — the contents of the *transitive* dependency closure, which only + the compiler can enumerate. It selects a leaf inside that directory. + +Layout:: + + $SGLANG_JIT_CACHE_DIR///build-/ + deps-/ one leaf per transitive-content state + sgl_deps.json the dependency list of *this* build + .so + .staging-/ a build in progress + +```` and ```` are for humans (``du -sh`` per arch, +"what is this directory") and carry no correctness weight — every distinguishing +input is folded into ``build_key`` as well. + +**Each leaf carries its own dependency list and is never modified after it is +published.** A leaf is a hit only when re-hashing its own list against the files +as they stand *now* reproduces the leaf's own name. That makes the recorded data +self-verifying rather than trusted: a truncated, tampered, or differently +formatted list simply fails to reproduce the name, and a leaf published by a +machine whose dependency graph differs is skipped instead of poisoning this one. +There is deliberately no shared, mutable manifest for writers to merge into. + +The known gap is ``__has_include``-style constructs, where the dependency graph +turns on a file's *existence* rather than any listed file's content. The +compiler-version and package-version components of ``build_key`` cover the +realistic instances; the residue is accepted (ccache has carried the same gap +for decades). +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import logging +import os +import pathlib +import shutil +import subprocess +import sysconfig +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as dist_version +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import msgspec + +from sglang.kernels.jit.utils.arch import get_jit_cuda_arch +from sglang.kernels.jit.utils.common import cache_once, is_hip_runtime, is_musa_runtime +from sglang.kernels.jit.utils.compile import toolchain +from sglang.kernels.jit.utils.compile.paths import KERNEL_PATH +from sglang.kernels.jit.utils.compile.spec import BuildSpec +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + +_DEPS_FILE = "sgl_deps.json" +_BUILD_KEY_PREFIX = "build-" +_DEPS_KEY_PREFIX = "deps-" +_KEY_HEX_LEN = 16 +# Leaves are examined newest-first and the search stops at the first match, so +# the common case costs one read. The cap only bounds the pathological tail: +# 100 leaves all missing costs ~32 ms, and a run that reaches it is about to +# spend seconds rebuilding anyway. +_MAX_LEAVES_SCANNED = 100 + + +class _DepEntry(msgspec.Struct, frozen=True, array_like=True): + """One dependency, stored install-location independently. + + ``root`` is an anchor token resolved against the *current* environment, so a + list written by one clone is readable from another. + """ + + root: str + relpath: str + digest: str + + +# --------------------------------------------------------------------------- +# Anchor roots — the only reason a dependency list survives a re-clone +# --------------------------------------------------------------------------- + + +def _package_dir(package: str) -> Optional[pathlib.Path]: + try: + spec = importlib.util.find_spec(package) + except (ImportError, ValueError): + return None + if spec is None: + return None + if spec.origin is not None: + return pathlib.Path(spec.origin).resolve().parent + locations = list(spec.submodule_search_locations or []) + return pathlib.Path(locations[0]).resolve() if locations else None + + +@cache_once +def _anchor_roots() -> Tuple[Tuple[str, pathlib.Path], ...]: + """Anchor tokens ordered most-specific first, so the longest match wins.""" + candidates: List[Tuple[str, Optional[pathlib.Path]]] = [ + ("kernels", KERNEL_PATH), + ("tvm_ffi", _package_dir("tvm_ffi")), + ("pkg:flashinfer", _package_dir("flashinfer")), + ("pkg:deep_gemm", _package_dir("deep_gemm")), + ("pkg:nvidia", _package_dir("nvidia")), + ("toolkit", toolchain.toolkit_home()), + ("sitepkgs", pathlib.Path(sysconfig.get_paths()["purelib"])), + ("sys", pathlib.Path("/usr")), + ] + # Resolved, because the paths being matched against them are resolved too. + # `/usr/local/cuda` is a symlink to `/usr/local/cuda-`; leaving the + # anchor unresolved makes every toolkit header miss it and fall through to + # `sys`, which bakes the CUDA version into the recorded relpath. + roots = [ + (token, path.resolve()) + for token, path in candidates + if path is not None and path.exists() + ] + roots.sort(key=lambda item: len(str(item[1])), reverse=True) + return tuple(roots) + + +def _normalize_path(path: pathlib.Path) -> Tuple[str, str]: + """Split *path* into ``(anchor token, path relative to that anchor)``. + + Falls back to ``("abs", )`` for anything outside every known + root — correct, just not portable across machines (which only costs a miss). + """ + for token, root in _anchor_roots(): + try: + return token, str(path.relative_to(root)) + except ValueError: + continue + return "abs", str(path) + + +def _resolve_path(*, root: str, relpath: str) -> Optional[pathlib.Path]: + if root == "abs": + return pathlib.Path(relpath) + for token, base in _anchor_roots(): + if token == root: + return base / relpath + return None + + +# --------------------------------------------------------------------------- +# Content digests +# --------------------------------------------------------------------------- + +_digest_cache: Dict[pathlib.Path, Optional[str]] = {} + + +def _file_digest(path: pathlib.Path) -> Optional[str]: + """Content digest of *path*, or None if it cannot be read. + + Memoized per process: a server hashes the same CUTLASS headers for every + kernel it loads, and the union is ~25 MB. The memo means a source edited + while the process is alive is not noticed, which is fine — modules are + resolved once at startup and never re-resolved. + """ + cached = _digest_cache.get(path) + if cached is not None or path in _digest_cache: + return cached + try: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + except OSError: + digest = None + _digest_cache[path] = digest + return digest + + +def clear_digest_cache() -> None: + """Drop the memo. For tests that mutate sources between lookups.""" + _digest_cache.clear() + + +def _hash_parts(parts: Iterable[object]) -> str: + digest = hashlib.sha256() + for part in parts: + digest.update(repr(part).encode()) + digest.update(b"\0") + return digest.hexdigest()[:_KEY_HEX_LEN] + + +# --------------------------------------------------------------------------- +# build_key +# --------------------------------------------------------------------------- + +_VERSIONED_PACKAGES = ( + "apache-tvm-ffi", + "flashinfer-python", + "deep_gemm", + "nvidia-mathdx", + "torch", +) + + +@cache_once +def _target_tag() -> str: + """Short, human-readable target name used as a cache path segment. + + HIP deliberately uses ``gcnArchName`` rather than the CUDA-shaped + ``(major, minor)`` capability: the latter maps gfx940/gfx941/gfx942 onto a + single ``9.4``, which are three different compile targets. + """ + if is_hip_runtime(): + return toolchain.gpu_arch_name().split(":")[0] or "unknown" + arch = get_jit_cuda_arch() + prefix = "mp" if is_musa_runtime() else "sm" + return f"{prefix}{arch.major}{arch.minor}{arch.suffix}" + + +@cache_once +def _environment_fingerprint() -> str: + """Process-wide part of ``build_key``: target, compilers, package versions. + + The target here is the unabridged one — unlike ``_target_tag`` it keeps the + gfx feature suffixes, because this half has to be exact. + + Both compilers are fingerprinted: nvcc hands all host code to ``c++``, so a + different host compiler means different system headers and different host + codegen for otherwise identical inputs. + """ + if is_hip_runtime(): + target = f"hip:{toolchain.gpu_arch_name()}" + else: + arch = get_jit_cuda_arch() + target = f"{'musa' if is_musa_runtime() else 'cuda'}:{arch.target_name}" + + compilers = [] + for path in (toolchain.device_compiler_path(), toolchain.host_compiler_path()): + try: + compilers.append(subprocess.check_output([path, "--version"], text=True)) + except (OSError, subprocess.SubprocessError) as error: + logger.warning("Cannot fingerprint compiler %s: %s", path, error) + compilers.append("unknown") + + versions: List[Tuple[str, str]] = [] + for name in _VERSIONED_PACKAGES: + try: + versions.append((name, dist_version(name))) + except (PackageNotFoundError, ValueError): + versions.append((name, "absent")) + + return _hash_parts([target, compilers, versions]) + + +def compute_build_key(spec: BuildSpec, *, build_file: str) -> str: + """Everything that is knowable before the compiler runs. + + Two things are hashed, and together they are everything the compiler sees: + *build_file* — the generated ninja text, so no flag can reach the compiler + without reaching the key — and the translation units, whose generated + wrapper source carries the exports that never appear in the build file. + Absolute paths in both are anchor-normalized, so the same tree under a + different clone directory still keys the same. + """ + source_digests = [ + ( + _normalize_path(pathlib.Path(path).resolve()), + _file_digest(pathlib.Path(path).resolve()), + ) + for path in sorted(spec.sources) + ] + units = [ + (unit.filename if unit.source is None else "", unit.is_cuda, unit.source) + for unit in spec.translation_units() + ] + return _hash_parts( + [ + spec.module_args, + tuple(source_digests), + _normalize_text(build_file), + _normalize_text(repr(units)), + _environment_fingerprint(), + ] + ) + + +def _normalize_text(text: str) -> str: + """Replace every known root in *text* with its anchor token. + + Longest root first, so `/usr/local/cuda-12.9` wins over `/usr`. + """ + for token, root in _anchor_roots(): + text = text.replace(str(root), f"<{token}>") + return text + + +def build_key_dir(*, module_name: str, build_key: str) -> 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}" + + +# --------------------------------------------------------------------------- +# deps_key — a leaf that reproduces its own name +# --------------------------------------------------------------------------- + + +def _deps_key(entries: Sequence[_DepEntry]) -> str: + return _hash_parts([(e.root, e.relpath, e.digest) for e in entries]) + + +def _read_deps(leaf: pathlib.Path) -> Optional[List[_DepEntry]]: + try: + raw = (leaf / _DEPS_FILE).read_bytes() + except OSError: + return None + try: + return msgspec.json.decode(raw, type=List[_DepEntry]) + except msgspec.DecodeError: + return None + + +def _refresh( + entries: Sequence[_DepEntry], +) -> Tuple[Optional[List[_DepEntry]], List[str]]: + """Re-read every recorded dependency as it stands now. + + Returns ``(entries with current digests, names that changed)``; the entries + are None when a recorded dependency has vanished, which is itself a change. + """ + current: List[_DepEntry] = [] + changed: List[str] = [] + for entry in entries: + path = _resolve_path(root=entry.root, relpath=entry.relpath) + digest = _file_digest(path) if path is not None else None + if digest is None: + return None, [f"{entry.root}:{entry.relpath}"] + if digest != entry.digest: + changed.append(f"{entry.root}:{entry.relpath}") + current.append(_DepEntry(root=entry.root, relpath=entry.relpath, digest=digest)) + return current, changed + + +def find_prebuilt(*, scope: pathlib.Path, module_name: str) -> Optional[pathlib.Path]: + """The leaf whose recorded dependencies still hash to its own name, if any. + + Newest leaves are examined first, so the common case costs one read. + """ + try: + leaves = sorted( + (p for p in scope.iterdir() if p.name.startswith(_DEPS_KEY_PREFIX)), + key=lambda p: p.stat().st_mtime, + reverse=True, + ) + except OSError: + return None + if len(leaves) > _MAX_LEAVES_SCANNED: + logger.debug( + "JIT module %s has %d cached leaves; only the newest %d are considered.", + module_name, + len(leaves), + _MAX_LEAVES_SCANNED, + ) + + reason: List[str] = [] + for leaf in leaves[:_MAX_LEAVES_SCANNED]: + entries = _read_deps(leaf) + if entries is None: + continue + current, changed = _refresh(entries) + if current is None: + reason = reason or changed + continue + if f"{_DEPS_KEY_PREFIX}{_deps_key(current)}" != leaf.name: + reason = reason or changed + continue + candidate = leaf / f"{module_name}.so" + if candidate.is_file(): + # Keep the mtime ordering meaningful: it is what puts this leaf + # first next time, and what a size-bounded GC would evict by. + # Bookkeeping only, so it must never turn a hit into a failure: the + # cache root can be a read-only mount, and a prune racing the + # is_file() above leaves nothing here to touch. + try: + os.utime(leaf, None) + except OSError: + pass + return candidate + + if reason: + log = logger.info if envs.SGLANG_JIT_CACHE_DEBUG.get() else logger.debug + log("Rebuilding JIT module %s: %s changed", module_name, reason[0]) + return None + + +# --------------------------------------------------------------------------- +# Publication +# --------------------------------------------------------------------------- + + +def commit_build( + spec: BuildSpec, + *, + scope: pathlib.Path, + staging: pathlib.Path, + dependencies: Sequence[pathlib.Path], +) -> Optional[pathlib.Path]: + """Publish a freshly built *staging* directory into the cache. + + Never raises: publication is opportunistic, and a module that fails to + publish is merely rebuilt next time. + """ + entries = _to_entries(dependencies=dependencies, build_dir=staging) + if not _covers_direct_sources(entries=entries, direct_sources=spec.sources): + logger.warning( + "JIT module %s produced no usable dependency information; it will be " + "rebuilt on every load. This is a build-rule problem, not a cache problem.", + spec.module_name, + ) + return None + + # Written *before* the rename, so the leaf carries its own list the moment + # it becomes visible. Nothing ever rewrites it afterwards. + (staging / _DEPS_FILE).write_bytes(msgspec.json.encode(entries)) + leaf = _publish( + staging=staging, leaf=scope / f"{_DEPS_KEY_PREFIX}{_deps_key(entries)}" + ) + _prune(scope=scope, keep_newest=leaf) + return leaf / f"{spec.module_name}.so" + + +def _prune(*, scope: pathlib.Path, keep_newest: pathlib.Path) -> None: + """Drop the oldest builds of this variant past ``SGLANG_JIT_CACHE_KEEP``. + + Unset keeps everything, which is what makes reverting an edit an instant + hit rather than a rebuild — the leaves *are* the history. + + Deleting a leaf another process is using is safe: an unlinked ``.so`` stays + mapped for whoever already loaded it, and a lookup that loses its leaf + mid-flight falls through to a rebuild. + """ + keep = envs.SGLANG_JIT_CACHE_KEEP.get() + if keep is None: + return + leaves = sorted( + ( + path + for path in scope.iterdir() + if path.name.startswith(_DEPS_KEY_PREFIX) and path != keep_newest + ), + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + for stale in leaves[max(keep - 1, 0) :]: + logger.debug("Pruning JIT build %s/%s", scope.name, stale.name) + shutil.rmtree(stale, ignore_errors=True) + + +def _to_entries( + *, dependencies: Sequence[pathlib.Path], build_dir: pathlib.Path +) -> List[_DepEntry]: + """Turn scanned paths into portable, sorted entries. + + Files under *build_dir* (its own generated units and objects) are dropped: + their paths are unstable, and their contents are already a function of + inputs ``build_key`` covers. + """ + seen: Dict[Tuple[str, str], _DepEntry] = {} + for candidate in dependencies: + path = ( + candidate.resolve() + if candidate.is_absolute() + else (build_dir / candidate).resolve() + ) + if path.is_relative_to(build_dir) or not path.is_file(): + continue + digest = _file_digest(path) + if digest is None: + continue + root, relpath = _normalize_path(path) + seen[(root, relpath)] = _DepEntry(root=root, relpath=relpath, digest=digest) + return [seen[key] for key in sorted(seen)] + + +def _covers_direct_sources( + *, entries: Sequence[_DepEntry], direct_sources: Sequence[str] +) -> bool: + """Reject a dependency list that does not even mention the direct sources. + + A truncated or empty scan would otherwise narrow the checked set below what + ``build_key`` already covers, which is the one way bad recorded data could + cause reuse instead of a rebuild. + """ + if not entries: + return False + recorded = {(entry.root, entry.relpath) for entry in entries} + return all( + _normalize_path(pathlib.Path(source).resolve()) in recorded + for source in direct_sources + ) + + +def _publish(*, staging: pathlib.Path, leaf: pathlib.Path) -> pathlib.Path: + """Move *staging* into place so the leaf appears complete or not at all. + + A directory rename is atomic; losing the race means another process built + the identical content first, so its result is used and ours is discarded. + Kept as its own function because that "failure is success" branch is the + part worth pinning with a test. + """ + leaf.parent.mkdir(parents=True, exist_ok=True) + try: + os.rename(staging, leaf) + except OSError: + if not leaf.is_dir(): + raise + logger.debug("JIT leaf %s already published by another process", leaf.name) + return leaf diff --git a/python/sglang/kernels/jit/utils/compile/cpp_args.py b/python/sglang/kernels/jit/utils/compile/cpp_args.py new file mode 100644 index 000000000..584d2c866 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/cpp_args.py @@ -0,0 +1,50 @@ +"""Rendering Python values as C++ template arguments.""" + +from __future__ import annotations + +from typing import TypeAlias, Union + +import torch + +CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, str, bool, torch.dtype] + + +class CPPArgList(list): + def __str__(self) -> str: + return ", ".join(self) + + +CPP_DTYPE_MAP = { + torch.float64: "double", + torch.float32: "fp32_t", + torch.float16: "fp16_t", + torch.bfloat16: "bf16_t", + # The fnuz variants are the ROCm-side torch dtypes; fp8_*_t resolves to + # the matching HIP type there (see HIP_FP8_TYPE_* in utils.cuh). + torch.float8_e4m3fn: "fp8_e4m3_t", + torch.float8_e4m3fnuz: "fp8_e4m3_t", + torch.float8_e5m2: "fp8_e5m2_t", + torch.float8_e5m2fnuz: "fp8_e5m2_t", + torch.int8: "int8_t", + torch.int16: "int16_t", + torch.int32: "int32_t", + torch.int64: "int64_t", + torch.uint8: "uint8_t", + torch.uint16: "uint16_t", + torch.uint32: "uint32_t", + torch.uint64: "uint64_t", + torch.bool: "bool", +} + + +def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: + def _convert(arg: CPP_TEMPLATE_TYPE) -> str: + if isinstance(arg, bool): + return "true" if arg else "false" + if isinstance(arg, (int, str, float)): + return str(arg) + if isinstance(arg, torch.dtype): + return CPP_DTYPE_MAP[arg] + raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}") + + return CPPArgList(_convert(arg) for arg in args) diff --git a/python/sglang/kernels/jit/utils/compile/loader.py b/python/sglang/kernels/jit/utils/compile/loader.py new file mode 100644 index 000000000..dca05aba0 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/loader.py @@ -0,0 +1,200 @@ +"""``load_jit``: resolve a request, reuse a cached build, or make one.""" + +from __future__ import annotations + +import contextlib +import fcntl +import logging +import os +import pathlib +import shutil +import uuid +from typing import TYPE_CHECKING, List, Tuple + +import torch + +from sglang.kernels.jit.utils.arch import get_default_target_flags +from sglang.kernels.jit.utils.common import is_hip_runtime +from sglang.kernels.jit.utils.compile import cache, ninja +from sglang.kernels.jit.utils.compile.paths import ( + DEFAULT_CFLAGS, + DEFAULT_INCLUDE, + DEFAULT_LDFLAGS, +) +from sglang.kernels.jit.utils.compile.spec import BuildSpec, resolve_sources +from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES + +if TYPE_CHECKING: + from tvm_ffi import Module + + _DISABLE_TORCH_COMPILE = lambda f: f +else: + # NOTE: this is not friendly to type checking + _DISABLE_TORCH_COMPILE = torch.compiler.disable + + +logger = logging.getLogger(__name__) + +_LOCK_FILE = ".lock" + + +# JIT compilation is pure Python/filesystem plumbing (path `.resolve()` calls +# `os.lstat`, etc.) that Dynamo cannot trace. When a lazily-loaded kernel is +# first reached from inside a `@torch.compile`d region, tracing into it produces +# spurious "Dynamo does not know how to trace the builtin `posix.lstat`" graph +# breaks. The load happens once and is memoized, so keep it out of the graph. +@_DISABLE_TORCH_COMPILE +def load_jit( + *args: str, + cpp_files: List[str] | None = None, + cuda_files: List[str] | None = None, + cpp_wrappers: List[Tuple[str, str]] | None = None, + cuda_wrappers: List[Tuple[str, str]] | None = None, + extra_cflags: List[str] | None = None, + extra_cuda_cflags: List[str] | None = None, + extra_ldflags: List[str] | None = None, + extra_include_paths: List[str] | None = None, + extra_dependencies: List[str] | None = None, + header_only: bool = True, +) -> Module: + """Load a JIT module, compiling it if no cached build still applies. + + A wrapper is a ``(export_name, kernel_name)`` pair: ``export_name`` is what + Python calls, ``kernel_name`` is the C++ class or function it resolves to. + + :param args: Unique marker of the module. Must differ between kernels. + :param cpp_files: C++ sources. Relative names resolve against the in-tree + ``csrc/`` directory; pass an absolute path for a source + outside the tree. + :param cuda_files: CUDA sources, resolved like `cpp_files`. + :param cpp_wrappers: C++ exports to generate. + :param cuda_wrappers: CUDA exports to generate. + :param extra_cflags: Extra host compiler flags. + :param extra_cuda_cflags: Extra device compiler flags. + :param extra_ldflags: Extra linker flags. + :param extra_include_paths: Extra include paths. + :param extra_dependencies: Registered header-only dependencies, e.g. cutlass. + :param header_only: Compile through a generated wrapper that exports the + given entry points. Otherwise the sources must export + from the C++ side themselves. + """ + if is_hip_runtime(): + extra_cuda_cflags = [ + flag + for flag in (extra_cuda_cflags or []) + if flag not in ("--use_fast_math", "-use_fast_math") + ] + + includes = list(DEFAULT_INCLUDE) + (extra_include_paths or []) + for dep in sorted(set(extra_dependencies or [])): + if dep not in REGISTERED_DEPENDENCIES: + raise ValueError(f"Dependency {dep} is not registered.") + includes += REGISTERED_DEPENDENCIES[dep]() + + module_args = tuple(str(arg) for arg in args) + spec = BuildSpec( + module_args=module_args, + cpp_files=resolve_sources(cpp_files), + cuda_files=resolve_sources(cuda_files), + cpp_wrappers=tuple(cpp_wrappers or ()), + cuda_wrappers=tuple(cuda_wrappers or ()), + cflags=tuple(DEFAULT_CFLAGS + (extra_cflags or [])), + cuda_cflags=tuple(get_default_target_flags() + (extra_cuda_cflags or [])), + ldflags=tuple(DEFAULT_LDFLAGS + (extra_ldflags or [])), + include_paths=tuple(includes), + header_only=header_only, + ) + + # Generated once and threaded through: the cache key is taken over this + # exact text, and this exact text is what gets compiled. + build_file = ninja.generate(spec) + + build_key = cache.compute_build_key(spec, build_file=build_file) + scope = cache.build_key_dir(module_name=spec.module_name, build_key=build_key) + + prebuilt = cache.find_prebuilt(scope=scope, module_name=spec.module_name) + if prebuilt is not None: + try: + return _load(prebuilt) + except Exception as e: + # Also the benign case where a concurrent GC unlinked the leaf + # between the lookup and the load. + logger.warning( + "Cached JIT module %s failed to load; rebuilding. " "Got error: %s", + spec.module_name, + e, + ) + + 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 + # tensor-parallel ranks starting together into one compile plus N-1 + # cache hits instead of N identical compiles. + prebuilt = cache.find_prebuilt(scope=scope, module_name=spec.module_name) + if prebuilt is not None: + try: + return _load(prebuilt) + except Exception as e: + logger.warning( + "Cached JIT module %s failed to load; rebuilding. Got error: %s", + spec.module_name, + e, + ) + # A leaf is named after the dependency list it was built from, + # so the rebuild below lands on this exact name and publishing + # it would find the directory already there and keep the broken + # one -- leaving every later process to fail twice and rebuild + # for nothing. We hold the lock, so drop it now. + shutil.rmtree(prebuilt.parent, ignore_errors=True) + + # Build into a private staging directory, then publish it by renaming. + # Building in place would let another process observe a leaf that exists + # but is still being linked — and that process is on the fast path, so it + # takes no lock. A rename makes the leaf appear complete or not at all. + # + # The staging name is random rather than pid-based: the cache root is + # meant to be a shared mount, where two machines can hold the same pid. + staging = scope / f".staging-{uuid.uuid4().hex}" + try: + library = ninja.build(spec=spec, build_dir=staging, build_file=build_file) + # Loaded before publishing, so a broken artifact never becomes a + # leaf that later runs have to discover and discard. + module = _load(library) + cache.commit_build( + spec, + scope=scope, + staging=staging, + dependencies=ninja.scan_dependencies(staging), + ) + return module + finally: + shutil.rmtree(staging, ignore_errors=True) + + +@contextlib.contextmanager +def _build_lock(scope: pathlib.Path): + """Serialize builds of one module variant across processes. + + Every tensor-parallel rank on a node reaches the same cold cache at the same + moment. Without this they each run a full compile: measured at 8 ranks, + 14.0 s of CPU-oversubscribed nvcc against 9.0 s when one builds and the rest + wait and then hit the cache. + + The lock only saves duplicated work; it is not what makes publication safe + (the atomic rename is), and readers on the fast path deliberately do not + take it — that would serialize every warm load. Deadlock is not reachable: + one lock, never nested, and the kernel drops it if the holder dies. + """ + scope.mkdir(parents=True, exist_ok=True) + handle = os.open(scope / _LOCK_FILE, os.O_CREAT | os.O_RDWR, 0o644) + try: + fcntl.flock(handle, fcntl.LOCK_EX) + yield + finally: + os.close(handle) # releases the lock + + +def _load(library: pathlib.Path) -> Module: + from tvm_ffi import load_module + + return load_module(str(library)) diff --git a/python/sglang/kernels/jit/utils/compile/ninja.py b/python/sglang/kernels/jit/utils/compile/ninja.py new file mode 100644 index 000000000..4045c7831 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/ninja.py @@ -0,0 +1,228 @@ +"""Generating and running the ``build.ninja`` for one JIT module. + +Owning this file is what makes the cache key exact. Every compiler, flag, +include path and link argument is written here, so the key can be taken over the +generated text itself instead of over an approximation of what a dependency +would have chosen. + +Two things fall out of that which were previously impossible: + +* ``deps = gcc`` is deliberately *not* emitted. That setting folds the depfile + into ninja's binary ``.ninja_deps`` log and deletes it; plain ``depfile =`` + keeps the ``.d`` on disk, which is what the cache reads to learn the real + dependency closure. The dep-log optimization it gives up only matters for + builds with far more translation units than a JIT module has. +* The device rule always writes a depfile. tvm-ffi's HIP branch declares + ``depfile = $out.d`` while running a command that never produces one, so ROCm + builds silently carried no header dependencies at all. +""" + +from __future__ import annotations + +import logging +import os +import pathlib +import shlex +import subprocess +from typing import List + +from sglang.kernels.jit.utils.compile import toolchain +from sglang.kernels.jit.utils.compile.spec import BuildSpec + +logger = logging.getLogger(__name__) + +_BUILD_FILE = "build.ninja" +_NINJA_TIMEOUT_S = 1800 + + +def _escape(path: str) -> str: + """Escape a path for a ninja *path* field (a build statement's in/out).""" + return path.replace("$", "$$").replace(":", "$:").replace(" ", "$ ") + + +def _arg(path: str) -> str: + """Render a path as one shell word inside a rule command. + + Ninja's own escaping only survives ninja's parser: `$ ` reaches the command + line as a plain space, and every command runs through a shell, so an + unquoted include or library directory with a space in it arrives at the + compiler as several arguments. Quote for the shell first, then escape what + ninja still reads -- `$` is special everywhere in a build file. + """ + return shlex.quote(path).replace("$", "$$") + + +def _quote_path_flags(flags: List[str]) -> List[str]: + """Shell-quote the directory carried by every ``-I``/``-L`` flag. + + Applied once, at the end, wherever the flag came from -- this layer, the + toolchain, or the caller -- so a directory with a space in it stays one + argument. Anything else is passed through untouched. + """ + quoted: List[str] = [] + for flag in flags: + for prefix in ("-I", "-L"): + if flag.startswith(prefix) and len(flag) > len(prefix): + quoted.append(prefix + _arg(flag[len(prefix) :])) + break + else: + quoted.append(flag) + return quoted + + +def generate(spec: BuildSpec) -> str: + """The complete build description for *spec*, as ninja syntax. + + Paths of the generated translation units are relative, so the file is + identical no matter which directory the build runs in. + """ + units = spec.translation_units() + with_device = any(unit.is_cuda for unit in units) + + host_cc, device_cc = toolchain.compilers() + includes = toolchain.base_include_paths() + list(spec.include_paths) + include_flags = [f"-I{path}" for path in includes] + + cxxflags = _quote_path_flags( + toolchain.base_cxx_flags() + list(spec.cflags) + include_flags + ) + cudaflags = _quote_path_flags( + toolchain.base_cuda_flags() + + toolchain.target_flags() + + list(spec.cuda_cflags) + + include_flags + ) + ldflags = _quote_path_flags( + toolchain.base_link_flags(with_device=with_device) + list(spec.ldflags) + ) + + lines = [ + "ninja_required_version = 1.3", + f"cxx = {_arg(host_cc)}", + f"nvcc = {_arg(device_cc)}", + f"cxxflags = {' '.join(cxxflags)}", + f"cudaflags = {' '.join(cudaflags)}", + f"ldflags = {' '.join(ldflags)}", + "", + "rule compile_cxx", + " depfile = $out.d", + ' command = $cxx -MD -MF "$out.d" $cxxflags -c "$in" -o "$out"', + "", + "rule compile_cuda", + " depfile = $out.d", + ' command = $nvcc -MD -MF "$out.d" $cudaflags -c "$in" -o "$out"', + "", + "rule link", + ' command = $cxx $in $ldflags -o "$out"', + "", + ] + + objects: List[str] = [] + for index, unit in enumerate(units): + obj = f"{unit.stem}_{index}.o" + rule = "compile_cuda" if unit.is_cuda else "compile_cxx" + lines.append(f"build {obj}: {rule} {_escape(unit.filename)}") + objects.append(obj) + + lines += [ + "", + f"build {spec.module_name}.so: link {' '.join(objects)}", + "", + f"default {spec.module_name}.so", + "", + ] + return "\n".join(lines) + + +def build(*, spec: BuildSpec, build_dir: pathlib.Path, build_file: str) -> pathlib.Path: + """Write the sources and *build_file* into *build_dir*, then run ninja. + + The caller passes the generated text rather than letting this regenerate it, + so the text the cache key was taken over is provably the text that gets + compiled. + + *build_dir* is always a fresh staging directory, so there is nothing here to + keep current — every file is written once and compiled once. + """ + build_dir.mkdir(parents=True, exist_ok=True) + for unit in spec.translation_units(): + # Only the generated wrappers are materialized; sources that already + # exist are compiled where they are. + if unit.source is not None: + (build_dir / unit.filename).write_text(unit.source) + (build_dir / _BUILD_FILE).write_text(build_file) + + command = ["ninja", "-f", _BUILD_FILE] + jobs = os.environ.get("MAX_JOBS") + if jobs: + command += ["-j", jobs] + completed = subprocess.run( + command, + cwd=str(build_dir), + capture_output=True, + text=True, + timeout=_NINJA_TIMEOUT_S, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError( + f"Failed to build JIT module {spec.module_name} in {build_dir}\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + return build_dir / f"{spec.module_name}.so" + + +def scan_dependencies(build_dir: pathlib.Path) -> List[pathlib.Path]: + """Every file the compiler read, taken from the depfiles the build left. + + Without ``deps = gcc`` ninja leaves each ``.d`` in place, so this is a plain + read of what the compiler itself reported — no preprocessing is re-run and + no include paths are re-guessed. + """ + paths: List[pathlib.Path] = [] + seen = set() + for depfile in sorted(build_dir.glob("*.o.d")): + try: + text = depfile.read_text(errors="ignore") + except OSError: + continue + for candidate in _parse_depfile(text): + if candidate in seen: + continue + seen.add(candidate) + paths.append(pathlib.Path(candidate)) + return paths + + +def _parse_depfile(text: str) -> List[str]: + """Prerequisites from a make-style ``.d`` file. + + Handles the escapes both producers emit: ``\\`` line continuations, and + ``\\ `` for spaces inside a path. + """ + joined = text.replace("\\\r\n", " ").replace("\\\n", " ") + tokens: List[str] = [] + for line in joined.split("\n"): + _, separator, prerequisites = line.partition(":") + if not separator: + continue + current: List[str] = [] + index = 0 + while index < len(prerequisites): + char = prerequisites[index] + if char == "\\" and index + 1 < len(prerequisites): + if prerequisites[index + 1] in " \\": + current.append(prerequisites[index + 1]) + index += 2 + continue + if char.isspace(): + if current: + tokens.append("".join(current)) + current = [] + index += 1 + continue + current.append(char) + index += 1 + if current: + tokens.append("".join(current)) + return tokens diff --git a/python/sglang/kernels/jit/utils/compile/paths.py b/python/sglang/kernels/jit/utils/compile/paths.py new file mode 100644 index 000000000..6e4caf221 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/paths.py @@ -0,0 +1,32 @@ +"""Where the in-tree JIT sources live, and the compile defaults applied to them. + +Kept in its own module so that both the build cache and the loader can depend +on it without depending on each other. +""" + +from __future__ import annotations + +import importlib.util +import pathlib +from typing import List + +from sglang.kernels.jit.utils.common import cache_once + + +@cache_once +def _resolve_kernel_path() -> pathlib.Path: + spec = importlib.util.find_spec("sglang.kernels.jit") + assert spec is not None and spec.origin is not None + cur_dir = pathlib.Path(spec.origin).parent.resolve() + + candidate = cur_dir.resolve() + if (candidate / "include").exists() and (candidate / "csrc").exists(): + return candidate + + raise RuntimeError("Cannot find sglang.kernels.jit path") + + +KERNEL_PATH = _resolve_kernel_path() +DEFAULT_INCLUDE: List[str] = [str(KERNEL_PATH / "include")] +DEFAULT_CFLAGS: List[str] = ["-std=c++20", "-O3"] +DEFAULT_LDFLAGS: List[str] = [] diff --git a/python/sglang/kernels/jit/utils/compile/spec.py b/python/sglang/kernels/jit/utils/compile/spec.py new file mode 100644 index 000000000..c435dea90 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/spec.py @@ -0,0 +1,134 @@ +"""One fully-resolved JIT build, described in one place. + +``BuildSpec`` is the hand-off between the halves of ``load_jit``: the cache has +to see *every* input that could change the output in order to key it, and the +ninja generator has to feed those same inputs to the compiler. + +Anything added here that affects the generated code must also reach +``cache.compute_build_key``. +""" + +from __future__ import annotations + +import pathlib +from typing import List, Optional, Tuple + +import msgspec + +from sglang.kernels.jit.utils.compile.paths import KERNEL_PATH + +_MODULE_NAME_PREFIX = "sgl_kernel_jit_" + + +class TranslationUnit(msgspec.Struct, frozen=True): + """One file handed to the compiler. + + ``source`` is set for the wrapper units sglang generates and left None for + sources that already exist on disk; the ninja layer writes the former into + the build directory and points at the latter where they are. + """ + + filename: str + is_cuda: bool + source: Optional[str] = None + + @property + def stem(self) -> str: + return pathlib.Path(self.filename).stem + + +class BuildSpec(msgspec.Struct, frozen=True): + """Everything needed to either key or run one build.""" + + module_args: Tuple[str, ...] + cpp_files: Tuple[str, ...] + cuda_files: Tuple[str, ...] + cpp_wrappers: Tuple[Tuple[str, str], ...] + cuda_wrappers: Tuple[Tuple[str, str], ...] + cflags: Tuple[str, ...] + cuda_cflags: Tuple[str, ...] + ldflags: Tuple[str, ...] + include_paths: Tuple[str, ...] + header_only: bool + + @property + def module_name(self) -> str: + """Derived: the args are the module's identity, the name just spells it.""" + return _MODULE_NAME_PREFIX + "_".join(self.module_args) + + @property + def sources(self) -> Tuple[str, ...]: + return self.cpp_files + self.cuda_files + + @property + def wrappers(self) -> Tuple[Tuple[str, str], ...]: + return self.cpp_wrappers + self.cuda_wrappers + + def translation_units(self) -> List[TranslationUnit]: + """What the compiler is actually invoked on. + + Header-only modules are compiled through a generated wrapper that + includes the sources and exports the requested entry points; everything + else is compiled in place and exports from the C++ side itself. + """ + if not self.header_only: + return [ + TranslationUnit(filename=path, is_cuda=path.endswith(".cu")) + for path in self.sources + ] + + units: List[TranslationUnit] = [] + for name, is_cuda, files, wrappers in ( + ("main.cpp", False, self.cpp_files, self.cpp_wrappers), + ("cuda.cu", True, self.cuda_files, self.cuda_wrappers), + ): + if not files and not wrappers: + continue + units.append( + TranslationUnit( + filename=name, + is_cuda=is_cuda, + source=_wrapper_source(files, wrappers), + ) + ) + return units + + +# What tvm-ffi's own `_decorate_with_tvm_ffi` prepends to every generated unit. +# The wrapper below uses TVM_FFI_DLL_EXPORT_TYPED_FUNC, so it has to include the +# header that defines it rather than rely on the kernel's own include chain +# happening to drag it in — that dependency held for every kernel in tree, but +# it is not something a new kernel's author would know to preserve. +_FFI_INCLUDES = ( + "#include ", + "#include ", + "#include ", + "#include ", + "#include ", +) + + +def _wrapper_source( + files: Tuple[str, ...], wrappers: Tuple[Tuple[str, str], ...] +) -> str: + lines = list(_FFI_INCLUDES) + lines += [f'#include "{path}"' for path in files] + lines.append("namespace sglang {") + lines += [ + f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" + for export_name, kernel_name in wrappers + ] + lines.append("} // namespace sglang") + return "\n".join(lines) + "\n" + + +def resolve_sources(files: List[str] | None) -> Tuple[str, ...]: + """Absolute paths pass through; anything else is relative to ``csrc/``.""" + return tuple( + str( + path.resolve() + if path.is_absolute() + else (KERNEL_PATH / "csrc" / path).resolve() + ) + for path in map(pathlib.Path, files or []) + ) diff --git a/python/sglang/kernels/jit/utils/compile/toolchain.py b/python/sglang/kernels/jit/utils/compile/toolchain.py new file mode 100644 index 000000000..4da0a53a7 --- /dev/null +++ b/python/sglang/kernels/jit/utils/compile/toolchain.py @@ -0,0 +1,168 @@ +"""The toolchain a JIT build runs on: compilers, tvm-ffi's headers, base flags. + +sglang generates its own ``build.ninja`` rather than going through +``tvm_ffi.cpp.load_inline``, so the flags tvm-ffi used to supply implicitly have +to be stated here. That is the point: every flag that reaches the compiler is +now visible in one place and therefore hashable into the build key, instead of +living inside a dependency whose defaults we could only approximate by version +number. + +Only tvm-ffi's *locations* are still consumed — its headers, its shared library, +and ``tvm_ffi.load_module`` for loading the result. +""" + +from __future__ import annotations + +import logging +import os +import pathlib +import shutil +from typing import List, Tuple + +import torch + +from sglang.kernels.jit.utils.arch import get_jit_cuda_arch +from sglang.kernels.jit.utils.common import cache_once, is_hip_runtime + +logger = logging.getLogger(__name__) + + +@cache_once +def cuda_home() -> str: + """CUDA install root, resolved the way tvm-ffi resolves it. + + `arch._jit_cuda_version` resolves nvcc the same way for its own purposes; + the two must stay in agreement, since one picks the target and the other + compiles for it. + """ + configured = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if configured is not None: + return configured + nvcc_path = shutil.which("nvcc") + if nvcc_path is not None: + return os.path.dirname(os.path.dirname(nvcc_path)) + return "/usr/local/cuda" + + +@cache_once +def rocm_home() -> str: + """ROCm install root, resolved the way tvm-ffi resolves it.""" + return os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") or "/opt/rocm" + + +@cache_once +def device_compiler_path() -> str: + """The nvcc/hipcc that JIT builds actually invoke. + + Resolved the same way tvm-ffi resolves it, so the binary the cache + fingerprints is the binary that does the compiling. + """ + if is_hip_runtime(): + return os.path.join(rocm_home(), "bin", "hipcc") + return os.path.join(cuda_home(), "bin", "nvcc") + + +@cache_once +def host_compiler_path() -> str: + """The C++ compiler host code is handed to. + + nvcc dispatches all host code to it, so its version decides both which + system headers are pulled in and how that half is codegen'd. + """ + return os.environ.get("CXX", "c++") + + +@cache_once +def gpu_arch_name() -> str: + """The compile target as the vendor names it. + + On ROCm this is ``gcnArchName`` (``gfx942:sramecc+:xnack-``) rather than the + CUDA-shaped ``(major, minor)`` capability: the latter maps gfx940/gfx941/ + gfx942 onto a single ``9.4``, which are three different compile targets. + """ + if not is_hip_runtime(): + return get_jit_cuda_arch().target_name + try: + device = torch.cuda.current_device() + return str(torch.cuda.get_device_properties(device).gcnArchName) + except Exception: + logger.warning("Cannot detect ROCm gcnArchName; the JIT cache target degrades.") + return "unknown" + + +@cache_once +def toolkit_home() -> pathlib.Path: + """The CUDA/ROCm root, derived from the compiler already resolved.""" + return pathlib.Path(device_compiler_path()).parent.parent + + +@cache_once +def tvm_ffi_paths() -> Tuple[Tuple[str, ...], str, str]: + """``(include dirs, library dir, library name)`` for linking against tvm-ffi.""" + from tvm_ffi.libinfo import ( + find_dlpack_include_path, + find_include_path, + find_libtvm_ffi, + ) + + lib = pathlib.Path(find_libtvm_ffi()) + includes = tuple(dict.fromkeys([find_include_path(), find_dlpack_include_path()])) + return includes, str(lib.parent), lib.stem.removeprefix("lib") + + +def target_flags() -> List[str]: + """The device flags that pin the build to this GPU. + + Emitted from the architecture sglang already detected, rather than left to + the compiler driver to probe: the value is part of the cache key, so it has + to be decided here and not rediscovered at build time. + """ + if is_hip_runtime(): + return [f"--offload-arch={gpu_arch_name()}"] + arch = get_jit_cuda_arch() + target = f"{arch.major}{arch.minor}{arch.suffix}" + return [f"-gencode=arch=compute_{target},code=sm_{target}"] + + +def base_cxx_flags() -> List[str]: + """Only what the platform requires; `-std`/`-O` arrive with the spec. + + Kept disjoint from ``arch.get_default_target_flags`` on purpose — supplying + `-std=c++20` from both is what used to make nvcc warn about an incompatible + redefinition on every single build. + """ + return ["-fPIC"] + + +def base_cuda_flags() -> List[str]: + if is_hip_runtime(): + return ["-fPIC", "-D__HIP_PLATFORM_AMD__=1", "-fno-gpu-rdc"] + return ["-Xcompiler", "-fPIC"] + + +def base_include_paths() -> List[str]: + includes, _, _ = tvm_ffi_paths() + if is_hip_runtime(): + return [*includes, f"{rocm_home()}/include"] + return list(includes) + + +def base_link_flags(*, with_device: bool) -> List[str]: + """Link flags for a module, with the GPU runtime only when it has device code. + + A module built purely from ``.cpp`` sources must not drag in libcudart: CPU + runners have no CUDA toolkit to link it from, and the module never calls it. + tvm-ffi keyed this off the presence of ``.cu`` sources for the same reason. + """ + _, lib_dir, lib_name = tvm_ffi_paths() + flags = ["-shared", f"-L{lib_dir}", f"-l{lib_name}"] + if not with_device: + return flags + if is_hip_runtime(): + return flags + [f"-L{rocm_home()}/lib", "-lamdhip64"] + return flags + [f"-L{cuda_home()}/lib64", "-lcudart"] + + +def compilers() -> Tuple[str, str]: + """``(host compiler, device compiler)`` as they will appear in build.ninja.""" + return host_compiler_path(), device_compiler_path() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 64e7826c4..8392e8533 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -972,6 +972,15 @@ class Envs: # Cache directories # =================================================================== SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang")) + # JIT kernel build cache. None = unset, resolving to ~/.cache/sglang/jit; + # point it at a persistent mount to share builds across CI jobs. + SGLANG_JIT_CACHE_DIR = EnvStr(None) + # Log, at INFO, which dependency changed whenever a module is rebuilt. + SGLANG_JIT_CACHE_DEBUG = EnvBool(False) + # How many builds to keep per module variant. None = unset = keep all, which + # 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) # =================================================================== # Expert-parallel dispatch and MoE execution diff --git a/test/registered/kernels/test_jit_cache.py b/test/registered/kernels/test_jit_cache.py new file mode 100644 index 000000000..7fb2e7fe0 --- /dev/null +++ b/test/registered/kernels/test_jit_cache.py @@ -0,0 +1,532 @@ +"""CPU-only tests for the JIT build pipeline: ninja generation and the cache. + +Everything here runs against synthetic files, so the invariants that matter — a +bad recorded dependency list never causes reuse, differing flags never share a +directory, a moved clone still hits — are checked without a GPU or a compiler. +""" + +from __future__ import annotations + +import os +import pathlib +import sys + +import msgspec +import pytest + +from sglang.kernels.jit.utils.compile import cache, ninja +from sglang.kernels.jit.utils.compile.paths import KERNEL_PATH +from sglang.kernels.jit.utils.compile.spec import BuildSpec +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +@pytest.fixture(autouse=True) +def _fresh_digests(): + cache.clear_digest_cache() + yield + cache.clear_digest_cache() + + +def _write(path: pathlib.Path, text: str) -> pathlib.Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + return path + + +def _spec(**overrides) -> BuildSpec: + base = dict( + module_args=("activation", "bf16_t"), + cpp_files=(), + cuda_files=(), + cpp_wrappers=(("run", "Kernel::run"),), + cuda_wrappers=(), + cflags=("-O3",), + cuda_cflags=("-O3",), + ldflags=(), + include_paths=(), + header_only=True, + ) + base.update(overrides) + return BuildSpec(**base) + + +def _build_key(**overrides) -> str: + spec = _spec(**overrides) + return cache.compute_build_key(spec, build_file=ninja.generate(spec)) + + +def _entries(paths) -> list: + out = [] + for path in paths: + root, relpath = cache._normalize_path(path) + out.append( + cache._DepEntry(root=root, relpath=relpath, digest=cache._file_digest(path)) + ) + return out + + +def _publish_leaf(scope: pathlib.Path, paths, *, module_name="m") -> pathlib.Path: + """Create a leaf the way commit_build would: name derived from its own list.""" + entries = _entries(paths) + leaf = scope / f"{cache._DEPS_KEY_PREFIX}{cache._deps_key(entries)}" + leaf.mkdir(parents=True) + (leaf / cache._DEPS_FILE).write_bytes(msgspec.json.encode(entries)) + (leaf / f"{module_name}.so").write_bytes(b"") + return leaf + + +# -------------------------------------------------------------------------- +# Anchor roots +# -------------------------------------------------------------------------- + + +def test_in_tree_paths_normalize_to_an_anchor(): + header = KERNEL_PATH / "include" / "sgl_kernel" / "utils.cuh" + root, relpath = cache._normalize_path(header) + assert root == "kernels" + assert relpath == "include/sgl_kernel/utils.cuh" + assert cache._resolve_path(root=root, relpath=relpath) == header + + +def test_unknown_paths_fall_back_to_absolute(tmp_path): + root, relpath = cache._normalize_path(tmp_path / "elsewhere.h") + assert root == "abs" + assert cache._resolve_path(root=root, relpath=relpath) == tmp_path / "elsewhere.h" + + +def test_unresolvable_anchor_is_a_miss_not_a_crash(): + assert cache._resolve_path(root="pkg:does-not-exist", relpath="x.h") is None + + +def test_anchor_roots_are_symlink_resolved(tmp_path, monkeypatch): + """Anchors must be symlink-resolved, since the paths matched against them are. + + `/usr/local/cuda` is a symlink to `/usr/local/cuda-`; an unresolved + anchor makes every toolkit header miss it and fall through to `sys`, whose + relpath then carries the CUDA version and breaks reuse across upgrades. + + The toolkit stands in for all of them — one comprehension resolves every + anchor — and it is faked rather than read off the machine so this still + guards on the CPU-only runners, which have no toolkit at all. + """ + from sglang.kernels.jit.utils.compile import toolchain + + versioned = tmp_path / "cuda-12.9" + versioned.mkdir() + (tmp_path / "cuda").symlink_to(versioned) + monkeypatch.setattr(toolchain, "toolkit_home", lambda: tmp_path / "cuda") + + # `__wrapped__` is the undecorated function: the anchors are memoized for + # the process, and the real ones were already computed by an earlier test. + roots = dict(cache._anchor_roots.__wrapped__()) + assert roots["toolkit"] == versioned + + +# -------------------------------------------------------------------------- +# build_key — what must and must not change it +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "overrides", + [ + {"module_args": ("activation", "fp16_t")}, + {"cuda_cflags": ("-O3", "--use_fast_math")}, + {"cflags": ("-O2",)}, + {"ldflags": ("-lfoo",)}, + {"cpp_wrappers": (("run", "Other::run"),)}, + {"include_paths": ("/opt/extra",)}, + {"header_only": False, "cpp_files": ("/tmp/x.cpp",), "cpp_wrappers": ()}, + ], + ids=["args", "cuda_cflags", "cflags", "ldflags", "wrappers", "includes", "mode"], +) +def test_build_key_separates_every_build_input(overrides): + """Two builds differing in any of these must never share a directory. + + Sharing one would let a lookup select a leaf produced under different flags, + which is the failure the whole two-key split exists to prevent. + """ + assert _build_key(**overrides) != _build_key() + + +def test_build_key_covers_the_whole_ninja_file(): + """Every flag reaching the compiler reaches the key, because the key is + taken over the generated build file itself rather than over a hand-listed + subset of inputs.""" + spec = _spec() + baseline = cache.compute_build_key(spec, build_file=ninja.generate(spec)) + tampered = ninja.generate(spec).replace("-O3", "-O0") + assert cache.compute_build_key(spec, build_file=tampered) != baseline + + +def test_build_key_tracks_direct_source_contents(tmp_path): + source = _write(tmp_path / "a.cu", "// v1") + before = _build_key(cuda_files=(str(source),)) + source.write_text("// v2") + cache.clear_digest_cache() + assert _build_key(cuda_files=(str(source),)) != before + + +def test_no_unordered_container_reaches_the_key(monkeypatch): + """Nothing hashed into a key may iterate in `PYTHONHASHSEED` order. + + A set or dict among the hashed parts would make the same tree key + differently in two processes: no error, no wrong result, the cache simply + never hits again. Verified end-to-end by running the key computation under + several hash seeds; this pins the property cheaply. + """ + recorded = [] + original = cache._hash_parts + monkeypatch.setattr( + cache, + "_hash_parts", + lambda parts: recorded.append(list(parts)) or original(recorded[-1]), + ) + _build_key() + + def walk(value, path="parts"): + assert not isinstance( + value, (set, frozenset, dict) + ), f"unordered container at {path}: {type(value).__name__}" + if isinstance(value, (list, tuple)): + for index, item in enumerate(value): + walk(item, f"{path}[{index}]") + + walk(recorded) + + +def test_build_key_is_independent_of_install_location(): + """Absolute paths are anchor-normalized, which is what lets a second clone + of the same tree reuse the first clone's builds.""" + text = ninja.generate(_spec()) + assert str(KERNEL_PATH) not in cache._normalize_text(text) + + +# -------------------------------------------------------------------------- +# deps_key — a leaf that reproduces its own name +# -------------------------------------------------------------------------- + + +def test_leaf_is_found_when_nothing_changed(tmp_path): + dep = _write(tmp_path / "dep.h", "// v1") + leaf = _publish_leaf(tmp_path, [dep]) + assert cache.find_prebuilt(scope=tmp_path, module_name="m") == leaf / "m.so" + + +def test_leaf_is_skipped_when_a_dependency_changed(tmp_path): + dep = _write(tmp_path / "dep.h", "// v1") + _publish_leaf(tmp_path, [dep]) + dep.write_text("// v2") + cache.clear_digest_cache() + assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None + + +def test_leaf_survives_an_mtime_only_change(tmp_path): + dep = _write(tmp_path / "dep.h", "// stable") + leaf = _publish_leaf(tmp_path, [dep]) + os.utime(dep, (0, 0)) + cache.clear_digest_cache() + assert cache.find_prebuilt(scope=tmp_path, module_name="m") == leaf / "m.so" + + +def test_leaf_is_skipped_when_a_dependency_disappeared(tmp_path): + dep = _write(tmp_path / "dep.h", "// here") + _publish_leaf(tmp_path, [dep]) + dep.unlink() + cache.clear_digest_cache() + assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None + + +def test_a_leaf_that_does_not_match_its_own_name_is_skipped(tmp_path): + """The recorded list is verified, not trusted. + + This is what replaces a shared manifest plus a format-version check: a + truncated, tampered, or foreign list simply fails to reproduce the leaf's + own name, so no schema field has to be believed. + """ + dep = _write(tmp_path / "dep.h", "// v1") + leaf = _publish_leaf(tmp_path, [dep]) + other = _write(tmp_path / "other.h", "// x") + (leaf / cache._DEPS_FILE).write_bytes(msgspec.json.encode(_entries([dep, other]))) + assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None + + +@pytest.mark.parametrize( + "payload", + [b"", b"not json", b'[["kernels", "a.h"]]'], + ids=["empty", "garbage", "wrong-shape"], +) +def test_an_unreadable_dependency_list_is_a_miss(tmp_path, payload): + dep = _write(tmp_path / "dep.h", "// v1") + leaf = _publish_leaf(tmp_path, [dep]) + (leaf / cache._DEPS_FILE).write_bytes(payload) + assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None + + +def test_a_foreign_leaf_does_not_block_a_valid_one(tmp_path): + """A leaf naming a file that does not exist here is skipped, not fatal. + + A shared, merged manifest could not do this: one unresolvable entry would + make every lookup fail permanently on this machine. + """ + dep = _write(tmp_path / "dep.h", "// v1") + good = _publish_leaf(tmp_path, [dep]) + + foreign = tmp_path / f"{cache._DEPS_KEY_PREFIX}{'0' * 16}" + foreign.mkdir() + (foreign / cache._DEPS_FILE).write_bytes( + msgspec.json.encode( + [cache._DepEntry(root="sys", relpath="include/c++/99/absent.h", digest="x")] + ) + ) + (foreign / "m.so").write_bytes(b"") + os.utime(foreign, None) # make the foreign leaf the newest + + assert cache.find_prebuilt(scope=tmp_path, module_name="m") == good / "m.so" + + +def test_a_hit_survives_an_unwritable_cache(tmp_path, monkeypatch): + """Touching the leaf is bookkeeping; it must never turn a hit into a crash. + + The cache root can be a read-only mount, and a prune racing the lookup + leaves nothing to touch -- either way `os.utime` raises, and before this it + escaped `find_prebuilt` and took `load_jit` down on an otherwise good hit. + """ + header = _write(tmp_path / "a.h", "x") + scope = tmp_path / "scope" + _publish_leaf(scope, [header], module_name="m") + + def deny(*args, **kwargs): + raise PermissionError("read-only file system") + + monkeypatch.setattr(cache.os, "utime", deny) + assert cache.find_prebuilt(scope=scope, module_name="m") is not None + + +def test_missing_library_is_not_a_hit(tmp_path): + dep = _write(tmp_path / "dep.h", "// v1") + leaf = _publish_leaf(tmp_path, [dep]) + (leaf / "m.so").unlink() + assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None + + +# -------------------------------------------------------------------------- +# Commit-side guards +# -------------------------------------------------------------------------- + + +def test_empty_dependency_scan_is_rejected(): + """An empty scan must not be recorded: it would narrow the checked set to + nothing, the one shape of bad recorded data that could cause reuse.""" + assert not cache._covers_direct_sources(entries=[], direct_sources=["/x/a.cuh"]) + + +def test_scan_missing_a_direct_source_is_rejected(tmp_path): + other = _write(tmp_path / "other.h", "// x") + entries = _entries([other]) + assert not cache._covers_direct_sources( + entries=entries, direct_sources=[str(tmp_path / "a.cuh")] + ) + assert cache._covers_direct_sources(entries=entries, direct_sources=[str(other)]) + + +def test_build_directory_entries_are_dropped(tmp_path): + """The generated translation unit is not a dependency of itself. + + Its path is unstable and its contents are already a function of inputs the + build key covers, so recording it would defeat reuse across clones. + """ + build_dir = (tmp_path / "build").resolve() + generated = _write(build_dir / "cuda.cu", "// generated") + outside = _write(tmp_path / "real.h", "// real") + entries = cache._to_entries(dependencies=[generated, outside], build_dir=build_dir) + assert [entry.relpath for entry in entries] == [str(outside)] + + +def test_publish_loses_the_race_gracefully(tmp_path): + """Two processes building identical content: the loser adopts the winner's leaf.""" + winner = tmp_path / "leaf" + winner.mkdir() + (winner / "m.so").write_bytes(b"winner") + staging = tmp_path / "staging" + staging.mkdir() + (staging / "m.so").write_bytes(b"loser") + + assert cache._publish(staging=staging, leaf=winner) == winner + assert (winner / "m.so").read_bytes() == b"winner" + + +def test_build_lock_excludes_a_second_holder(tmp_path): + """One compile per module variant per node, not one per process. + + Every tensor-parallel rank hits the same cold cache at the same instant. + Without exclusion each runs a full compile — measured with 8 ranks, all + eight compiled; with it, one compiled and seven took the cache. The lock + saves duplicated work only; the atomic rename is what makes publication + safe, so a missing lock is slow rather than wrong. + """ + import threading + + from sglang.kernels.jit.utils.compile import loader + + held, release, contender_entered = (threading.Event() for _ in range(3)) + + def holder(): + with loader._build_lock(tmp_path): + held.set() + release.wait(5) + + def contender(): + with loader._build_lock(tmp_path): + contender_entered.set() + + first = threading.Thread(target=holder) + first.start() + assert held.wait(5) + + second = threading.Thread(target=contender) + second.start() + assert not contender_entered.wait(0.3), "entered while the lock was held" + + release.set() + first.join(5) + second.join(5) + assert contender_entered.is_set(), "never entered after the lock was released" + + +# -------------------------------------------------------------------------- +# ninja generation +# -------------------------------------------------------------------------- + + +def test_generated_ninja_keeps_depfiles_on_disk(): + """`deps = gcc` must not be emitted. + + That setting folds each depfile into ninja's binary log and deletes it, + leaving the cache with nothing to record the dependency closure from. + """ + text = ninja.generate(_spec(cuda_wrappers=(("run", "K::run"),))) + assert "deps = gcc" not in text + assert "depfile = $out.d" in text + + +def test_generated_ninja_asks_both_compilers_for_dependencies(): + """Both rules must write a depfile on every backend. + + tvm-ffi's HIP branch declared `depfile =` while running a command that never + produced one, so ROCm builds silently carried no header dependencies. + """ + text = ninja.generate(_spec(cuda_wrappers=(("run", "K::run"),))) + compile_commands = [ + line + for line in text.splitlines() + if line.startswith(" command = ") and ' -c "$in"' in line + ] + assert len(compile_commands) == 2 + assert all('-MD -MF "$out.d"' in line for line in compile_commands) + + +def test_pure_cpp_module_does_not_link_the_gpu_runtime(): + """A module with no `.cu` sources must not ask the linker for libcudart. + + `ngram_corpus` is five .cpp files and no device code, and it is built on + CPU-only CI runners that have no CUDA toolkit — linking it there fails with + `cannot find -lcudart`. tvm-ffi keyed the runtime flags off the presence of + `.cu` sources for exactly this reason. + """ + cpu_only = _spec(cpp_files=("/tmp/a.cpp",), cpp_wrappers=(), header_only=False) + ldflags = next( + line + for line in ninja.generate(cpu_only).splitlines() + if line.startswith("ldflags = ") + ) + assert "cudart" not in ldflags and "amdhip" not in ldflags + + with_device = _spec(cuda_files=("/tmp/a.cu",), cpp_wrappers=(), header_only=False) + ldflags = next( + line + for line in ninja.generate(with_device).splitlines() + if line.startswith("ldflags = ") + ) + assert "cudart" in ldflags or "amdhip" in ldflags + + +def test_generated_ninja_is_deterministic(): + spec = _spec(cuda_wrappers=(("run", "K::run"),)) + assert ninja.generate(spec) == ninja.generate(spec) + + +def test_header_only_module_compiles_through_a_generated_wrapper(tmp_path): + source = str(tmp_path / "kernel.cuh") + units = _spec( + cuda_files=(source,), cuda_wrappers=(("run", "K::run"),) + ).translation_units() + assert [unit.filename for unit in units] == ["main.cpp", "cuda.cu"] + generated = next(unit for unit in units if unit.filename == "cuda.cu").source + assert f'#include "{source}"' in generated + assert "TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, (K::run));" in generated + # The wrapper must include the header defining the macro it uses. Every + # kernel in tree happens to drag it in transitively, so dropping this would + # only break whichever future kernel does not. + assert "#include " in generated + + +def test_non_header_only_module_compiles_its_sources_in_place(tmp_path): + source = str(tmp_path / "kernel.cu") + units = _spec( + cuda_files=(source,), cpp_wrappers=(), header_only=False + ).translation_units() + assert [(unit.filename, unit.source, unit.is_cuda) for unit in units] == [ + (source, None, True) + ] + + +# -------------------------------------------------------------------------- +# depfile parsing +# -------------------------------------------------------------------------- + + +def test_depfile_parsing_handles_continuations_and_escaped_spaces(): + text = "cuda_0.o: /a/cuda.cu \\\n /a/with\\ space.cuh \\\n /b/plain.h\n" + assert ninja._parse_depfile(text) == [ + "/a/cuda.cu", + "/a/with space.cuh", + "/b/plain.h", + ] + + +def test_depfile_parsing_ignores_a_target_with_no_prerequisites(): + assert ninja._parse_depfile("a.o:\n") == [] + + +# -------------------------------------------------------------------------- +# Layout +# -------------------------------------------------------------------------- + + +def test_layout_is_readable_and_scoped_by_build_key(): + scope = cache.build_key_dir( + module_name="sgl_kernel_jit_activation_bf16_t", build_key="abc123" + ) + assert scope.name == "build-abc123" + assert scope.parent.name == "sgl_kernel_jit_activation_bf16_t" + assert scope.parent.parent.name == cache._target_tag() + + +def test_module_name_is_derived_from_the_args(): + assert _spec().module_name == "sgl_kernel_jit_activation_bf16_t" + + +def test_relative_sources_resolve_against_csrc(): + from sglang.kernels.jit.utils.compile.spec import resolve_sources + + assert resolve_sources(["elementwise/activation.cuh"]) == ( + str(KERNEL_PATH / "csrc" / "elementwise" / "activation.cuh"), + ) + assert resolve_sources(["/usr/include/stdio.h"]) == ("/usr/include/stdio.h",) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/test_kernel_inventory.py b/test/registered/kernels/test_kernel_inventory.py index e2f3dd6b2..99069e51f 100644 --- a/test/registered/kernels/test_kernel_inventory.py +++ b/test/registered/kernels/test_kernel_inventory.py @@ -153,10 +153,19 @@ def test_internal_registry_target_attributes_are_declared(): assert not missing, f"KernelSpec targets missing attributes: {missing}" +# `load_jit` takes in-tree names and absolute paths on the same keyword, so this +# check can only reach the declarations spelled out in the source. A module that +# assembles its file list at runtime from a package outside `jit/csrc` has no +# in-tree name to verify and belongs here; there is none at the moment. +_RUNTIME_JIT_SOURCE_MODULES: set[str] = set() + + def test_jit_source_declarations_exist(): missing = [] unsupported = [] for python_file in OPS_ROOT.rglob("*.py"): + if python_file.relative_to(OPS_ROOT).as_posix() in _RUNTIME_JIT_SOURCE_MODULES: + continue tree = ast.parse(python_file.read_text()) constants = _module_string_constants(tree) for call in (node for node in ast.walk(tree) if isinstance(node, ast.Call)):