Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153{ lib, buildPythonPackage, fetchFromGitHub, replaceVars, symlinkJoin,
# build-system setuptools, torch,
# buildInputs fmt, pybind11,
# nativeBuildInputs autoAddDriverRunpath,
# tests pytestCheckHook, writableTmpDirAsHomeHook,
# passthru deep-gemm,
config, cudaPackages, cudaSupport ? config.cudaSupport,}:
let inherit (lib) getBin optionalAttrs optionals ;inbuildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { pname = "deep-gemm"; version = "2.1.1.post3"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "deepseek-ai"; repo = "DeepGEMM"; tag = "v${finalAttrs.version}"; hash = "sha256-2yEHiuTaNUodWlZk7waqBsVMip2qiVJPgQHwsY0I63k="; };
patches = [ ./use-system-libraries.patch
# DeepGEMM does JIT compilation and needs to access the NVIDIA compiler and some libraries at # runtime. # Instead of letting it search for the cuda toolkit on the host, hardcode the path to a custom # closure. (replaceVars ./patch-runtime-cuda-home-path.patch { cuda_home = symlinkJoin { name = "cuda-toolkit"; paths = with cudaPackages; [ (lib.getBin cuda_nvcc) # bin/nvcc, bin/ptxas, nvvm/, nvcc.profile (lib.getBin cutlass) # include/cute, include/cutlass (lib.getInclude cccl) # include/cuda/std/* (libcu++) (lib.getInclude cuda_cudart) # include/cuda_runtime.h, cuda_bf16.h, cuda_fp8.h (lib.getInclude cuda_cuobjdump) # bin/cuobjdump ]; }; }) ];
postPatch = # Upstream hardcodes `__version__ = '2.1.1'` in `deep_gemm/__init__.py`, which does not match the # release tag and fails the metadata check. '' substituteInPlace deep_gemm/__init__.py \ --replace-fail \ "__version__ = '2.1.1'" \ "__version__ = '${finalAttrs.version}'" '' # torch >= 2.14 headers require C++20, while setup.py builds the host extension as C++17. # The JIT-compiled kernels already default to C++20. + '' substituteInPlace setup.py \ --replace-fail \ "'-std=c++17'" \ "'-std=c++20'" '';
env = { # Do not append a `+local` revision to the version (no git repository available) DG_USE_LOCAL_VERSION = "0"; } // optionalAttrs cudaSupport { CUDA_HOME = (getBin cudaPackages.cuda_nvcc).outPath;
LDFLAGS = toString [ # Fake libcuda.so (the real one is deployed impurely) "-L${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" ]; };
build-system = [ setuptools torch ];
nativeBuildInputs = [ autoAddDriverRunpath ];
buildInputs = [ fmt pybind11 ] ++ optionals cudaSupport ( with cudaPackages; [ cuda_cudart # cuda_runtime_api.h cuda_nvrtc # nvrtc.h cutlass # cute/arch/mma_sm100_desc.hpp libcublas # cublas_v2.h libcusolver # cusolverDn.h libcusparse # cusparse.h ] );
nativeCheckInputs = [ pytestCheckHook writableTmpDirAsHomeHook ];
# Tests require GPU access doCheck = false;
passthru.gpuCheck = deep-gemm.overridePythonAttrs { requiredSystemFeatures = [ "cuda" ];
# dlopens libcuda.so at import time pythonImportsCheck = [ "deep_gemm" ];
doCheck = true; };
meta = { description = "Clean and efficient FP8 GEMM kernels with fine-grained scaling"; homepage = "https://github.com/deepseek-ai/DeepGEMM"; license = lib.licenses.mit; maintainers = with lib.maintainers; [ GaetanLepage ]; broken = !cudaSupport; };})