From 5a37c481f287d102e0549ff077ecb6a273b1dcd0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Wed, 8 Jul 2026 07:33:57 +0000
Subject: [PATCH] Update nixpkgs

---
 flake.lock                                 | 14 +++++++-------
 flake.nix                                  |  5 ++++-
 nix-builder/lib/build-sets.nix             |  2 --
 nix-builder/lib/extension/torch/arch.nix   |  2 +-
 nix-builder/lib/extension/tvm-ffi/arch.nix |  2 +-
 nix-builder/lib/mk-build-set.nix           |  1 -
 nix-builder/overlay.nix                    | 12 ++++++++++++
 7 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/flake.lock b/flake.lock
index bd4e5bfa0..b90b72d81 100644
--- a/flake.lock
+++ b/flake.lock
@@ -35,17 +35,17 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1776927958,
-        "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=",
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "fec2c46cca5bf9767486a290abae51200b656d69",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
         "type": "github"
       },
       "original": {
         "owner": "NixOS",
-        "ref": "nixos-unstable-small",
         "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
         "type": "github"
       }
     },
@@ -64,11 +64,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1776914043,
-        "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=",
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
         "type": "github"
       },
       "original": {
diff --git a/flake.nix b/flake.nix
index aa4b29c57..c9fcf9bc7 100644
--- a/flake.nix
+++ b/flake.nix
@@ -10,7 +10,10 @@
 
   inputs = {
     flake-utils.url = "github:numtide/flake-utils";
-    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable-small";
+    # Temporarily pin this commit to fix ucx issues:
+    # https://github.com/NixOS/nixpkgs/pull/538738
+    nixpkgs.url = "github:NixOS/nixpkgs/ec1a11210589d294f0ac99d3290a27e6c73dfa1d";
+    #nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable-small";
     flake-compat.url = "github:edolstra/flake-compat";
     rust-overlay = {
       url = "github:oxalica/rust-overlay";
diff --git a/nix-builder/lib/build-sets.nix b/nix-builder/lib/build-sets.nix
index 348781357..ed16d1c0f 100644
--- a/nix-builder/lib/build-sets.nix
+++ b/nix-builder/lib/build-sets.nix
@@ -6,8 +6,6 @@
 let
   inherit (nixpkgs) lib;
 
-  overlay = import ../nix/overlay.nix;
-
   inherit (import ./torch-version-utils.nix { inherit lib; })
     backend
     flattenSystems
diff --git a/nix-builder/lib/extension/torch/arch.nix b/nix-builder/lib/extension/torch/arch.nix
index d23655916..13dc9277d 100644
--- a/nix-builder/lib/extension/torch/arch.nix
+++ b/nix-builder/lib/extension/torch/arch.nix
@@ -198,7 +198,7 @@ stdenv.mkDerivation (prevAttrs: {
 
       # Make dependent on build configuration dependencies once
       # the Torch dependency is gone.
-      cuda_cccl
+      cccl
       libcublas
       libcusolver
       libcusparse
diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix
index ffc877cd9..fdadd9024 100644
--- a/nix-builder/lib/extension/tvm-ffi/arch.nix
+++ b/nix-builder/lib/extension/tvm-ffi/arch.nix
@@ -194,7 +194,7 @@ stdenv.mkDerivation (prevAttrs: {
 
       # Make dependent on build configuration dependencies once
       # the Torch dependency is gone.
-      cuda_cccl
+      cccl
       libcublas
       libcusolver
       libcusparse
diff --git a/nix-builder/lib/mk-build-set.nix b/nix-builder/lib/mk-build-set.nix
index dd8d54254..550c39b51 100644
--- a/nix-builder/lib/mk-build-set.nix
+++ b/nix-builder/lib/mk-build-set.nix
@@ -48,7 +48,6 @@ let
             });
           }
         );
-
   };
 
   overlayForRocmVersion = rocmVersion: self: super: {
diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix
index 9fab8d225..7b727b03f 100644
--- a/nix-builder/overlay.nix
+++ b/nix-builder/overlay.nix
@@ -1,5 +1,6 @@
 final: prev:
 {
+  # Local packages/hooks.
   kernel-builder = final.callPackage ./pkgs/kernel-builder { };
 
   cmakeNvccThreadsHook = final.callPackage ./pkgs/cmake-nvcc-threads-hook { };
@@ -37,6 +38,17 @@ final: prev:
     }
   );
 
+  # nixpkgs overrides
+  ucc = prev.ucc.overrideAttrs (
+    _: prevAttrs: {
+      # Concurrent build flags do not seem to get passed through? So instead, let's
+      # add parallelism to building CUDA files.
+      preConfigure = prevAttrs.preConfigure + ''
+        export NVCC_CFLAGS="--threads $NIX_BUILD_CORES"
+      '';
+    }
+  );
+
   ucx = prev.ucx.overrideAttrs (
     _: prevAttrs: {
       buildInputs = prevAttrs.buildInputs ++ [ final.cudaPackages.cuda_nvcc ];