Migrate the rest of hf-nix to kernel-builder

Author: danieldk

flake.lock

@@ -3,17 +3,17 @@
 
   inputs = {
     flake-utils.url = "github:numtide/flake-utils";
-    nixpkgs.follows = "hf-nix/nixpkgs";
+    # Put back to nixos-unstable-small the next bump. Exact revision is
+    # to avoid a rebuild during the hf-nix -> kernel-builder transition.
+    nixpkgs.url = "github:NixOS/nixpkgs/c543a59edf25ada193719764f3bc0c6ba835f94d";
     flake-compat.url = "github:edolstra/flake-compat";
-    hf-nix.url = "github:huggingface/hf-nix";
   };
 
   outputs =
     {
       self,
       flake-compat,
       flake-utils,
-      hf-nix,
       nixpkgs,
     }:
     let
@@ -33,7 +33,6 @@
             name = system;
             value = import ./lib/build-sets.nix {
               inherit nixpkgs system torchVersions;
-              hf-nix = hf-nix.overlays.default;
             };
           }) systems
         );

flake.nix

@@ -1,7 +1,6 @@
 {
   nixpkgs,
   system,
-  hf-nix,
   torchVersions,
 }:
 
@@ -121,7 +120,6 @@ let
             xpuSupport = true;
           };
           overlays = [
-            hf-nix
             overlay
             (overlayForXpuVersion xpuVersion)
           ];
@@ -138,7 +136,6 @@ let
       allowUnfree = true;
     };
     overlays = [
-      hf-nix
       overlay
     ];
   };
@@ -157,7 +154,6 @@ let
             cudaSupport = true;
           };
           overlays = [
-            hf-nix
             overlay
             (overlayForCudaVersion cudaVersion)
           ];
@@ -179,7 +175,6 @@ let
             rocmSupport = true;
           };
           overlays = [
-            hf-nix
             overlay
             (overlayForRocmVersion rocmVersion)
           ];

lib/build-sets.nix

@@ -1,22 +1,58 @@
-final: prev: {
-  cmakeNvccThreadsHook = prev.callPackage ./pkgs/cmake-nvcc-threads-hook { };
-
-  # Local packages
+final: prev:
+let
+  # For XPU we use MKL from the joined oneAPI toolkit.
+  useMKL = final.stdenv.isx86_64 && !(final.config.xpuSupport or false);
+in
+{
+  # Use MKL for BLAS/LAPACK on x86_64.
+  blas = if useMKL then prev.blas.override { blasProvider = prev.mkl; } else prev.blas;
+  lapack = if useMKL then prev.lapack.override { lapackProvider = prev.mkl; } else prev.blas;
 
   build2cmake = prev.callPackage ./pkgs/build2cmake { };
 
+  cmakeNvccThreadsHook = prev.callPackage ./pkgs/cmake-nvcc-threads-hook { };
+
   get-kernel-check = prev.callPackage ./pkgs/get-kernel-check { };
 
   kernel-abi-check = prev.callPackage ./pkgs/kernel-abi-check { };
 
   kernel-layout-check = prev.callPackage ./pkgs/kernel-layout-check { };
 
+  # Used by ROCm.
+  libffi_3_2 = final.libffi_3_3.overrideAttrs (
+    finalAttrs: _: {
+      version = "3.2.1";
+      src = final.fetchurl {
+        url = with finalAttrs; "https://gcc.gnu.org/pub/${pname}/${pname}-${version}.tar.gz";
+        hash = "sha256-0G67jh2aItGeONY/24OVQlPzm+3F1GIyoFZFaFciyjc=";
+      };
+    }
+  );
+
+  magma = (prev.callPackage ./pkgs/magma { }).magma;
+
+  magma-hip =
+    (prev.callPackage ./pkgs/magma {
+      cudaSupport = false;
+      rocmSupport = true;
+    }).magma;
+
+  nvtx = final.callPackage ./pkgs/nvtx { };
+
+  metal-cpp = final.callPackage ./pkgs/metal-cpp { };
+
   rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
 
   remove-bytecode-hook = prev.callPackage ./pkgs/remove-bytecode-hook { };
 
   stdenvGlibc_2_27 = prev.callPackage ./pkgs/stdenv-glibc-2_27 { };
 
+  ucx = prev.ucx.overrideAttrs (
+    _: prevAttrs: {
+      buildInputs = prevAttrs.buildInputs ++ [ final.cudaPackages.cuda_nvcc ];
+    }
+  );
+
   # Python packages
   pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
     (
@@ -53,6 +89,14 @@ final: prev: {
 
         mkTorch = callPackage ./pkgs/python-modules/torch/binary { };
 
+        scipy = python-super.scipy.overrideAttrs (
+          _: prevAttrs: {
+            # Three tests have a slight deviance.
+            doCheck = false;
+            doInstallCheck = false;
+          }
+        );
+
         torch-bin_2_8 = mkTorch {
           version = "2.8";
           xpuPackages = final.xpuPackages_2025_1;
@@ -70,7 +114,60 @@ final: prev: {
         torch_2_9 = callPackage ./pkgs/python-modules/torch/source/2_9 {
           xpuPackages = final.xpuPackages_2025_2;
         };
+
+        triton-xpu_2_8 = callPackage ./pkgs/python-modules/triton-xpu {
+          torchVersion = "2.8";
+          xpuPackages = final.xpuPackages_2025_1;
+        };
+
+        triton-xpu_2_9 = callPackage ./pkgs/python-modules/triton-xpu {
+          torchVersion = "2.9";
+          xpuPackages = final.xpuPackages_2025_2;
+        };
       }
     )
+    (import ./pkgs/python-modules/hooks)
   ];
+
+  xpuPackages = final.xpuPackages_2025_1;
 }
+// (import ./pkgs/cutlass { pkgs = final; })
+// (
+  let
+    flattenVersion = prev.lib.strings.replaceStrings [ "." ] [ "_" ];
+    readPackageMetadata = path: (builtins.fromJSON (builtins.readFile path));
+    versions = [
+      "6.3.4"
+      "6.4.2"
+      "7.0.1"
+    ];
+    newRocmPackages = final.callPackage ./pkgs/rocm-packages { };
+  in
+  builtins.listToAttrs (
+    map (version: {
+      name = "rocmPackages_${flattenVersion (prev.lib.versions.majorMinor version)}";
+      value = newRocmPackages {
+        packageMetadata = readPackageMetadata ./pkgs/rocm-packages/rocm-${version}-metadata.json;
+      };
+    }) versions
+  )
+)
+// (
+  let
+    flattenVersion = prev.lib.strings.replaceStrings [ "." ] [ "_" ];
+    readPackageMetadata = path: (builtins.fromJSON (builtins.readFile path));
+    xpuVersions = [
+      "2025.1.3"
+      "2025.2.1"
+    ];
+    newXpuPackages = final.callPackage ./pkgs/xpu-packages { };
+  in
+  builtins.listToAttrs (
+    map (version: {
+      name = "xpuPackages_${flattenVersion (prev.lib.versions.majorMinor version)}";
+      value = newXpuPackages {
+        packageMetadata = readPackageMetadata ./pkgs/xpu-packages/intel-deep-learning-${version}.json;
+      };
+    }) xpuVersions
+  )
+)