cudaPackages.cudnn-frontend: 1.16.0 -> 1.24.0 (#526689)

pkgs/by-name/dl/dlpack/package.nix

@@ -0,0 +1,37 @@
+{
+  lib,
+  fetchFromGitHub,
+  stdenv,
+
+  # nativeBuildInputs
+  cmake,
+}:
+stdenv.mkDerivation (finalAttrs: {
+  pname = "dlpack";
+  version = "1.3";
+  __structuredAttrs = true;
+  strictDeps = true;
+
+  src = fetchFromGitHub {
+    owner = "dmlc";
+    repo = "dlpack";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-kIHBgTYaHEmweRBFtRl1pXhOyQ5TEwU8dLUssTMEnpc=";
+  };
+
+  nativeBuildInputs = [
+    cmake
+  ];
+
+  # no tests
+  doCheck = false;
+
+  meta = {
+    description = "Open in-memory tensor structure for sharing tensors among frameworks";
+    homepage = "https://github.com/dmlc/dlpack";
+    downloadPage = "https://github.com/dmlc/dlpack/releases";
+    changelog = "https://github.com/dmlc/dlpack/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ prince213 ];
+  };
+})

pkgs/development/cuda-modules/packages/cudnn-frontend/package.nix

@@ -1,9 +1,7 @@
 {
-  autoAddDriverRunpath,
   backendStdenv,
   catch2_3,
   cmake,
-  cuda_cccl,
   cuda_cudart,
   cuda_nvcc,
   cuda_nvrtc,
@@ -13,19 +11,19 @@
   gitUpdater,
   lib,
   libcublas,
-  ninja,
   nlohmann_json,
+
+  withSamples ? true,
+  withTests ? true,
 }:
 let
   inherit (lib) licenses maintainers teams;
   inherit (lib.lists) optionals;
   inherit (lib.strings)
     cmakeBool
-    cmakeFeature
     optionalString
     ;
 in
-# TODO(@connorbaker): This should be a hybrid C++/Python package.
 backendStdenv.mkDerivation (finalAttrs: {
   __structuredAttrs = true;
   strictDeps = true;
@@ -34,13 +32,13 @@ backendStdenv.mkDerivation (finalAttrs: {
   name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
 
   pname = "cudnn-frontend";
-  version = "1.16.0";
+  version = "1.24.0";
 
   src = fetchFromGitHub {
     owner = "NVIDIA";
     repo = "cudnn-frontend";
     tag = "v${finalAttrs.version}";
-    hash = "sha256-+8aBl9dKd2Uz50XoOr91NRyJ4OGJtzfDNNNYGQJ9b94=";
+    hash = "sha256-I6el8e6Jo1l/S5eqxWiH2KksNxw4hJ+av5qj/yqJjI8=";
   };
 
   # nlohmann_json should be the only vendored dependency.
@@ -59,59 +57,56 @@ backendStdenv.mkDerivation (finalAttrs: {
   outputs = [
     "out"
   ]
-  ++ optionals finalAttrs.doCheck [
+  ++ optionals withSamples [
     "legacy_samples"
     "samples"
+  ]
+  ++ optionals withTests [
     "tests"
   ];
 
   nativeBuildInputs = [
-    autoAddDriverRunpath # Needed for samples because it links against CUDA::cuda_driver
     cmake
     cuda_nvcc
-    ninja
   ];
 
   buildInputs = [
-    cuda_cccl
     cuda_cudart
+  ]
+  ++ optionals (withSamples || withTests) [
+    catch2_3
+    cuda_nvrtc
+    cudnn
+    libcublas
   ];
 
   cmakeFlags = [
-    (cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
-    (cmakeFeature "FETCHCONTENT_TRY_FIND_PACKAGE_MODE" "ALWAYS")
-    (cmakeBool "CUDNN_FRONTEND_BUILD_SAMPLES" finalAttrs.doCheck)
-    (cmakeBool "CUDNN_FRONTEND_BUILD_TESTS" finalAttrs.doCheck)
-    (cmakeBool "CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS" false)
-  ];
-
-  checkInputs = [
-    cudnn
-    cuda_nvrtc
-    catch2_3
-    libcublas
+    (cmakeBool "CUDNN_FRONTEND_BUILD_SAMPLES" withSamples)
+    (cmakeBool "CUDNN_FRONTEND_BUILD_TESTS" withTests)
   ];
 
   enableParallelBuilding = true;
 
   propagatedBuildInputs = [
     nlohmann_json
+    cuda_nvrtc # nvrtc.h
   ];
 
-  # TODO(@connorbaker): I'm using this incorrectly to build the executables which would allow us to test functionality,
-  # rather than to indicate the checkPhase will actually run.
-  doCheck = true;
-
-  postInstall = optionalString finalAttrs.doCheck ''
-    moveToOutput "bin/legacy_samples" "$legacy_samples"
-    moveToOutput "bin/samples" "$samples"
-    moveToOutput "bin/tests" "$tests"
-    if [[ -e "$out/bin" ]]
-    then
-      nixErrorLog "The bin directory in \$out should no longer exist."
-      exit 1
-    fi
-  '';
+  postInstall =
+    optionalString withSamples ''
+      moveToOutput "bin/legacy_samples" "$legacy_samples"
+      moveToOutput "bin/samples" "$samples"
+    ''
+    + optionalString withTests ''
+      moveToOutput "bin/tests" "$tests"
+    ''
+    + ''
+      if [[ -e "$out/bin" ]]
+      then
+        nixErrorLog "The bin directory in \$out should no longer exist."
+        exit 1
+      fi
+    '';
 
   passthru.updateScript = gitUpdater {
     inherit (finalAttrs) pname version;
@@ -119,11 +114,13 @@ backendStdenv.mkDerivation (finalAttrs: {
   };
 
   meta = {
-    description = "A c++ wrapper for the cudnn backend API";
+    description = "Python and C++ Graph API with SOTA attention (SDPA / Flash Attention), MoE grouped GEMM fusions, and FP8/MXFP8 kernels for Hopper and Blackwell GPUs";
     homepage = "https://github.com/NVIDIA/cudnn-frontend";
+    downloadPage = "https://github.com/NVIDIA/cudnn-frontend/releases";
+    changelog = "https://github.com/NVIDIA/cudnn-frontend/releases/tag/${finalAttrs.src.tag}";
     license = licenses.mit;
     # Supports cuDNN 8.5.0 and newer:
-    # https://github.com/NVIDIA/cudnn-frontend/blob/11b51e9c5ad6cc71cd66cb873e34bc922d97d547/README.md?plain=1#L32
+    # https://github.com/NVIDIA/cudnn-frontend/blob/v1.24.0/README.md?plain=1#L83
     platforms = [
       "aarch64-linux"
       "x86_64-linux"

pkgs/development/python-modules/nvidia-cudnn-frontend/default.nix

@@ -0,0 +1,95 @@
+{
+  lib,
+  buildPythonPackage,
+  nvidia-cudnn-frontend,
+
+  # build-system
+  cmake,
+  ninja,
+  pybind11,
+  setuptools,
+
+  # nativeBuildInputs
+  cudaPackages,
+
+  # buildInputs
+  dlpack,
+
+  # propagatedBuildInputs
+  nlohmann_json,
+
+  # tests
+  looseversion,
+  pytestCheckHook,
+  torch,
+}:
+buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: {
+  inherit (cudaPackages.cudnn-frontend)
+    version
+    src
+    meta
+    ;
+
+  pname = "nvidia-cudnn-frontend";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  postPatch =
+    cudaPackages.cudnn-frontend.postPatch
+    + ''
+      substituteInPlace pyproject.toml \
+        --replace-fail '"ninja==1.11.1.1"' '"ninja"' \
+        --replace-fail '"pybind11[global]>=2.13,<3"' '"pybind11"'
+
+      sed -i '/cmake_args =/a\\f"-DCUDNN_FRONTEND_USE_SYSTEM_DLPACK=ON",' setup.py
+    ''
+    + ''
+      substituteInPlace python/cudnn/__init__.py \
+        --replace-fail \
+          'os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/lib/libcudnn.so.*[0-9]")' \
+          '"${lib.getLib cudaPackages.cudnn}/lib/libcudnn.so"'
+    '';
+
+  build-system = [
+    cmake
+    ninja
+    pybind11
+    setuptools
+  ];
+
+  dontUseCmakeConfigure = true;
+
+  nativeBuildInputs = [
+    cudaPackages.cuda_nvcc
+  ];
+
+  buildInputs = [
+    cudaPackages.cuda_cudart
+    cudaPackages.cuda_nvrtc # nvrtc.h
+    cudaPackages.cudnn
+    dlpack
+  ];
+
+  propagatedBuildInputs = [
+    nlohmann_json
+  ];
+
+  pythonImportsCheck = [ "cudnn" ];
+
+  # requires GPU
+  doCheck = false;
+  nativeCheckInputs = [
+    looseversion
+    pytestCheckHook
+    torch
+  ];
+
+  enabledTestPaths = [
+    "test/"
+  ];
+
+  passthru.gpuCheck = nvidia-cudnn-frontend.overridePythonAttrs {
+    requiredSystemFeatures = [ "cuda" ];
+    doCheck = true;
+  };
+})

pkgs/top-level/python-packages.nix

@@ -11462,6 +11462,8 @@ self: super: with self; {
 
   nvdlib = callPackage ../development/python-modules/nvdlib { };
 
+  nvidia-cudnn-frontend = callPackage ../development/python-modules/nvidia-cudnn-frontend { };
+
   nvidia-cutlass = callPackage ../development/python-modules/nvidia-cutlass { };
 
   nvidia-cutlass-dsl = callPackage ../development/python-modules/nvidia-cutlass-dsl { };