···

       
15
       
 
       
  ninja,

     

       
16
       
 
       
  numpy,

     

       
17
       
 
       
  torch,

     

       
0
       
 
       

     

       
0
       
 
       

     

       
18
       
 
       
}:

     

       
19
       
 
       


     

       
20
       
-
       
let

     

       
21
       
 
       
  pname = "flashinfer";

     

       
22
       
-
       
  version = "0.2.9";

     

       
23
       
-
       


     

       
24
       
-
       
  src_cutlass = fetchFromGitHub {

     

       
25
       
-
       
    owner = "NVIDIA";

     

       
26
       
-
       
    repo = "cutlass";

     

       
27
       
-
       
    # Using the revision obtained in submodule inside flashinfer's `3rdparty`.

     

       
28
       
-
       
    tag = "v${version}";

     

       
29
       
-
       
    hash = "sha256-d4czDoEv0Focf1bJHOVGX4BDS/h5O7RPoM/RrujhgFQ=";

     

       
30
       
-
       
  };

     

       
31
       
-
       


     

       
32
       
-
       
in

     

       
33
       
-
       
buildPythonPackage {

     

       
34
       
-
       
  format = "setuptools";

     

       
35
       
-
       
  inherit pname version;

     

       
36
       
 
       


     

       
37
       
 
       
  src = fetchFromGitHub {

     

       
38
       
 
       
    owner = "flashinfer-ai";

     

       
39
       
 
       
    repo = "flashinfer";

     

       
40
       
 
       
    tag = "v${version}";

     

       
41
       
-
       
    hash = "sha256-M0q6d+EpuTehbw68AQ73Fhwmw2tzjymYjSXaol9QC7Y=";

     

       
0
       
 
       

     

       
42
       
 
       
  };

     

       
43
       
 
       


     

       
44
       
 
       
  build-system = [ setuptools ];

     
···

       
48
       
 
       
    ninja

     

       
49
       
 
       
    (lib.getBin cudaPackages.cuda_nvcc)

     

       
50
       
 
       
  ];

     

       
0
       
 
       

     

       
51
       
 
       
  dontUseCmakeConfigure = true;

     

       
52
       
 
       


     

       
53
       
-
       
  buildInputs = [

     

       
54
       
-
       
    cudaPackages.cuda_cudart

     

       
55
       
-
       
    cudaPackages.libcublas

     

       
56
       
-
       
    cudaPackages.cuda_cccl

     

       
57
       
-
       
    cudaPackages.libcurand

     

       
58
       
 
       
  ];

     

       
59
       
-
       


     

       
60
       
-
       
  postPatch = ''

     

       
61
       
-
       
    rmdir 3rdparty/cutlass

     

       
62
       
-
       
    ln -s ${src_cutlass} 3rdparty/cutlass

     

       
63
       
-
       
  '';

     

       
64
       
 
       


     

       
65
       
 
       
  # FlashInfer offers two installation modes:

     

       
66
       
 
       
  #

     
···

       
86
       
 
       
  dependencies = [

     

       
87
       
 
       
    numpy

     

       
88
       
 
       
    torch

     

       
0
       
 
       

     

       
0
       
 
       

     

       
89
       
 
       
  ];

     

       
90
       
 
       


     

       
91
       
 
       
  meta = with lib; {

     

···

       
15
       
 
       
  ninja,

     

       
16
       
 
       
  numpy,

     

       
17
       
 
       
  torch,

     

       
18
       
+
       
  pynvml,

     

       
19
       
+
       
  einops,

     

       
20
       
 
       
}:

     

       
21
       
 
       


     

       
22
       
+
       
buildPythonPackage rec {

     

       
23
       
 
       
  pname = "flashinfer";

     

       
24
       
+
       
  version = "0.2.14";

     

       
25
       
+
       
  pyproject = true;

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
26
       
 
       


     

       
27
       
 
       
  src = fetchFromGitHub {

     

       
28
       
 
       
    owner = "flashinfer-ai";

     

       
29
       
 
       
    repo = "flashinfer";

     

       
30
       
 
       
    tag = "v${version}";

     

       
31
       
+
       
    hash = "sha256-MZiZwdedz+Vxa1+VBfHDKf4NVSiOAytGboIJ0DvCXmk=";

     

       
32
       
+
       
    fetchSubmodules = true;

     

       
33
       
 
       
  };

     

       
34
       
 
       


     

       
35
       
 
       
  build-system = [ setuptools ];

     
···

       
39
       
 
       
    ninja

     

       
40
       
 
       
    (lib.getBin cudaPackages.cuda_nvcc)

     

       
41
       
 
       
  ];

     

       
42
       
+
       


     

       
43
       
 
       
  dontUseCmakeConfigure = true;

     

       
44
       
 
       


     

       
45
       
+
       
  buildInputs = with cudaPackages; [

     

       
46
       
+
       
    cuda_cccl

     

       
47
       
+
       
    cuda_cudart

     

       
48
       
+
       
    libcublas

     

       
49
       
+
       
    libcurand

     

       
50
       
 
       
  ];

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
51
       
 
       


     

       
52
       
 
       
  # FlashInfer offers two installation modes:

     

       
53
       
 
       
  #

     
···

       
73
       
 
       
  dependencies = [

     

       
74
       
 
       
    numpy

     

       
75
       
 
       
    torch

     

       
76
       
+
       
    pynvml

     

       
77
       
+
       
    einops

     

       
78
       
 
       
  ];

     

       
79
       
 
       


     

       
80
       
 
       
  meta = with lib; {