build mlx cpu for now
This commit is contained in:
Evan
@@ -81,9 +81,11 @@
|
||||
_module.args.cudaPkgs = import inputs.nixpkgs
|
||||
{
|
||||
inherit system;
|
||||
config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
|
||||
cudaSupport = true;
|
||||
cudaCapabilities = [ "12.1" ];
|
||||
config = {
|
||||
allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
|
||||
cudaSupport = true;
|
||||
cudaCapabilities = [ "12.1" ];
|
||||
};
|
||||
};
|
||||
# Allow unfree for metal-toolchain (needed for Darwin Metal packages)
|
||||
_module.args.pkgs = import inputs.nixpkgs {
|
||||
@@ -174,7 +176,7 @@
|
||||
shellHook = ''
|
||||
unset PYTHONPATH
|
||||
export REPO_ROOT=$(git rev-parse --show-toplevel)
|
||||
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.python}/lib"
|
||||
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.editable-venv}/lib"
|
||||
${lib.optionalString stdenv.isLinux ''
|
||||
export LD_LIBRARY_PATH="${openssl.out}/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH"
|
||||
''}
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh
|
||||
index b0b79628..bbc377be 100644
|
||||
--- a/mlx/backend/cuda/device/binary_ops.cuh
|
||||
+++ b/mlx/backend/cuda/device/binary_ops.cuh
|
||||
@@ -46,6 +46,15 @@ struct Remainder {
|
||||
}
|
||||
} else if constexpr (is_complex_v<T>) {
|
||||
return x % y;
|
||||
+ } else if constexpr (
|
||||
+ cuda::std::is_same_v<T, __half> ||
|
||||
+ cuda::std::is_same_v<T, __nv_bfloat16>) {
|
||||
+ float yf = static_cast<float>(y);
|
||||
+ float r = cuda::std::fmod(static_cast<float>(x), y);
|
||||
+ if (r != 0.0f && ((r < 0.0f) != (yf < 0.0f))) {
|
||||
+ r += yf;
|
||||
+ }
|
||||
+ return static_cast<T>(r);
|
||||
} else {
|
||||
T r = cuda::std::fmod(x, y);
|
||||
if (r != 0 && (r < 0 != y < 0)) {
|
||||
+5
-1
@@ -18,13 +18,14 @@ dependencies = [
|
||||
"exo_pyo3_bindings", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"mlx",
|
||||
"mlx[cpu]; sys_platform == 'linux'",
|
||||
"mlx-lm",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"hypercorn>=0.18.0",
|
||||
"openai-harmony>=0.0.8",
|
||||
"httpx>=0.28.1",
|
||||
"tomlkit>=0.14.0",
|
||||
"mflux==0.16.9; sys_platform == 'darwin'",
|
||||
"mflux==0.17.2; sys_platform == 'darwin'",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
@@ -50,6 +51,7 @@ cuda = [
|
||||
"torch>=2.10.0; sys_platform == 'linux'",
|
||||
"vllm>=0.13.0; sys_platform == 'linux'",
|
||||
"fastsafetensors>=0.1.10; sys_platform == 'linux'",
|
||||
"mlx[cuda13]; sys_platform == 'linux'"
|
||||
]
|
||||
|
||||
###
|
||||
@@ -164,6 +166,8 @@ vllm = [
|
||||
]
|
||||
fastsafetensors = ["setuptools", "pybind11"]
|
||||
torch = ["typing-extensions"]
|
||||
torchvision = ["torch"]
|
||||
torchaudio = ["torch"]
|
||||
|
||||
###
|
||||
# ruff configuration
|
||||
|
||||
+362
-370
@@ -1,8 +1,329 @@
|
||||
{ inputs, ... }:
|
||||
{
|
||||
perSystem =
|
||||
{ self', pkgs, cudaPkgs, lib, ... }:
|
||||
let
|
||||
mkPythonSet = { self', pkgs, lib, editable ? false }:
|
||||
let
|
||||
inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux isx86_64;
|
||||
inherit (pkgs.config) cudaSupport;
|
||||
inherit (pkgs) cudaPackages;
|
||||
cudaLibs = with cudaPackages; [
|
||||
cuda_cudart
|
||||
cuda_cccl
|
||||
cuda_cupti
|
||||
cuda_nvrtc
|
||||
cuda_nvtx
|
||||
cudnn
|
||||
libcufile
|
||||
libcublas
|
||||
libcufft
|
||||
libcurand
|
||||
libcusolver
|
||||
libcusparse
|
||||
libcusparse_lt
|
||||
libnvjitlink
|
||||
libnvshmem
|
||||
nccl
|
||||
];
|
||||
cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
|
||||
mkdir -p $out/include
|
||||
ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
|
||||
'';
|
||||
|
||||
cudaRoot = pkgs.symlinkJoin {
|
||||
name = "cuda-merged-exo";
|
||||
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
|
||||
};
|
||||
exoOverlay = final: prev:
|
||||
{
|
||||
mlx = prev.mlx.overrideAttrs (old:
|
||||
let
|
||||
# Static dependencies included directly during compilation
|
||||
gguf-tools = pkgs.fetchFromGitHub {
|
||||
owner = "antirez";
|
||||
repo = "gguf-tools";
|
||||
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
|
||||
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
|
||||
};
|
||||
|
||||
metal_cpp = pkgs.fetchzip {
|
||||
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
|
||||
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
|
||||
};
|
||||
|
||||
nanobind = pkgs.fetchFromGitHub {
|
||||
owner = "wjakob";
|
||||
repo = "nanobind";
|
||||
rev = "v2.10.2";
|
||||
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
|
||||
fetchSubmodules = true;
|
||||
};
|
||||
|
||||
nvtx = pkgs.fetchFromGitHub {
|
||||
name = "nvtx3";
|
||||
owner = "NVIDIA";
|
||||
repo = "NVTX";
|
||||
rev = "v3.1.1";
|
||||
hash = "sha256-sx72N+Gskg9Vtqc3sXsWoE/2PHFI2Hq08lEaw0sll5Y=";
|
||||
};
|
||||
cudnn = pkgs.fetchFromGitHub {
|
||||
name = "cudnn_frontend";
|
||||
owner = "NVIDIA";
|
||||
repo = "cudnn-frontend";
|
||||
rev = "v1.16.0";
|
||||
hash = "sha256-+8aBl9dKd2Uz50XoOr91NRyJ4OGJtzfDNNNYGQJ9b94=";
|
||||
};
|
||||
mlx_cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
|
||||
mkdir -p $out/include
|
||||
${pkgs.tree}/bin/tree ${cudaPackages.cuda_cccl} -L 3
|
||||
exit 1
|
||||
ln -s ${cudaPackages.cuda_cccl}/include/cuda $out/include/cuda
|
||||
ln -s ${cudaPackages.cuda_cccl}/include/nv $out/include/nv
|
||||
'';
|
||||
in
|
||||
{
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake ] ++ lib.optionals isDarwin [ self'.packages.metal-toolchain ] ++ lib.optionals cudaSupport [
|
||||
cudaPackages.cuda_nvcc
|
||||
pkgs.autoAddDriverRunpath
|
||||
pkgs.autoPatchelfHook
|
||||
];
|
||||
# TODO: non-sdk_26 support
|
||||
buildInputs = (old.buildInputs or [ ])
|
||||
++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.openblas ]
|
||||
++ lib.optionals isDarwin [ pkgs.apple-sdk_26 ]
|
||||
++ lib.optionals cudaSupport (cudaLibs ++ [ cudaPackages.cudnn ]);
|
||||
prePatch = ''${pkgs.tree}/bin/tree'';
|
||||
patches = (old.patches or [ ]) ++ lib.optionals cudaSupport [ ../nix/mlx_patch_fmod.patch ];
|
||||
|
||||
postPatch = ''
|
||||
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
|
||||
--replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
|
||||
'';
|
||||
|
||||
DEV_RELEASE = 1;
|
||||
CMAKE_ARGS = toString ([
|
||||
(lib.cmakeBool "USE_SYSTEM_FMT" true)
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
|
||||
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
|
||||
(lib.cmakeBool "MLX_BUILD_CPU" true)
|
||||
(lib.cmakeBool "MLX_BUILD_METAL" isDarwin)
|
||||
(lib.cmakeBool "MLX_BUILD_CUDA" false)
|
||||
(lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
|
||||
] ++ lib.optionals cudaSupport [
|
||||
(lib.cmakeOptionType "filepath" "CUDAToolkit_ROOT" "${cudaRoot}")
|
||||
(lib.cmakeOptionType "string" "MLX_CUDA_ARCHITECTURES" "121")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${cudaPackages.cutlass}")
|
||||
# TODO: replace with cudaPackages.cudnn
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUDNN" "${cudnn}")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CCCL" "${cudaPackages.cuda_cccl}")
|
||||
# TODO: replace with cudaPackages.nvtx
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NVTX3" "${nvtx}")
|
||||
] ++ lib.optionals isDarwin [
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
|
||||
] ++ lib.optionals (isDarwin && isx86_64) [
|
||||
(lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
|
||||
]);
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
|
||||
});
|
||||
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
|
||||
pname = "exo-pyo3-bindings";
|
||||
version = "0.1.0";
|
||||
src = self'.packages.exo_pyo3_bindings;
|
||||
# Install from pre-built wheel
|
||||
nativeBuildInputs = [ final.pyprojectWheelHook ];
|
||||
dontStrip = true;
|
||||
passthru = prev.exo-pyo3-bindings.passthru or { };
|
||||
postInstall = ''
|
||||
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
|
||||
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
|
||||
touch $siteDir/py.typed
|
||||
'';
|
||||
};
|
||||
torch = prev.torch.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
|
||||
});
|
||||
torchaudio = prev.torchaudio.overrideAttrs (old:
|
||||
{
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport [
|
||||
cudaPackages.cuda_cudart
|
||||
];
|
||||
preFixup = (old.preFixup or "") + ''
|
||||
addAutoPatchelfSearchPath "${final.torch}"
|
||||
'';
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
|
||||
});
|
||||
torchvision = prev.torchvision.overrideAttrs (old:
|
||||
{
|
||||
nativebuildInputs = (old.nativebuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
preFixup = (old.preFixup or "") + ''
|
||||
addAutoPatchelfSearchPath "${final.torch}"
|
||||
'';
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
|
||||
});
|
||||
xgrammar = prev.xgrammar.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake pkgs.autoPatchelfHook ];
|
||||
});
|
||||
# Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
|
||||
vllm = prev.vllm.overrideAttrs (old:
|
||||
let
|
||||
cutlass = pkgs.fetchFromGitHub {
|
||||
name = "cutlass-source";
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
tag = "v4.2.1";
|
||||
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
|
||||
};
|
||||
triton-kernels = pkgs.fetchFromGitHub {
|
||||
owner = "triton-lang";
|
||||
repo = "triton";
|
||||
tag = "v3.6.0";
|
||||
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
|
||||
};
|
||||
|
||||
cutlass-flashmla = pkgs.fetchFromGitHub {
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
|
||||
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
|
||||
};
|
||||
|
||||
flashmla = pkgs.stdenv.mkDerivation {
|
||||
pname = "flashmla";
|
||||
version = "1.0.0";
|
||||
|
||||
src = pkgs.fetchFromGitHub {
|
||||
name = "FlashMLA-source";
|
||||
owner = "vllm-project";
|
||||
repo = "FlashMLA";
|
||||
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
|
||||
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
|
||||
};
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass-flashmla} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
qutlass = pkgs.fetchFromGitHub {
|
||||
name = "qutlass-source";
|
||||
owner = "IST-DASLab";
|
||||
repo = "qutlass";
|
||||
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
|
||||
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
|
||||
};
|
||||
vllm-flash-attn = pkgs.stdenv.mkDerivation {
|
||||
pname = "vllm-flash-attn";
|
||||
version = "2.7.2.post1";
|
||||
|
||||
src = pkgs.fetchFromGitHub {
|
||||
name = "flash-attention-source";
|
||||
owner = "vllm-project";
|
||||
repo = "flash-attention";
|
||||
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
|
||||
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
|
||||
};
|
||||
|
||||
patches = [
|
||||
(pkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
|
||||
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
|
||||
})
|
||||
(pkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
|
||||
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
|
||||
})
|
||||
];
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
in
|
||||
{
|
||||
patches = (old.patches or [ ]) ++ [ ../nix/vllm_uv2nix_cmake.patch ];
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
pkgs.cmake
|
||||
pkgs.ninja
|
||||
pkgs.autoAddDriverRunpath
|
||||
] ++ lib.optionals cudaSupport [
|
||||
cudaPackages.cuda_nvcc
|
||||
];
|
||||
# TODO: vllm rocm/cpu
|
||||
VLLM_TARGET_DEVICE = "empty";
|
||||
# TODO: vllm non cuda13 support, more arch's, etc.
|
||||
} // lib.optionalAttrs cudaSupport {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ cudaLibs;
|
||||
|
||||
CUDA_HOME = "${cudaRoot}";
|
||||
CUDAToolkit_ROOT = "${cudaRoot}";
|
||||
CUDACXX = "${cudaRoot}/bin/nvcc";
|
||||
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
|
||||
VLLM_TARGET_DEVICE = "cuda";
|
||||
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
|
||||
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
|
||||
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
|
||||
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
|
||||
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
|
||||
CAFFE2_USE_CUDNN = "ON";
|
||||
CAFFE2_USE_CUFILE = "ON";
|
||||
CUTLASS_ENABLE_CUBLAS = "ON";
|
||||
CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
|
||||
|
||||
UV2NIX_CMAKE_FLAGS_JSON = builtins.toJSON [
|
||||
"-DCUDAToolkit_ROOT=${cudaRoot}"
|
||||
"-DCMAKE_CUDA_COMPILER=${cudaRoot}/bin/nvcc"
|
||||
"-DCMAKE_PREFIX_PATH=${cudaRoot}"
|
||||
"-DFETCHCONTENT_SOURCE_DIR_CUTLASS=${lib.getDev cutlass}"
|
||||
"-DFLASH_MLA_SRC_DIR=${lib.getDev flashmla}"
|
||||
"-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
|
||||
"-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
|
||||
"-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
|
||||
"-DCUTLASS_NVCC_ARCHS_ENABLED=${cudaPackages.flags.cmakeCudaArchitecturesString}"
|
||||
"-DCAFFE2_USE_CUDNN=ON"
|
||||
"-DCAFFE2_USE_CUFILE=ON"
|
||||
"-DCUTLASS_ENABLE_CUBLAS=ON"
|
||||
];
|
||||
});
|
||||
} // lib.optionalAttrs cudaSupport {
|
||||
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core ];
|
||||
propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ pkgs.util-linux ];
|
||||
});
|
||||
nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ (with cudaPackages; [ libnvjitlink libcublas libcusparse ]);
|
||||
});
|
||||
nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ cudaPackages.libnvjitlink ];
|
||||
});
|
||||
nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
|
||||
});
|
||||
nvidia-cutlass-dsl-libs-base = prev.nvidia-cutlass-dsl-libs-base.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
|
||||
});
|
||||
};
|
||||
|
||||
# Load workspace from uv.lock
|
||||
workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
|
||||
workspaceRoot = ../.;
|
||||
@@ -11,381 +332,55 @@
|
||||
# Create overlay from workspace
|
||||
# Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
|
||||
overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
|
||||
|
||||
# Override overlay to inject Nix-built components
|
||||
exoOverlay = final: prev: {
|
||||
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
|
||||
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
|
||||
# Copy .pyi stub + py.typed marker so basedpyright can find the types.
|
||||
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
|
||||
pname = "exo-pyo3-bindings";
|
||||
version = "0.1.0";
|
||||
src = self'.packages.exo_pyo3_bindings;
|
||||
# Install from pre-built wheel
|
||||
nativeBuildInputs = [ final.pyprojectWheelHook ];
|
||||
dontStrip = true;
|
||||
passthru = prev.exo-pyo3-bindings.passthru or { };
|
||||
postInstall = ''
|
||||
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
|
||||
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
|
||||
touch $siteDir/py.typed
|
||||
'';
|
||||
};
|
||||
};
|
||||
|
||||
inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux;
|
||||
|
||||
# Overlay to provide build systems and custom packages
|
||||
cudaOverlay = final: prev:
|
||||
let
|
||||
torchLibs = [
|
||||
final.nvidia-cuda-runtime
|
||||
final.nvidia-cuda-nvrtc
|
||||
final.nvidia-cuda-cupti
|
||||
final.nvidia-nvjitlink
|
||||
final.nvidia-cudnn-cu13
|
||||
final.nvidia-cusparse
|
||||
final.nvidia-cusparselt-cu13
|
||||
final.nvidia-cufile
|
||||
final.nvidia-nvshmem-cu13
|
||||
final.nvidia-nccl-cu13
|
||||
final.nvidia-cublas
|
||||
final.nvidia-cufft
|
||||
final.nvidia-curand
|
||||
final.nvidia-cusolver
|
||||
];
|
||||
cutlass = cudaPkgs.fetchFromGitHub {
|
||||
name = "cutlass-source";
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
tag = "v4.2.1";
|
||||
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
|
||||
};
|
||||
triton-kernels = cudaPkgs.fetchFromGitHub {
|
||||
owner = "triton-lang";
|
||||
repo = "triton";
|
||||
tag = "v3.6.0";
|
||||
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
|
||||
};
|
||||
|
||||
cutlass-flashmla = cudaPkgs.fetchFromGitHub {
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
|
||||
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
|
||||
};
|
||||
|
||||
flashmla = cudaPkgs.stdenv.mkDerivation {
|
||||
pname = "flashmla";
|
||||
version = "1.0.0";
|
||||
|
||||
src = cudaPkgs.fetchFromGitHub {
|
||||
name = "FlashMLA-source";
|
||||
owner = "vllm-project";
|
||||
repo = "FlashMLA";
|
||||
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
|
||||
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
|
||||
};
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass-flashmla} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
qutlass = cudaPkgs.fetchFromGitHub {
|
||||
name = "qutlass-source";
|
||||
owner = "IST-DASLab";
|
||||
repo = "qutlass";
|
||||
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
|
||||
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
|
||||
};
|
||||
vllm-flash-attn = cudaPkgs.stdenv.mkDerivation {
|
||||
pname = "vllm-flash-attn";
|
||||
version = "2.7.2.post1";
|
||||
|
||||
src = cudaPkgs.fetchFromGitHub {
|
||||
name = "flash-attention-source";
|
||||
owner = "vllm-project";
|
||||
repo = "flash-attention";
|
||||
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
|
||||
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
|
||||
};
|
||||
|
||||
patches = [
|
||||
(cudaPkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
|
||||
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
|
||||
})
|
||||
(cudaPkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
|
||||
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
|
||||
})
|
||||
];
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
|
||||
cudaLibs = with cudaPkgs.cudaPackages_13; [
|
||||
cuda_cudart
|
||||
cuda_cccl
|
||||
cuda_cupti
|
||||
cuda_nvrtc
|
||||
cuda_nvtx
|
||||
cudnn
|
||||
libcufile
|
||||
libcublas
|
||||
libcublas
|
||||
libcufft
|
||||
libcurand
|
||||
libcusolver
|
||||
libcusparse
|
||||
libcusparse_lt
|
||||
libcufile
|
||||
libnvjitlink
|
||||
libnvshmem
|
||||
nccl
|
||||
];
|
||||
|
||||
cuda_cccl_compat = cudaPkgs.runCommand "cuda-cccl-compat" { } ''
|
||||
mkdir -p $out/include
|
||||
ln -s ${cudaPkgs.cudaPackages_13.cuda_cccl}/include $out/include/cccl
|
||||
'';
|
||||
|
||||
cudaRoot = cudaPkgs.symlinkJoin {
|
||||
name = "cuda-merged-exo";
|
||||
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPkgs.cudaPackages_13.cuda_nvcc cuda_cccl_compat ]);
|
||||
};
|
||||
in
|
||||
{
|
||||
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [
|
||||
cudaPkgs.rdma-core
|
||||
cudaPkgs.autoAddDriverRunpath
|
||||
];
|
||||
propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ cudaPkgs.util-linux ];
|
||||
});
|
||||
nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ (with cudaPkgs.cudaPackages_13; [ libnvjitlink libcublas libcusparse cudaPkgs.autoAddDriverRunpath ]);
|
||||
});
|
||||
nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ cudaPkgs.cudaPackages_13.libnvjitlink cudaPkgs.autoAddDriverRunpath ];
|
||||
});
|
||||
nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ cudaPkgs.rdma-core cudaPkgs.pmix cudaPkgs.libfabric cudaPkgs.ucx cudaPkgs.openmpi cudaPkgs.autoAddDriverRunpath ];
|
||||
});
|
||||
nvidia-cutlass-dsl-libs-base = prev.nvidia-cutlass-dsl-libs-base.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ cudaPkgs.autoAddDriverRunpath ];
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
|
||||
});
|
||||
|
||||
torch = prev.torch.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ cudaLibs ++ [ cudaPkgs.autoAddDriverRunpath ];
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
|
||||
});
|
||||
torchaudio = prev.torchaudio.overrideAttrs (old:
|
||||
{
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [
|
||||
final.torch
|
||||
cudaPkgs.cudaPackages_13.cuda_cudart
|
||||
cudaPkgs.autoAddDriverRunpath
|
||||
];
|
||||
preFixup = (old.preFixup or "") + ''
|
||||
addAutoPatchelfSearchPath "${final.torch}"
|
||||
'';
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
|
||||
});
|
||||
torchvision = prev.torchvision.overrideAttrs (old:
|
||||
{
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [
|
||||
final.torch
|
||||
cudaPkgs.autoAddDriverRunpath
|
||||
];
|
||||
preFixup = (old.preFixup or "") + ''
|
||||
addAutoPatchelfSearchPath "${final.torch}"
|
||||
'';
|
||||
autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
|
||||
});
|
||||
xgrammar = prev.xgrammar.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ cudaPkgs.cmake cudaPkgs.autoPatchelfHook ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ ];
|
||||
# patches = (old.patches or [ ]) ++ [ ../nix/xgrammar_cmake.patch ];
|
||||
});
|
||||
vllm = prev.vllm.overrideAttrs (old: {
|
||||
patches = (old.patches or [ ]) ++ [ ../nix/vllm_uv2nix_cmake.patch ];
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
cudaPkgs.cmake
|
||||
cudaPkgs.ninja
|
||||
cudaPkgs.autoAddDriverRunpath
|
||||
cudaPkgs.cudaPackages_13.cuda_nvcc
|
||||
];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ cudaLibs;
|
||||
|
||||
CUDA_HOME = "${cudaRoot}";
|
||||
CUDAToolkit_ROOT = "${cudaRoot}";
|
||||
CUDACXX = "${cudaRoot}/bin/nvcc";
|
||||
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
|
||||
VLLM_TARGET_DEVICE = "cuda";
|
||||
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
|
||||
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
|
||||
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
|
||||
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
|
||||
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
|
||||
CAFFE2_USE_CUDNN = "ON";
|
||||
CAFFE2_USE_CUFILE = "ON";
|
||||
CUTLASS_ENABLE_CUBLAS = "ON";
|
||||
CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
|
||||
|
||||
UV2NIX_CMAKE_FLAGS_JSON = builtins.toJSON [
|
||||
"-DCUDAToolkit_ROOT=${cudaRoot}"
|
||||
"-DCMAKE_CUDA_COMPILER=${cudaRoot}/bin/nvcc"
|
||||
"-DCMAKE_PREFIX_PATH=${cudaRoot}"
|
||||
"-DFETCHCONTENT_SOURCE_DIR_CUTLASS=${lib.getDev cutlass}"
|
||||
"-DFLASH_MLA_SRC_DIR=${lib.getDev flashmla}"
|
||||
"-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
|
||||
"-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
|
||||
"-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
|
||||
"-DCUTLASS_NVCC_ARCHS_ENABLED=${cudaPkgs.cudaPackages_13.flags.cmakeCudaArchitecturesString}"
|
||||
"-DCAFFE2_USE_CUDNN=ON"
|
||||
"-DCAFFE2_USE_CUFILE=ON"
|
||||
"-DCUTLASS_ENABLE_CUBLAS=ON"
|
||||
];
|
||||
});
|
||||
};
|
||||
|
||||
# Additional overlay for Linux-specific fixes (type checking env).
|
||||
# Native wheels have shared lib dependencies we don't need at type-check time.
|
||||
mlxOverlay = final: prev:
|
||||
let
|
||||
ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
|
||||
nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
|
||||
|
||||
# Static dependencies included directly during compilation
|
||||
gguf-tools = pkgs.fetchFromGitHub {
|
||||
owner = "antirez";
|
||||
repo = "gguf-tools";
|
||||
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
|
||||
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
|
||||
};
|
||||
|
||||
metal_cpp = pkgs.fetchzip {
|
||||
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
|
||||
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
|
||||
};
|
||||
|
||||
nanobind = pkgs.fetchFromGitHub {
|
||||
owner = "wjakob";
|
||||
repo = "nanobind";
|
||||
rev = "v2.10.2";
|
||||
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
|
||||
fetchSubmodules = true;
|
||||
};
|
||||
|
||||
in
|
||||
{
|
||||
mlx = prev.mlx.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake ] ++ lib.optionals isDarwin [ self'.packages.metal-toolchain ];
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ gguf-tools pkgs.openblas pkgs.fmt ] ++ lib.optionals isDarwin [ pkgs.apple-sdk_26 ];
|
||||
|
||||
postPatch = ''
|
||||
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
|
||||
--replace-fail "g++" "$CXX"
|
||||
'';
|
||||
|
||||
DEV_RELEASE = 1;
|
||||
CMAKE_ARGS = toString ([
|
||||
(lib.cmakeBool "USE_SYSTEM_FMT" true)
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
|
||||
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
|
||||
(lib.cmakeBool "MLX_BUILD_CPU" true)
|
||||
(lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
|
||||
] ++ lib.optionals isDarwin [
|
||||
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
|
||||
(lib.cmakeBool "MLX_BUILD_METAL" true)
|
||||
] ++ lib.optionals pkgs.config.cudaSupport [
|
||||
(lib.cmakeBool "MLX_BUILD_CUDA" true)
|
||||
]);
|
||||
|
||||
});
|
||||
};
|
||||
|
||||
|
||||
in
|
||||
(pkgs.callPackage inputs.pyproject-nix.build.packages {
|
||||
python = pkgs.python313;
|
||||
}).overrideScope (
|
||||
lib.composeManyExtensions ([
|
||||
inputs.pyproject-build-systems.overlays.default
|
||||
overlay
|
||||
exoOverlay
|
||||
] ++ lib.optionals editable [
|
||||
(workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "bench" ]; })
|
||||
])
|
||||
);
|
||||
|
||||
pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
|
||||
inherit python;
|
||||
}).overrideScope (
|
||||
lib.composeManyExtensions [
|
||||
inputs.pyproject-build-systems.overlays.default
|
||||
overlay
|
||||
exoOverlay
|
||||
mlxOverlay
|
||||
]
|
||||
);
|
||||
cudaPythonSet = (cudaPkgs.callPackage inputs.pyproject-nix.build.packages {
|
||||
python = cudaPkgs.python313;
|
||||
}).overrideScope (
|
||||
lib.composeManyExtensions [
|
||||
inputs.pyproject-build-systems.overlays.default
|
||||
overlay
|
||||
exoOverlay
|
||||
mlxOverlay
|
||||
cudaOverlay
|
||||
]
|
||||
);
|
||||
# mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
|
||||
# mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
|
||||
venvCollisionPaths = lib.optionals isLinux [
|
||||
"lib/python3.13/site-packages/cv2*"
|
||||
"lib/python3.13/site-packages/mlx*"
|
||||
"lib/python3.13/site-packages/nvidia*"
|
||||
];
|
||||
|
||||
editablePythonSet = cudaPythonSet.overrideScope (
|
||||
workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "bench" ]; }
|
||||
);
|
||||
in
|
||||
{
|
||||
perSystem =
|
||||
{ self', pkgs, cudaPkgs, lib, ... }:
|
||||
let
|
||||
inherit (pkgs.stdenv.hostPlatform) isDarwin;
|
||||
pythonSet = mkPythonSet { inherit self' pkgs lib; };
|
||||
# taking cudaPkgs.cudaPackages_13.pkgs creates a new nixpkgs that defaults to cuda 13
|
||||
cudaPythonSet = mkPythonSet { inherit self' lib; pkgs = cudaPkgs.cudaPackages_13.pkgs; };
|
||||
|
||||
editablePythonSet = mkPythonSet { inherit self' lib; pkgs = cudaPkgs.cudaPackages_13.pkgs; editable = true; };
|
||||
evenv = (editablePythonSet.mkVirtualEnv "exo-dev-env"
|
||||
{
|
||||
exo = [ "dev" "cuda" ];
|
||||
exo-pyo3-bindings = [ ];
|
||||
exo-bench = [ ];
|
||||
}).overrideAttrs {
|
||||
venvIgnoreCollisions = venvCollisionPaths;
|
||||
#venvIgnoreCollisions = venvCollisionPaths;
|
||||
venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
|
||||
};
|
||||
exoVenv = (pythonSet.mkVirtualEnv "exo-env" {
|
||||
exo = [ ];
|
||||
exo-pyo3-bindings = [ ];
|
||||
}).overrideAttrs {
|
||||
venvIgnoreCollisions = venvCollisionPaths;
|
||||
# venvIgnoreCollisions = venvCollisionPaths;
|
||||
venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
|
||||
};
|
||||
exoCudaVenv = (cudaPythonSet.mkVirtualEnv "exo-env" {
|
||||
exo = [ "cuda" ];
|
||||
exo-pyo3-bindings = [ ];
|
||||
}).overrideAttrs {
|
||||
venvIgnoreCollisions = venvCollisionPaths;
|
||||
# venvIgnoreCollisions = venvCollisionPaths;
|
||||
venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
|
||||
};
|
||||
|
||||
|
||||
|
||||
# Virtual environment with dev dependencies for testing
|
||||
testVenv = (pythonSet.mkVirtualEnv "exo-test-env"
|
||||
{
|
||||
@@ -393,7 +388,7 @@
|
||||
exo-pyo3-bindings = [ ];
|
||||
}
|
||||
).overrideAttrs {
|
||||
venvIgnoreCollisions = venvCollisionPaths;
|
||||
# venvIgnoreCollisions = venvCollisionPaths;
|
||||
};
|
||||
|
||||
mkPythonScript = name: path: pkgs.writeShellApplication {
|
||||
@@ -438,26 +433,21 @@
|
||||
|
||||
exoCudaPackage = cudaPkgs.runCommand "exo"
|
||||
{
|
||||
nativeBuildInputs = [ pkgs.makeWrapper ];
|
||||
nativeBuildInputs = [ cudaPkgs.makeWrapper ];
|
||||
}
|
||||
''
|
||||
mkdir -p $out/bin
|
||||
mkdir -p $out/bin
|
||||
|
||||
# Create wrapper script
|
||||
makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo \
|
||||
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
|
||||
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
|
||||
--prefix LD_LIBRARY_PATH : ${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib \
|
||||
${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
|
||||
# Create wrapper script
|
||||
makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo \
|
||||
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
|
||||
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
|
||||
--prefix LD_LIBRARY_PATH : /run/opengl-driver/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib \
|
||||
${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
|
||||
'';
|
||||
in
|
||||
{
|
||||
# Python package only available on macOS (requires MLX/Metal)
|
||||
packages = (lib.optionalAttrs isDarwin
|
||||
{
|
||||
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
|
||||
exo-test-env = testVenv;
|
||||
}) // {
|
||||
packages = {
|
||||
exo = exoPackage;
|
||||
exo-cuda = exoCudaPackage;
|
||||
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
|
||||
@@ -465,7 +455,9 @@
|
||||
exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
|
||||
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
|
||||
editable-venv = evenv;
|
||||
inherit python;
|
||||
} // lib.optionalAttrs isDarwin {
|
||||
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
|
||||
exo-test-env = testVenv;
|
||||
};
|
||||
|
||||
checks = {
|
||||
|
||||
@@ -806,6 +806,7 @@ dependencies = [
|
||||
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "mlx", marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
@@ -825,6 +826,7 @@ build = [
|
||||
]
|
||||
cuda = [
|
||||
{ name = "fastsafetensors", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx", marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "torch", version = "2.10.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "vllm", marker = "sys_platform == 'linux'" },
|
||||
@@ -853,8 +855,10 @@ requires-dist = [
|
||||
{ name = "huggingface-hub", specifier = ">=0.33.4" },
|
||||
{ name = "hypercorn", specifier = ">=0.18.0" },
|
||||
{ name = "loguru", specifier = ">=0.7.3" },
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin'", specifier = "==0.16.9" },
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin'", specifier = "==0.17.2" },
|
||||
{ name = "mlx", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
|
||||
{ name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
|
||||
{ name = "mlx", extras = ["cuda13"], marker = "sys_platform == 'linux' and extra == 'cuda'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
|
||||
{ name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation" },
|
||||
{ name = "msgspec", specifier = ">=0.19.0" },
|
||||
{ name = "nanobind", marker = "extra == 'build'" },
|
||||
@@ -896,7 +900,8 @@ dependencies = [
|
||||
{ name = "loguru", marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "math-verify", marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "7.34.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "5.29.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "7.34.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "tiktoken", marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "transformers", marker = "(sys_platform == 'darwin' and extra == 'project-9-exo-bench') or (sys_platform == 'linux' and extra == 'project-9-exo-bench') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
]
|
||||
@@ -2198,7 +2203,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "mflux"
|
||||
version = "0.16.9"
|
||||
version = "0.17.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "filelock", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
@@ -2212,6 +2217,7 @@ dependencies = [
|
||||
{ name = "piexif", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "pillow", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "platformdirs", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "5.29.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "regex", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "requests", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "safetensors", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
@@ -2224,9 +2230,9 @@ dependencies = [
|
||||
{ name = "twine", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "urllib3", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e3/dc/73dfccae395df22623fbb6f748657700fd9372be22904a0dafb825dc42d2/mflux-0.16.9.tar.gz", hash = "sha256:ff35e9386b5f026a6b97c786b3426e5341105de3356e55ec1b0c5951ff0ac8c8", size = 764296, upload-time = "2026-03-07T11:54:45.031Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/1b/75/65f791f54c9531b524813d18ff87a0d34866ef220a8fe2ad637437f3cb54/mflux-0.17.2.tar.gz", hash = "sha256:52dee2d27cf438a84648e5c1861b92ceb63a9ac06823d14452a78646a1d30ee7", size = 779264, upload-time = "2026-03-23T13:08:18.377Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/47/69/091df2d0f4a2677d5e2a8b326d568bd884d7ca37f15d3ff78e574aca2271/mflux-0.16.9-py3-none-any.whl", hash = "sha256:988e925653a024e81b46401fd5a2e423dd283428ba64949af4d9257e7df8ddd9", size = 1018804, upload-time = "2026-03-07T11:54:43.47Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/68/02/f94eca4e77b7d12685060461eb793cbc8c00e96cc7fe0ce376374201aed2/mflux-0.17.2-py3-none-any.whl", hash = "sha256:be1642b04847413c0a8ed1dae82ce1ca023e155b057d82a8301eca9c3fe08339", size = 1037451, upload-time = "2026-03-23T13:08:16.747Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2255,8 +2261,8 @@ image = [
|
||||
|
||||
[[package]]
|
||||
name = "mlx"
|
||||
version = "0.30.7.dev20260312+257d5692"
|
||||
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
|
||||
version = "0.31.2.dev20260330+e5e64331"
|
||||
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }
|
||||
|
||||
[[package]]
|
||||
name = "mlx-lm"
|
||||
@@ -2267,8 +2273,9 @@ dependencies = [
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra != 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "5.29.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "6.33.5", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "7.34.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra != 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "protobuf", version = "7.34.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-3-exo-cuda') or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
{ name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cuda' and extra == 'project-9-exo-bench')" },
|
||||
@@ -3304,6 +3311,26 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "protobuf"
|
||||
version = "5.29.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
resolution-markers = [
|
||||
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"python_full_version < '3.14' and sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"sys_platform == 'darwin' and extra == 'extra-3-exo-cuda' and extra != 'project-9-exo-bench'",
|
||||
"sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra != 'project-9-exo-bench'",
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7e/57/394a763c103e0edf87f0938dafcd918d53b4c011dfc5c8ae80f3b0452dbb/protobuf-5.29.6.tar.gz", hash = "sha256:da9ee6a5424b6b30fd5e45c5ea663aef540ca95f9ad99d1e887e819cdf9b8723", size = 425623, upload-time = "2026-02-04T22:54:40.584Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/88/9ee58ff7863c479d6f8346686d4636dd4c415b0cbeed7a6a7d0617639c2a/protobuf-5.29.6-cp310-abi3-win32.whl", hash = "sha256:62e8a3114992c7c647bce37dcc93647575fc52d50e48de30c6fcb28a6a291eb1", size = 423357, upload-time = "2026-02-04T22:54:25.805Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1c/66/2dc736a4d576847134fb6d80bd995c569b13cdc7b815d669050bf0ce2d2c/protobuf-5.29.6-cp310-abi3-win_amd64.whl", hash = "sha256:7e6ad413275be172f67fdee0f43484b6de5a904cc1c3ea9804cb6fe2ff366eda", size = 435175, upload-time = "2026-02-04T22:54:28.592Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/06/db/49b05966fd208ae3f44dcd33837b6243b4915c57561d730a43f881f24dea/protobuf-5.29.6-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:b5a169e664b4057183a34bdc424540e86eea47560f3c123a0d64de4e137f9269", size = 418619, upload-time = "2026-02-04T22:54:30.266Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/d7/48cbf6b0c3c39761e47a99cb483405f0fde2be22cf00d71ef316ce52b458/protobuf-5.29.6-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:a8866b2cff111f0f863c1b3b9e7572dc7eaea23a7fae27f6fc613304046483e6", size = 320284, upload-time = "2026-02-04T22:54:31.782Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e3/dd/cadd6ec43069247d91f6345fa7a0d2858bef6af366dbd7ba8f05d2c77d3b/protobuf-5.29.6-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:e3387f44798ac1106af0233c04fb8abf543772ff241169946f698b3a9a3d3ab9", size = 320478, upload-time = "2026-02-04T22:54:32.909Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/cb/e3065b447186cb70aa65acc70c86baf482d82bf75625bf5a2c4f6919c6a3/protobuf-5.29.6-py3-none-any.whl", hash = "sha256:6b9edb641441b2da9fa8f428760fc136a49cf97a52076010cf22a2ff73438a86", size = 173126, upload-time = "2026-02-04T22:54:39.462Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "protobuf"
|
||||
version = "6.33.5"
|
||||
@@ -3330,10 +3357,6 @@ name = "protobuf"
|
||||
version = "7.34.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
resolution-markers = [
|
||||
"python_full_version >= '3.14' and sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"python_full_version < '3.14' and sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"sys_platform == 'darwin' and extra == 'extra-3-exo-cuda' and extra != 'project-9-exo-bench'",
|
||||
"sys_platform == 'darwin' and extra != 'extra-3-exo-cuda' and extra != 'project-9-exo-bench'",
|
||||
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"python_full_version < '3.14' and sys_platform == 'linux' and extra != 'extra-3-exo-cuda' and extra == 'project-9-exo-bench'",
|
||||
"sys_platform == 'linux' and extra != 'extra-3-exo-cuda' and extra != 'project-9-exo-bench'",
|
||||
|
||||
Reference in New Issue
Block a user