···

       
9
       
 
       
  config,

     

       
10
       
 
       
  buildPythonPackage,

     

       
11
       
 
       
  fetchFromGitHub,

     

       
0
       
 
       

     

       
0
       
 
       

     

       
12
       
 
       
  setuptools,

     

       
13
       
-
       
  cudaPackages,

     

       
0
       
 
       

     

       
14
       
 
       
  cmake,

     

       
15
       
 
       
  ninja,

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
16
       
 
       
  numpy,

     

       
0
       
 
       

     

       
0
       
 
       

     

       
17
       
 
       
  torch,

     

       
18
       
-
       
  pynvml,

     

       
19
       
-
       
  einops,

     

       
20
       
 
       
}:

     

       
21
       
 
       


     

       
22
       
 
       
buildPythonPackage rec {

     

       
23
       
 
       
  pname = "flashinfer";

     

       
24
       
-
       
  version = "0.2.14";

     

       
25
       
 
       
  pyproject = true;

     

       
26
       
 
       


     

       
27
       
 
       
  src = fetchFromGitHub {

     

       
28
       
 
       
    owner = "flashinfer-ai";

     

       
29
       
 
       
    repo = "flashinfer";

     

       
30
       
 
       
    tag = "v${version}";

     

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

     

       
32
       
 
       
    fetchSubmodules = true;

     

       
0
       
 
       

     

       
33
       
 
       
  };

     

       
34
       
 
       


     

       
35
       
 
       
  build-system = [ setuptools ];

     
···

       
68
       
 
       
    export MAX_JOBS="$NIX_BUILD_CORES"

     

       
69
       
 
       
  '';

     

       
70
       
 
       


     

       
71
       
-
       
  TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;

     

       
72
       
 
       


     

       
0
       
 
       

     

       
0
       
 
       

     

       
0
       
 
       

     

       
73
       
 
       
  dependencies = [

     

       
0
       
 
       

     

       
0
       
 
       

     

       
74
       
 
       
    numpy

     

       
75
       
-
       
    torch

     

       
76
       
 
       
    pynvml

     

       
77
       
-
       
    einops

     

       
0
       
 
       

     

       
0
       
 
       

     

       
78
       
 
       
  ];

     

       
79
       
 
       


     

       
80
       
-
       
  meta = with lib; {

     

       
81
       
 
       
    broken = !torch.cudaSupport || !config.cudaSupport;

     

       
82
       
 
       
    homepage = "https://flashinfer.ai/";

     

       
83
       
 
       
    description = "Library and kernel generator for Large Language Models";

     
···

       
88
       
 
       
      and inference, and delivers state-of-the-art performance across diverse

     

       
89
       
 
       
      scenarios.

     

       
90
       
 
       
    '';

     

       
91
       
-
       
    license = licenses.asl20;

     

       
92
       
-
       
    maintainers = with maintainers; [ breakds ];

     

       
93
       
 
       
  };

     

       
94
       
 
       
}

     

···

       
9
       
 
       
  config,

     

       
10
       
 
       
  buildPythonPackage,

     

       
11
       
 
       
  fetchFromGitHub,

     

       
12
       
+
       


     

       
13
       
+
       
  # build-system

     

       
14
       
 
       
  setuptools,

     

       
15
       
+
       


     

       
16
       
+
       
  # nativeBuildInputs

     

       
17
       
 
       
  cmake,

     

       
18
       
 
       
  ninja,

     

       
19
       
+
       
  cudaPackages,

     

       
20
       
+
       


     

       
21
       
+
       
  # dependencies

     

       
22
       
+
       
  click,

     

       
23
       
+
       
  einops,

     

       
24
       
 
       
  numpy,

     

       
25
       
+
       
  pynvml,

     

       
26
       
+
       
  tabulate,

     

       
27
       
 
       
  torch,

     

       
28
       
+
       
  tqdm,

     

       
0
       
 
       

     

       
29
       
 
       
}:

     

       
30
       
 
       


     

       
31
       
 
       
buildPythonPackage rec {

     

       
32
       
 
       
  pname = "flashinfer";

     

       
33
       
+
       
  version = "0.3.1";

     

       
34
       
 
       
  pyproject = true;

     

       
35
       
 
       


     

       
36
       
 
       
  src = fetchFromGitHub {

     

       
37
       
 
       
    owner = "flashinfer-ai";

     

       
38
       
 
       
    repo = "flashinfer";

     

       
39
       
 
       
    tag = "v${version}";

     

       
0
       
 
       

     

       
40
       
 
       
    fetchSubmodules = true;

     

       
41
       
+
       
    hash = "sha256-e9PfLfU0DdoLKlXiHylCbGd125c7Iw9y4NDIOAP0xHs=";

     

       
42
       
 
       
  };

     

       
43
       
 
       


     

       
44
       
 
       
  build-system = [ setuptools ];

     
···

       
77
       
 
       
    export MAX_JOBS="$NIX_BUILD_CORES"

     

       
78
       
 
       
  '';

     

       
79
       
 
       


     

       
80
       
+
       
  FLASHINFER_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;

     

       
81
       
 
       


     

       
82
       
+
       
  pythonRemoveDeps = [

     

       
83
       
+
       
    "nvidia-cudnn-frontend"

     

       
84
       
+
       
  ];

     

       
85
       
 
       
  dependencies = [

     

       
86
       
+
       
    click

     

       
87
       
+
       
    einops

     

       
88
       
 
       
    numpy

     

       
0
       
 
       

     

       
89
       
 
       
    pynvml

     

       
90
       
+
       
    tabulate

     

       
91
       
+
       
    torch

     

       
92
       
+
       
    tqdm

     

       
93
       
 
       
  ];

     

       
94
       
 
       


     

       
95
       
+
       
  meta = {

     

       
96
       
 
       
    broken = !torch.cudaSupport || !config.cudaSupport;

     

       
97
       
 
       
    homepage = "https://flashinfer.ai/";

     

       
98
       
 
       
    description = "Library and kernel generator for Large Language Models";

     
···

       
103
       
 
       
      and inference, and delivers state-of-the-art performance across diverse

     

       
104
       
 
       
      scenarios.

     

       
105
       
 
       
    '';

     

       
106
       
+
       
    license = lib.licenses.asl20;

     

       
107
       
+
       
    maintainers = with lib.maintainers; [ breakds ];

     

       
108
       
 
       
  };

     

       
109
       
 
       
}