Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 4 additions & 76 deletions flake.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 3 additions & 4 deletions flake.nix
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -33,7 +33,6 @@
name = system;
value = import ./lib/build-sets.nix {
inherit nixpkgs system torchVersions;
hf-nix = hf-nix.overlays.default;
};
}) systems
);
Expand Down
5 changes: 0 additions & 5 deletions lib/build-sets.nix
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
{
nixpkgs,
system,
hf-nix,
torchVersions,
}:

Expand Down Expand Up @@ -121,7 +120,6 @@ let
xpuSupport = true;
};
overlays = [
hf-nix
overlay
(overlayForXpuVersion xpuVersion)
];
Expand All @@ -138,7 +136,6 @@ let
allowUnfree = true;
};
overlays = [
hf-nix
overlay
];
};
Expand All @@ -157,7 +154,6 @@ let
cudaSupport = true;
};
overlays = [
hf-nix
overlay
(overlayForCudaVersion cudaVersion)
];
Expand All @@ -179,7 +175,6 @@ let
rocmSupport = true;
};
overlays = [
hf-nix
overlay
(overlayForRocmVersion rocmVersion)
];
Expand Down
105 changes: 101 additions & 4 deletions overlay.nix
Original file line number Diff line number Diff line change
@@ -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 ++ [
(
Expand Down Expand Up @@ -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;
Expand All @@ -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
)
)
Loading
Loading