Add Metal extension build support

[XXXXRT666]

Proposed changes

Add a setuptools-based workflow for building MLX extensions from C++ and Metal source files without requiring a project CMakeLists.txt.

• Add MetalExtension and BuildExtension with Ninja-by-default CMake builds, automatic fallback, generator-specific build directories, and optional nanobind stub generation.
• Compile Metal sources independently to AIR files with automatic header dependency tracking before linking the final metallib.
• Support Debug Metal builds, propagate the macOS deployment target, and detect logging support from the active Metal compiler.
• Inject the extension module and Metal library names automatically.
• Extend mlx_build_metallib to accept custom Metal compiler options.
• Add macOS CI coverage, tests, documentation, and a standalone activation example.
• Rename the existing CMake-based example to cmake_extension so both workflows remain separate.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[XXXXRT666]

One note on the Python Stable ABI: the existing CMake extension example passes STABLE_ABI to nanobind_add_module, but find_package(Python ...) only requests Development.Module. nanobind requires CPython 3.12+ and Python::SABIModule; otherwise, it silently disables STABLE_ABI. MLX’s core target currently has the same configuration, so both are effectively built against the regular CPython ABI.

Enabling Python::SABIModule only for MetalExtension could produce an abi3 binary, but it would not interoperate with nanobind-bound types from a regular-ABI mlx.core, such as array, Stream, and Device. Supporting abi3 properly would require building mlx.core with the Stable ABI as well and producing the corresponding abi3 wheel metadata.

The Python Stable ABI would still not provide a stable MLX C++ ABI across MLX releases. For now, this PR therefore keeps MetalExtension on the regular CPython ABI, and the example pins its mlx dependency to the version used for the build.

[XXXXRT666]

Hi, I wanted to check whether there are any updates on this PR or any remaining concerns I can address.

The main motivation is to reduce the amount of project-specific build machinery required for MLX Metal extensions. Some projects currently maintain their own compiler, linker, nanobind, and metallib build logic, for example:

• SGLang’s setup_metal.py
• vLLM Metal’s build.py

For example, I successfully built and tested SGLang’s rope_pool_fused Metal kernel using this API, reducing its custom build setup to roughly:

setup(
    name="sglang-kernel",
    version=_get_version(),
    packages=find_packages(where="python"),
    package_dir={"": "python"},
    package_data={"sgl_kernel": ["*.metallib", "*.pyi"]},
    ext_modules=[
        extension.MetalExtension(
            "sgl_kernel._metal",
            sources=[
                "csrc/metal/rope_pool_fused.cpp",
                "csrc/metal/rope_pool_fused.metal",
            ],
            include_dirs=["csrc", "csrc/metal"],
            extra_compile_args={
                "cxx": ["-O3", "-fvisibility=hidden"],
                "metal": [f"-std={metal_std}", "-O3"],
            },
        )
    ],
    cmdclass={"build_ext": extension.BuildExtension},
    install_requires=[f"mlx=={package_version('mlx')}"],
)

This allows other projects to describe their sources and project-specific flags while MLX handles the common extension build plumbing.

[zcbenz]

@jundot @WindChimeRan @yeahdongcn I see you have been heavily using custom Metal extensions, would you mind checking whether this feature would be helpful to you?

[yeahdongcn]

@jundot @WindChimeRan @yeahdongcn I see you have been heavily using custom Metal extensions, would you mind checking whether this feature would be helpful to you?

No problem. I can take a look. Thanks!

[WindChimeRan]

Thanks for the ping. @zcbenz @XXXXRT666

The direction looks useful, but the current API does not fully cover vllm-metal: we have one Python extension with multiple metallibs, including a separate NAX library targeting macOS 26.2 while the others target 15.0. MetalExtension currently assumes one same-named metallib and one deployment target, so it cannot replace our custom build as written.

[yeahdongcn]

Thanks for putting this together. My main question is whether MetalExtension should belong to MLX or PyTorch.

For CUDA/ROCm/MUSA, we use from torch.utils.cpp_extension import BuildExtension, CUDAExtension, so I'd prefer to follow a similar approach here. The current implementation always discovers and links against MLX, which makes it tightly coupled to the MLX framework.

I'd prefer a framework-neutral MetalExtension that provides the common build functionality, while allowing the specific framework dependencies to be added separately when needed.

[yeahdongcn]

SGLang's target architecture has a Torch-owned ModelRunner with MLX as an optional operator provider, so the downstream build contract needs an explicit mlx, torch, both, or metallib choice rather than an MLX-only build graph.

[yeahdongcn]

Requiring both a C++ source and a Metal source rules out a framework-neutral metallib and Torch's public torch.mps.load_metallib path, which needs no host extension. Could the API model a MetalLibrary as a first-class artifact and attach zero, one, or two backend-specific adapters? metallib would build only the library; mlx and torch would add their own adapters; both would build two separately named adapters over one validated-compatible metallib, or backend-specific library variants when the shader contracts differ.

This needs an artifact split rather than only making the C++ list optional:
setuptools.Extension, get_ext_fullpath, stub generation and sidecar copy
currently all assume an importable host module. A single host binary should
also not be expected to accept both mlx.core.array and torch.Tensor.

[yeahdongcn]

Could dependency discovery be conditional so Torch/metallib modes do not discover or link MLX, MLX mode does not require Torch, and both builds and tests two isolated adapters? If the facade remains in mlx.extension, could we also clarify whether requiring the MLX Python package solely to access a Torch/metallib builder is intentional? Native Torch support also needs an explicit composition contract with torch.utils.cpp_extension.BuildExtension: setup() has only one cmdclass["build_ext"], while this command currently delegates non-MetalExtension objects only to setuptools.

[XXXXRT666]

Thanks for the ping. @zcbenz @XXXXRT666

The direction looks useful, but the current API does not fully cover vllm-metal: we have one Python extension with multiple metallibs, including a separate NAX library targeting macOS 26.2 while the others target 15.0. MetalExtension currently assumes one same-named metallib and one deployment target, so it cannot replace our custom build as written.

It’s supported now. I’ve put together a minimal reproduction here that demonstrates how vLLM-Metal kernels can be compiled. I think python setup.py build_ext --inplace should work well for that.

[XXXXRT666]

Thanks for putting this together. My main question is whether MetalExtension should belong to MLX or PyTorch.

For CUDA/ROCm/MUSA, we use from torch.utils.cpp_extension import BuildExtension, CUDAExtension, so I'd prefer to follow a similar approach here. The current implementation always discovers and links against MLX, which makes it tightly coupled to the MLX framework.

I'd prefer a framework-neutral MetalExtension that provides the common build functionality, while allowing the specific framework dependencies to be added separately when needed.

MLX and PyTorch can coexist within the same extension here