vllm-cpp-rs/flake.nix at main · querymt/vllm-cpp-rs · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
{
  description = "Rust bindings for vllm.cpp";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    flake-parts.url = "github:hercules-ci/flake-parts";
    rust-overlay.url = "github:oxalica/rust-overlay";
  };

  outputs = inputs:
    inputs.flake-parts.lib.mkFlake {inherit inputs;} {
      systems = inputs.nixpkgs.lib.systems.flakeExposed;

      perSystem = {system, ...}: let
        overlays = [inputs.rust-overlay.overlays.default];
        pkgs = import inputs.nixpkgs {
          inherit system overlays;
        };
        cudaPkgs = import inputs.nixpkgs {
          inherit system overlays;
          config.allowUnfree = true;
        };

        rustToolchain = pkgs.rust-bin.fromRustupToolchainFile ./rust-toolchain.toml;
        msrvToolchain = pkgs.rust-bin.stable."1.85.0".default.override {
          extensions = ["clippy" "rustfmt"];
        };
      in {
        devShells =
          {
            default = pkgs.mkShell {
              packages = [
                rustToolchain
                pkgs.cmake
                pkgs.git
                pkgs.just
                pkgs.jq
                pkgs.ninja
                pkgs.pkg-config
                pkgs.gnutar
                pkgs.llvmPackages.clang
                pkgs.llvmPackages.bintools
                pkgs.rust-bindgen
              ];

              shellHook = ''
                export PS1="(dev:vllm-cpp-rs) $PS1"
              '';
            };

            msrv = pkgs.mkShell {
              packages = [
                msrvToolchain
                pkgs.cmake
                pkgs.just
                pkgs.ninja
                pkgs.pkg-config
                pkgs.llvmPackages.clang
                pkgs.llvmPackages.bintools
              ];

              shellHook = ''
                export PS1="(msrv:vllm-cpp-rs) $PS1"
              '';
            };
          }
          // pkgs.lib.optionalAttrs
          (builtins.elem system [
            "x86_64-linux"
            "aarch64-linux"
          ]) {
            cuda = let
              toolkit = cudaPkgs.cudaPackages.cudatoolkit;
              cutlass = cudaPkgs.cudaPackages.cutlass;
            in
              pkgs.mkShell {
                packages = [
                  rustToolchain
                  pkgs.cmake
                  pkgs.git
                  pkgs.just
                  pkgs.jq
                  pkgs.ninja
                  pkgs.pkg-config
                  pkgs.gnutar
                  pkgs.llvmPackages.clang
                  pkgs.llvmPackages.bintools
                  pkgs.rust-bindgen
                  toolkit
                  cutlass
                ];

                shellHook = ''
                  export PS1="(cuda:vllm-cpp-rs) $PS1"
                  export CUDA_PATH="${toolkit}"
                  export CUDA_HOME="$CUDA_PATH"
                  export CUDAToolkit_ROOT="$CUDA_PATH"
                  export VLLM_CPP_CUTLASS_DIR="${cutlass.src}"
                  if [ -d /run/opengl-driver/lib ]; then
                    export LD_LIBRARY_PATH="/run/opengl-driver/lib:${toolkit}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
                  else
                    export LD_LIBRARY_PATH="${toolkit}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
                  fi
                '';
              };

            vulkan = let
              vulkanLibraryPath = pkgs.lib.makeLibraryPath [
                pkgs.vulkan-loader
                pkgs.vulkan-validation-layers
              ];
            in
              pkgs.mkShell {
                packages = [
                  rustToolchain
                  pkgs.cmake
                  pkgs.git
                  pkgs.just
                  pkgs.jq
                  pkgs.ninja
                  pkgs.pkg-config
                  pkgs.gnutar
                  pkgs.llvmPackages.clang
                  pkgs.llvmPackages.bintools
                  pkgs.rust-bindgen
                  pkgs.vulkan-tools
                  pkgs.vulkan-loader
                  pkgs.vulkan-validation-layers
                ];

                shellHook = ''
                  export PS1="(vulkan:vllm-cpp-rs) $PS1"
                  if [ -d /run/opengl-driver/lib ]; then
                    export LD_LIBRARY_PATH="/run/opengl-driver/lib:${vulkanLibraryPath}''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
                  else
                    export LD_LIBRARY_PATH="${vulkanLibraryPath}''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
                  fi
                  export VK_LAYER_PATH="${pkgs.vulkan-validation-layers}/share/vulkan/explicit_layer.d"
                '';
              };
          };
      };
    };
}