From c47302c9517d3d949bd6a3fbdfbda0eca3095eac Mon Sep 17 00:00:00 2001
From: polwex <polwex@sortug.com>
Date: Tue, 14 Oct 2025 06:36:38 +0700
Subject: m

---
 derivations/flash-attn/nvidia.nix | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)
 create mode 100644 derivations/flash-attn/nvidia.nix

(limited to 'derivations/flash-attn/nvidia.nix')

diff --git a/derivations/flash-attn/nvidia.nix b/derivations/flash-attn/nvidia.nix
new file mode 100644
index 0000000..7b83e84
--- /dev/null
+++ b/derivations/flash-attn/nvidia.nix
@@ -0,0 +1,24 @@
+{
+  symlinkJoin,
+  cudaPackages,
+  pkgs,
+  cudaCapabilities ? pkgs.cudaPackages.flags.cudaCapabilities,
+  lib,
+}: {
+  BUILD_CUDA_EXT = "1";
+
+  CUDA_HOME = symlinkJoin {
+    name = "cuda-redist";
+    paths = with cudaPackages; [
+      cuda_cudart # cuda_runtime.h
+      cuda_nvcc
+    ];
+  };
+
+  CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
+
+  preBuild = ''
+    export PATH=${pkgs.gcc13Stdenv.cc}/bin:$PATH
+    export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
+  '';
+}
-- 
cgit v1.2.3