From 7aa90c58dc38bb228898fe154c348b578fcb0b7f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 22 Oct 2024 07:55:07 +0000
Subject: [PATCH 1/2] Add `impureWithCuda` dev shell

This shell is handy when developing some kernels jointly with TGI - it
adds nvcc and a bunch of commonly-used CUDA libraries to the environment.

We don't add this to the normal impure shell to keep the development
environment as clean as possible (avoid accidental dependencies, etc.).
---
 flake.nix            |  5 +++++
 nix/impure-shell.nix | 44 +++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/flake.nix b/flake.nix
index edef442f2cb..f26a983ed93 100644
--- a/flake.nix
+++ b/flake.nix
@@ -137,6 +137,11 @@
 
           impure = callPackage ./nix/impure-shell.nix { inherit server; };
 
+          impureWithCuda = callPackage ./nix/impure-shell.nix {
+            inherit server;
+            withCuda = true;
+          };
+
           impure-flash-attn-v1 = callPackage ./nix/impure-shell.nix {
             server = server.override { flash-attn = python3.pkgs.flash-attn-v1; };
           };
diff --git a/nix/impure-shell.nix b/nix/impure-shell.nix
index abed544a155..b0987e70589 100644
--- a/nix/impure-shell.nix
+++ b/nix/impure-shell.nix
@@ -1,7 +1,12 @@
 {
+  lib,
   mkShell,
   black,
+  cmake,
   isort,
+  ninja,
+  which,
+  cudaPackages,
   openssl,
   pkg-config,
   protobuf,
@@ -11,14 +16,17 @@
   ruff,
   rust-bin,
   server,
+
+  # Enable dependencies for building CUDA packages. Useful for e.g.
+  # developing marlin/moe-kernels in-place.
+  withCuda ? false,
 }:
 
 mkShell {
-  buildInputs =
+  nativeBuildInputs =
     [
       black
       isort
-      openssl.dev
       pkg-config
       (rust-bin.stable.latest.default.override {
         extensions = [
@@ -31,6 +39,19 @@ mkShell {
       redocly
       ruff
     ]
+    ++ (lib.optionals withCuda [
+      cmake
+      ninja
+      which
+
+      # For most Torch-based extensions, setting CUDA_HOME is enough, but
+      # some custom CMake builds (e.g. vLLM) also need to have nvcc in PATH.
+      cudaPackages.cuda_nvcc
+    ]);
+  buildInputs =
+    [
+      openssl.dev
+    ]
     ++ (with python3.pkgs; [
       venvShellHook
       docker
@@ -40,10 +61,26 @@ mkShell {
       pytest
       pytest-asyncio
       syrupy
-    ]);
+    ])
+    ++ (lib.optionals withCuda (
+      with cudaPackages;
+      [
+        cuda_cccl
+        cuda_cudart
+        cuda_nvtx
+        libcublas
+        libcusolver
+        libcusparse
+      ]
+    ));
 
   inputsFrom = [ server ];
 
+  env = lib.optionalAttrs withCuda {
+    CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}";
+    TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" python3.pkgs.torch.cudaCapabilities;
+  };
+
   venvDir = "./.venv";
 
   postVenvCreation = ''
@@ -51,6 +88,7 @@ mkShell {
     ( cd server ; python -m pip install --no-dependencies -e . )
     ( cd clients/python ; python -m pip install --no-dependencies -e . )
   '';
+
   postShellHook = ''
     unset SOURCE_DATE_EPOCH
     export PATH=$PATH:~/.cargo/bin

From 6d60aa97b9849debc211c4ba55b7c717e69e8094 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 22 Oct 2024 08:02:48 +0000
Subject: [PATCH 2/2] Add cuDNN

---
 nix/impure-shell.nix | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nix/impure-shell.nix b/nix/impure-shell.nix
index b0987e70589..9df4b111995 100644
--- a/nix/impure-shell.nix
+++ b/nix/impure-shell.nix
@@ -68,6 +68,7 @@ mkShell {
         cuda_cccl
         cuda_cudart
         cuda_nvtx
+        cudnn
         libcublas
         libcusolver
         libcusparse