From 622c9c367a6705dfd934b39bb8e4a2b10c52dac3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Thu, 29 Aug 2024 16:25:25 +0200
Subject: [PATCH] nix: build Torch against MKL and various other improvements
 (#2469)

Updates tgi-nix input:

- Move Torch closer to upstream by building against MKL.
- Remove compute capability 8.7 from Torch (Jetson).
- Sync nixpkgs cumpute capabilities with Torch (avoids
  compiling too mana capabilities for MAGMA).
- Use nixpkgs configuration passed through by `tgi-nix`.
---
 flake.lock     | 6 +++---
 flake.nix      | 9 +++------
 nix/server.nix | 1 +
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/flake.lock b/flake.lock
index b40f51b3..14011768 100644
--- a/flake.lock
+++ b/flake.lock
@@ -944,11 +944,11 @@
         "nixpkgs": "nixpkgs_6"
       },
       "locked": {
-        "lastModified": 1724270760,
-        "narHash": "sha256-KX566x0+3HZcB20HPdvdwyMm7ZJg21M+iqVrs/HCimA=",
+        "lastModified": 1724784743,
+        "narHash": "sha256-NdEoWeNwR/ZstYnHaiQWIYZvr7VsrAh7g3+ZHUPrxuI=",
         "owner": "danieldk",
         "repo": "tgi-nix",
-        "rev": "12cbaa76ff258351741d3b5afb7161f617fe7b4c",
+        "rev": "c9580c3e39a855246bb87b584bbea1885b44f524",
         "type": "github"
       },
       "original": {
diff --git a/flake.nix b/flake.nix
index 83feb26a..0739c90a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -31,15 +31,12 @@
           src = ./.;
           additionalCargoNixArgs = [ "--all-features" ];
         };
-        config = {
-          allowUnfree = true;
-          cudaSupport = true;
-        };
         pkgs = import nixpkgs {
-          inherit config system;
+          inherit system;
+          inherit (tgi-nix.lib) config;
           overlays = [
             rust-overlay.overlays.default
-            tgi-nix.overlay
+            tgi-nix.overlays.default
           ];
         };
         crateOverrides = import ./nix/crate-overrides.nix { inherit pkgs nix-filter; };
diff --git a/nix/server.nix b/nix/server.nix
index 4e0fdaa1..6ee088e0 100644
--- a/nix/server.nix
+++ b/nix/server.nix
@@ -28,6 +28,7 @@
   peft,
   safetensors,
   tokenizers,
+  torch,
   sentencepiece,
   transformers,
   typer,