python3Packages.exllamav2: init at 0.3.2; python3Packages.exllamav3: init at 0.0.20 (#481121)

2026-01-20 21:45:50 +00:00
parent 1b8a75653d 75b6cf2082
commit 63ff7e6fd5
3 changed files with 213 additions and 0 deletions
@@ -0,0 +1,107 @@
+{
+  fetchFromGitHub,
+  lib,
+  cudaPackages,
+  buildPythonPackage,
+  nix-update-script,
+
+  setuptools,
+
+  pybind11,
+
+  which,
+
+  fastparquet,
+  flash-attn,
+  ninja,
+  numpy,
+  pandas,
+  pillow,
+  pygments,
+  regex,
+  rich,
+  safetensors,
+  tokenizers,
+  torch,
+  websockets,
+}:
+buildPythonPackage (finalAttrs: {
+  pname = "exllamav2";
+  version = "0.3.2";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "turboderp-org";
+    repo = "exllamav2";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-WbpbANenOuy6F0qAKVKAmolHjgRKfPxSVud8FZG1TXw=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  buildInputs = [
+    pybind11
+  ]
+  ++ lib.optionals torch.cudaSupport [
+    cudaPackages.libcusparse # cusparse.h
+    cudaPackages.libcublas # cublas_v2.h
+    cudaPackages.libcusolver # cusolverDn.h
+    cudaPackages.libcurand # curand_kernel.h
+    cudaPackages.cuda_cudart # cuda_runtime.h
+  ];
+
+  env = lib.optionalAttrs torch.cudaSupport {
+    CUDA_HOME = lib.getDev cudaPackages.cuda_nvcc;
+    TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" torch.cudaCapabilities;
+  };
+
+  nativeBuildInputs = [
+    ninja
+    which
+  ];
+
+  pythonRelaxDeps = [
+    "numpy" # Wants numpy 1.26.4
+  ];
+
+  dependencies = [
+    fastparquet
+    flash-attn
+    ninja
+    numpy
+    pandas
+    pillow
+    pygments
+    regex
+    rich
+    safetensors
+    tokenizers
+    torch
+    websockets
+  ];
+
+  pythonImportsCheck = [ "exllamav2" ];
+
+  # Tests require GPU hardware and external model files
+  doCheck = false;
+
+  passthru.updateScript = nix-update-script { };
+
+  meta = {
+    homepage = "https://github.com/turboderp-org/exllamav2";
+    description = "Inference library for running LLMs locally on modern consumer-class GPUs";
+    changelog = "https://github.com/turboderp-org/exllamav2/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.mit;
+    platforms = [
+      "x86_64-windows"
+      "x86_64-linux"
+    ];
+
+    # Package requires CUDA or ROCm for functionality
+    # ROCm support is partially implemented but untested
+    broken = !torch.cudaSupport;
+    maintainers = with lib.maintainers; [ BatteredBunny ];
+  };
+})
@@ -0,0 +1,102 @@
+{
+  lib,
+  fetchFromGitHub,
+  buildPythonPackage,
+  cudaPackages,
+  nix-update-script,
+
+  setuptools,
+
+  flash-attn,
+  formatron,
+  kbnf,
+  marisa-trie,
+  ninja,
+  numpy,
+  pillow,
+  pydantic,
+  pyyaml,
+  rich,
+  safetensors,
+  tokenizers,
+  torch,
+  typing-extensions,
+}:
+let
+  newerThanTuring = lib.filter (version: lib.versionOlder "7.9" version) torch.cudaCapabilities;
+in
+buildPythonPackage (finalAttrs: {
+  pname = "exllamav3";
+  version = "0.0.20";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "turboderp-org";
+    repo = "exllamav3";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-G3PtxKU/J4JEQQOwFmrSWuSr/hA4uyxRci3khXCwEqE=";
+  };
+
+  pythonRelaxDeps = [
+    "pydantic"
+  ];
+
+  build-system = [
+    setuptools
+  ];
+
+  nativeBuildInputs = [
+    ninja
+  ];
+
+  buildInputs = lib.optionals torch.cudaSupport [
+    cudaPackages.cuda_cudart # cuda_runtime.h
+    cudaPackages.libcusparse # cusparse.h
+    cudaPackages.libcublas # cublas_v2.h
+    cudaPackages.libcusolver # cusolverDn.h
+    cudaPackages.libcurand # curand_kernel.h
+  ];
+
+  dependencies = [
+    flash-attn
+    formatron
+    kbnf
+    marisa-trie
+    numpy
+    pillow
+    pydantic
+    pyyaml
+    rich
+    safetensors
+    tokenizers
+    torch
+    typing-extensions
+  ];
+
+  env = lib.optionalAttrs torch.cudaSupport {
+    CUDA_HOME = lib.getDev cudaPackages.cuda_nvcc;
+    # exllamav3 only supports turing or newer GPUs
+    # https://github.com/turboderp-org/exllamav3/issues/44
+    TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" newerThanTuring;
+  };
+
+  pythonImportsCheck = [ "exllamav3" ];
+
+  # Tests require GPU hardware and external model files
+  doCheck = false;
+
+  passthru.updateScript = nix-update-script { };
+
+  meta = {
+    description = "Quantization and inference library for running LLMs locally on modern consumer-class GPUs";
+    homepage = "https://github.com/turboderp-org/exllamav3";
+    changelog = "https://github.com/turboderp-org/exllamav3/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.mit;
+    platforms = [
+      "x86_64-windows"
+      "x86_64-linux"
+    ];
+    broken = !torch.cudaSupport; # Package requires CUDA for functionality
+    maintainers = with lib.maintainers; [ BatteredBunny ];
+  };
+})
@@ -5121,6 +5121,10 @@ self: super: with self; {

  exiv2 = callPackage ../development/python-modules/exiv2 { inherit (pkgs) exiv2; };

+  exllamav2 = callPackage ../development/python-modules/exllamav2 { };
+
+  exllamav3 = callPackage ../development/python-modules/exllamav3 { };
+
  expandvars = callPackage ../development/python-modules/expandvars { };

  expects = callPackage ../development/python-modules/expects { };