Commit 6e46f918 authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python3Packages.deep-ep: cleanup, fix build on CUDA>=13.0

parent 0db35111

pkgs/development/python-modules/deep-ep/default.nix

+17 −20
Original line number Diff line number Diff line
@@ -8,7 +8,6 @@ 
  setuptools,



  # env

  symlinkJoin,

  cudaPackages,



  # buildInputs
@@ -20,11 +19,6 @@ 
  cudaSupport ? config.cudaSupport,

}:

let

  inherit (lib)

    getBin

    getInclude

    ;



  minSupportedCudaCapability = "8.0"; # build fails with 7.5



  minCudaCapability = builtins.head (
@@ -66,31 +60,34 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { 


      DISABLE_SM90_FEATURES = if disableSm90Features then "1" else "0";



      CUDA_HOME = symlinkJoin {

        name = "cuda-redist";

        paths = with cudaPackages; [

          (getBin cuda_nvcc)



          (getInclude cuda_cccl) # <nv/target>

          (getInclude cuda_cudart) # cuda_runtime.h

          (getInclude libcublas) # cublas_v2.h

          (getInclude libcusolver) # cusolverDn.h

          (getInclude libcusparse) # cusparse.h

        ];

      };

      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 = (getInclude cudaPackages.libnvshmem).outPath;

      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" ];