diff --git a/pkgs/development/python-modules/deep-ep/default.nix b/pkgs/development/python-modules/deep-ep/default.nix new file mode 100644 index 000000000000..a26eafdd22e9 --- /dev/null +++ b/pkgs/development/python-modules/deep-ep/default.nix @@ -0,0 +1,114 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + torch, + setuptools, + + # env + symlinkJoin, + cudaPackages, + + # buildInputs + pybind11, + rdma-core, + + config, + cudaCapabilities ? torch.cudaCapabilities, + cudaSupport ? config.cudaSupport, +}: +let + inherit (lib) + getBin + getInclude + ; + + minSupportedCudaCapability = "8.0"; # build fails with 7.5 + + minCudaCapability = builtins.head ( + builtins.sort (a: b: builtins.compareVersions a b < 0) cudaCapabilities + ); + + cudaCapabilities' = + if lib.versionOlder minCudaCapability minSupportedCudaCapability then + throw '' + CUDA capability "${minCudaCapability}" from `cudaCapabilities` is incompatible with DeepEP. + Only CUDA capabilities must be "${minSupportedCudaCapability}" or newer. + Build `python3Packages.deep-ep` with a compatible `cudaCapabilities` list. + '' + else + cudaCapabilities; + + disableSm90Features = lib.versionOlder minCudaCapability "9.0"; +in +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { + pname = "deep-ep"; + version = "1.2.1"; + pyproject = true; + + src = fetchFromGitHub { + owner = "deepseek-ai"; + repo = "DeepEP"; + tag = "v${finalAttrs.version}"; + hash = "sha256-xURR3uBAwKjDTNEG9p/vRRhH4Ldiz/u6kD/a+DPn5/Q="; + }; + + build-system = [ + setuptools + torch + ]; + + env = lib.optionalAttrs cudaSupport ( + { + TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep " " cudaCapabilities'}"; + + DISABLE_SM90_FEATURES = + if disableSm90Features then + lib.warn '' + python3Packages.deepep: Disabling SM90 features as the provided `cudaCapabilities` list include '${minCudaCapability}' + '' "1" + else + "0"; + + CUDA_HOME = symlinkJoin { + name = "cuda-redist"; + paths = with cudaPackages; [ + (getBin cuda_nvcc) + + (getInclude cuda_cccl) # + (getInclude cuda_cudart) # cuda_runtime.h + (getInclude libcublas) # cublas_v2.h + (getInclude libcusolver) # cusolverDn.h + (getInclude libcusparse) # cusparse.h + ]; + }; + } + + # nvshmem must be disabled (unsetting NVSHMEM_DIR) when supporting <9.0 capabilities + # https://github.com/deepseek-ai/DeepEP/blob/v1.2.1/setup.py#L65 + // lib.optionalAttrs (!disableSm90Features) { + NVSHMEM_DIR = (getInclude cudaPackages.libnvshmem).outPath; + } + ); + + buildInputs = [ + pybind11 + rdma-core + ]; + + pythonImportsCheck = [ "deep_ep" ]; + + # Tests check internode communications which is not possible in the sandbox + doCheck = false; + + meta = { + description = "Efficient expert-parallel communication library"; + homepage = "https://github.com/deepseek-ai/DeepEP"; + changelog = "https://github.com/deepseek-ai/DeepEP/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ GaetanLepage ]; + broken = !cudaSupport; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 115f36295c15..fc5a11fe7acb 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -3756,6 +3756,8 @@ self: super: with self; { deep-chainmap = callPackage ../development/python-modules/deep-chainmap { }; + deep-ep = callPackage ../development/python-modules/deep-ep { }; + deep-translator = callPackage ../development/python-modules/deep-translator { }; deepdiff = callPackage ../development/python-modules/deepdiff { };