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:
parent
4875dd88d8
commit
d40d3b0fe7
1 changed files with 3 additions and 8 deletions
|
@ -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 = [
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue