1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
|
{
lib,
buildPythonPackage,
fetchFromGitHub,
# build-system
torch,
setuptools,
# env
cudaPackages,
# buildInputs
pybind11,
rdma-core,
config,
cudaCapabilities ? torch.cudaCapabilities,
cudaSupport ? config.cudaSupport,
}:
let
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 "1" else "0";
CUDA_HOME = (lib.getBin cudaPackages.cuda_nvcc).outPath;
}
# 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 = (lib.getInclude cudaPackages.libnvshmem).outPath;
}
);
buildInputs = [
pybind11
rdma-core
]
++ lib.optionals cudaSupport (
with cudaPackages;
[
cuda_cccl # <nv/target>
cuda_cudart # cuda_runtime.h
libcublas # cublas_v2.h
libcusolver # cusolverDn.h
libcusparse # cusparse.h
]
# On CUDA >=13.0, crt/host_config.h is shipped in cudaPackages.cuda_crt
++ lib.optionals cuda_crt.meta.available [
cuda_crt # crt/host_config.h
]
);
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;
};
})
|
