From 6e46f9188422f358fa76ab742d610e2055feba61 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 23 Apr 2026 14:41:19 +0000
Subject: [PATCH] python3Packages.deep-ep: cleanup, fix build on CUDA>=13.0

---
 .../python-modules/deep-ep/default.nix        | 37 +++++++++----------
 1 file changed, 17 insertions(+), 20 deletions(-)

diff --git a/pkgs/development/python-modules/deep-ep/default.nix b/pkgs/development/python-modules/deep-ep/default.nix
index 38309a06f988..0ad6abb30b53 100644
--- a/pkgs/development/python-modules/deep-ep/default.nix
+++ b/pkgs/development/python-modules/deep-ep/default.nix
@@ -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" ];