[kernel] Content-addressed JIT build cache, generated from our own ninja (#34274)
Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
DarkSharpness
co-authored by
Claude
BBuf
Claude
BBuf
parent
240a12b302
commit
b784726863
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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/<target>/<module_name>/build-<build_key>/
|
||||
deps-<deps_key>/ one leaf per transitive-content state
|
||||
sgl_deps.json the dependency list of *this* build
|
||||
<module_name>.so
|
||||
.staging-<uuid>/ a build in progress
|
||||
|
||||
``<target>`` and ``<module_name>`` 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-<version>`; 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", <absolute path>)`` 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
|
||||
@@ -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)
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -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] = []
|
||||
@@ -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 <tvm/ffi/container/tensor.h>",
|
||||
"#include <tvm/ffi/dtype.h>",
|
||||
"#include <tvm/ffi/error.h>",
|
||||
"#include <tvm/ffi/extra/c_env_api.h>",
|
||||
"#include <tvm/ffi/function.h>",
|
||||
)
|
||||
|
||||
|
||||
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 [])
|
||||
)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user