diff --git a/pkgs/development/python-modules/nvidia-cutlass/default.nix b/pkgs/development/python-modules/nvidia-cutlass/default.nix new file mode 100644 index 000000000000..9bb17e3cf1a0 --- /dev/null +++ b/pkgs/development/python-modules/nvidia-cutlass/default.nix @@ -0,0 +1,76 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + + # dependencies + cuda-bindings, + cuda-pathfinder, + networkx, + numpy, + pydot, + scipy, + treelib, + + # tests + pytestCheckHook, + torch, +}: + +buildPythonPackage (finalAttrs: { + pname = "nvidia-cutlass"; + version = "4.2.0"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "NVIDIA"; + repo = "cutlass"; + tag = "v${finalAttrs.version}"; + hash = "sha256-XicHeV9ni9bSOWcUJM8HrCuz61mVK1EdZ9uxNvgWmvk="; + }; + + build-system = [ + setuptools + ]; + + pythonRemoveDeps = [ + # Replaced with the cuda-python sub-packages we actually need + "cuda-python" + ]; + dependencies = [ + cuda-bindings + cuda-pathfinder + networkx + numpy + pydot + scipy + treelib + ]; + + pythonImportsCheck = [ + "cutlass_cppgen" + "cutlass_library" + "pycute" + ]; + + nativeCheckInputs = [ + pytestCheckHook + torch + ]; + + enabledTestPaths = [ + "test" + ]; + + meta = { + description = "Python bindings for NVIDIA's CUTLASS library"; + homepage = "https://github.com/NVIDIA/cutlass"; + changelog = "https://github.com/NVIDIA/cutlass/blob/${finalAttrs.src.tag}/CHANGELOG.md"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 32df071d53a4..ce82af5e3179 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11439,6 +11439,8 @@ self: super: with self; { nvdlib = callPackage ../development/python-modules/nvdlib { }; + nvidia-cutlass = callPackage ../development/python-modules/nvidia-cutlass { }; + nvidia-dlprof-pytorch-nvtx = callPackage ../development/python-modules/nvidia-dlprof-pytorch-nvtx { };