diff --git a/flake.nix b/flake.nix index 34d00b8f..e8d3ba1d 100644 --- a/flake.nix +++ b/flake.nix @@ -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" ''} diff --git a/nix/mlx_patch_fmod.patch b/nix/mlx_patch_fmod.patch new file mode 100644 index 00000000..323de7a3 --- /dev/null +++ b/nix/mlx_patch_fmod.patch @@ -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) { + return x % y; ++ } else if constexpr ( ++ cuda::std::is_same_v || ++ cuda::std::is_same_v) { ++ float yf = static_cast(y); ++ float r = cuda::std::fmod(static_cast(x), y); ++ if (r != 0.0f && ((r < 0.0f) != (yf < 0.0f))) { ++ r += yf; ++ } ++ return static_cast(r); + } else { + T r = cuda::std::fmod(x, y); + if (r != 0 && (r < 0 != y < 0)) { diff --git a/pyproject.toml b/pyproject.toml index 8c805eb8..0465ccec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/python/parts.nix b/python/parts.nix index 75fc1057..1b675e74 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -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 = { diff --git a/uv.lock b/uv.lock index 490aa038..c03c9a32 100644 --- a/uv.lock +++ b/uv.lock @@ -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'",