0
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2025-07-14 06:00:33 +03:00

A few updates

This commit is contained in:
Break Yang 2025-05-20 15:29:12 -07:00
parent 4875dd88d8
commit d40d3b0fe7

View file

@ -6,10 +6,11 @@
# environment varaible `CUDA_HOME` to `cudatoolkit`.
{
lib,
config,
buildPythonPackage,
symlinkJoin,
fetchFromGitHub,
setuptools,
cudaPackages,
cmake,
ninja,
numpy,
@ -20,11 +21,6 @@ let
pname = "flashinfer";
version = "0.2.5";
inherit (torch) cudaPackages;
inherit (cudaPackages) cudaMajorMinorVersion;
cudaMajorMinorVersionString = lib.replaceStrings [ "." ] [ "" ] cudaMajorMinorVersion;
src_cutlass = fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
@ -49,7 +45,7 @@ buildPythonPackage {
nativeBuildInputs = [
cmake
ninja
cudaPackages.cudatoolkit
(lib.getBin cudaPackages.cuda_nvcc)
];
dontUseCmakeConfigure = true;
@ -76,7 +72,6 @@ buildPythonPackage {
export TORCH_NVCC_FLAGS="--maxrregcount=64"
'';
CUDA_HOME = "${cudaPackages.cudatoolkit}";
TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
dependencies = [