# Copyright 2026 SGLang Team. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import importlib import os import platform import shutil import subprocess import sys import sysconfig from pathlib import Path root = Path(__file__).parent.resolve() _BUILD_REQUIRES = [ ("setuptools", "setuptools"), ("mlx", "mlx"), ("nanobind", "nanobind"), ] def _ensure_toolchain(): if sys.platform != "darwin" or platform.machine() != "arm64": raise SystemExit("setup_metal.py only supports macOS (Apple Silicon).") if shutil.which("c++") is None or shutil.which("xcrun") is None: raise SystemExit( "Apple toolchain not found. Install the Xcode Command Line Tools " "with `xcode-select --install` (or a full Xcode install) and retry." ) try: subprocess.check_output( ["xcrun", "-sdk", "macosx", "metal", "--version"], stderr=subprocess.STDOUT, ) except (subprocess.CalledProcessError, FileNotFoundError) as exc: raise SystemExit( "Apple Metal shader compiler not found. Install a full Xcode " "(not just Command Line Tools) so that `xcrun -sdk macosx metal` " "is available, then retry." ) from exc def _ensure_build_requires(): missing = [] for import_name, pip_name in _BUILD_REQUIRES: try: importlib.import_module(import_name) except ImportError: missing.append(pip_name) if not missing: return print( f"[sgl-kernel:metal] installing build requirements: {missing}", flush=True, ) subprocess.check_call( [sys.executable, "-m", "pip", "install", "--upgrade", *missing] ) # Section 1: Prerequisites _ensure_toolchain() _ensure_build_requires() os.chdir(root) # Section 2: Build and install from setuptools import Extension, find_packages, setup # noqa: E402 from setuptools.command.build_ext import build_ext # noqa: E402 def _get_version(): with open(root / "pyproject.toml") as f: for line in f: if line.startswith("version"): return line.split("=")[1].strip().strip('"') operator_namespace = "sgl_kernel" metallib_name = "sgl_metal_kernels.metallib" # Metal shader sources (compiled with `xcrun metal`) and C++ host sources # (compiled with `c++`). Add new kernels by appending to these lists. metal_shader_sources = [ "csrc/metal/rope_pool_fused.metal", ] cxx_sources = [ "csrc/metal/rope_pool_fused.cpp", ] # Header search paths shared by both the Metal shader compiler and the C++ # host compiler. include_dirs = [ root / "csrc", root / "csrc" / "metal", ] cxx_flags = ["-std=c++17", "-O3", "-fvisibility=hidden"] metal_flags = ["-O3"] frameworks = ["Metal", "Foundation", "QuartzCore"] libraries = ["mlx"] class BuildMetalExtension(build_ext): def build_extension(self, ext): if sys.platform != "darwin" or platform.machine() != "arm64": raise RuntimeError("setup_metal.py only supports macOS (Apple Silicon).") ext_path = Path(self.get_ext_fullpath(ext.name)) ext_path.parent.mkdir(parents=True, exist_ok=True) # Use ccache for the C++ compiles when it is on PATH. ccache does not # support `.metal` sources (unsupported source language), so the Metal # shader compile is left untouched. ccache = shutil.which("ccache") cxx_cmd = [ccache, "c++"] if ccache else ["c++"] if ccache: print(f"[sgl-kernel:metal] using ccache at {ccache}", flush=True) python_exe = Path(sys.executable) python_include = Path(sysconfig.get_paths()["include"]) python_lib = Path(sysconfig.get_config_var("LIBDIR")) # Match the deployment target that Python itself was built against # unless the user overrides it. MLX's prebuilt wheels may require a # higher minimum; in that case set MACOSX_DEPLOYMENT_TARGET explicitly. deployment_target = os.environ.get( "MACOSX_DEPLOYMENT_TARGET", str(sysconfig.get_config_var("MACOSX_DEPLOYMENT_TARGET") or "11.0"), ) def _python_eval(expr: str) -> str: return subprocess.check_output( [str(python_exe), "-c", expr], text=True ).strip() nanobind_dir = Path( _python_eval("import nanobind; print(nanobind.__path__[0])") ) mlx_dir = Path(_python_eval("import mlx.core as mx; print(mx.__file__)")) mlx_site = mlx_dir.parent mlx_include = mlx_site / "include" mlx_lib = mlx_site / "lib" generated_dir = root / "build" / "metal" generated_dir.mkdir(parents=True, exist_ok=True) metallib_path = generated_dir / metallib_name metal_std = os.environ.get("SGL_METAL_STD", "metal3.1") ext_include_dirs = [Path(p) for p in (ext.include_dirs or [])] host_includes = [ python_include, nanobind_dir / "include", nanobind_dir / "ext" / "robin_map" / "include", mlx_include, mlx_include / "metal_cpp", ] all_includes = ext_include_dirs + host_includes include_args = [f"-I{p}" for p in all_includes] # `xcrun metal` accepts `-I` for header search; reuse the project # include dirs so shaders can include shared MSL headers. metal_include_args = [f"-I{p}" for p in ext_include_dirs] if not metal_shader_sources: raise RuntimeError("metal_shader_sources is empty; nothing to compile") air_paths = [] for rel in metal_shader_sources: metal_src = root / rel if not metal_src.is_file(): raise RuntimeError(f"metal shader source not found: {metal_src}") air_path = generated_dir / (metal_src.stem + ".air") self.spawn( [ "xcrun", "-sdk", "macosx", "metal", f"-std={metal_std}", *metal_flags, *metal_include_args, "-c", str(metal_src), "-o", str(air_path), ] ) air_paths.append(str(air_path)) self.spawn( [ "xcrun", "-sdk", "macosx", "metallib", *air_paths, "-o", str(metallib_path), ] ) cflags = [ *cxx_flags, f"-mmacosx-version-min={deployment_target}", *include_args, ] ldflags = [ "-shared", "-undefined", "dynamic_lookup", f"-mmacosx-version-min={deployment_target}", f"-L{python_lib}", f"-L{mlx_lib}", f"-Wl,-rpath,{mlx_lib}", *[f"-l{lib}" for lib in libraries], *[arg for fw in frameworks for arg in ("-framework", fw)], ] objects = [] for src in ext.sources: src_path = Path(src) obj_path = generated_dir / (src_path.stem + ".o") compile_cmd = [ *cxx_cmd, *cflags, "-c", str(src_path), "-o", str(obj_path), ] self.spawn(compile_cmd) objects.append(str(obj_path)) nanobind_src = nanobind_dir / "src" / "nb_combined.cpp" nanobind_obj = generated_dir / "nb_combined.o" nanobind_cmd = [ *cxx_cmd, *cflags, "-DNB_COMPACT_ASSERTIONS", "-DNB_BUILD", "-DNB_SHARED", "-c", str(nanobind_src), "-o", str(nanobind_obj), ] self.spawn(nanobind_cmd) objects.append(str(nanobind_obj)) link_cmd = [ "c++", *objects, *ldflags, "-o", str(ext_path), ] self.spawn(link_cmd) # Stage the metallib next to the freshly-linked extension so that # `install_lib` picks it up via `package_data={"sgl_kernel": ["*.metallib"]}`. staged_metallib = ext_path.parent / metallib_path.name if metallib_path.resolve() != staged_metallib.resolve(): shutil.copy2(metallib_path, staged_metallib) ext_modules = [ Extension( name=f"{operator_namespace}._metal", sources=cxx_sources, include_dirs=[str(p) for p in include_dirs], language="c++", ) ] setup( name="sglang-kernel", version=_get_version(), packages=find_packages(where="python"), package_dir={"": "python"}, package_data={"sgl_kernel": ["*.metallib"]}, include_package_data=True, ext_modules=ext_modules, cmdclass={"build_ext": BuildMetalExtension}, )