From d12e33217cdac4262745b93614e74a38bff808e1 Mon Sep 17 00:00:00 2001 From: Yaroslav Bolyukin Date: Mon, 9 Mar 2026 02:51:10 +0100 Subject: [PATCH 1/2] python3Packages.amd-aiter: init at 0.1.11-post1 Co-authored-by: Luna Nova --- .../python-modules/amd-aiter/default.nix | 134 ++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 136 insertions(+) create mode 100644 pkgs/development/python-modules/amd-aiter/default.nix diff --git a/pkgs/development/python-modules/amd-aiter/default.nix b/pkgs/development/python-modules/amd-aiter/default.nix new file mode 100644 index 000000000000..c3f1f3bd9db4 --- /dev/null +++ b/pkgs/development/python-modules/amd-aiter/default.nix @@ -0,0 +1,134 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + setuptools-scm, + packaging, + psutil, + pybind11, + einops, + ninja, + numpy, + pandas, + rocmPackages, + torch, + writableTmpDirAsHomeHook, + + gpuTargets ? rocmPackages.clr.localGpuTargets or rocmPackages.clr.gpuTargets, + # From validate_and_update_archs at https://github.com/ROCm/aiter/blob/main/aiter/jit/core.py + supportedGpuTargets ? [ + "gfx90a" + "gfx940" + "gfx941" + "gfx942" + "gfx1100" + "gfx1101" + "gfx1102" + "gfx1103" + "gfx1150" + "gfx1151" + "gfx1152" + "gfx1153" + "gfx1200" + "gfx1201" + "gfx950" + ], +}: +buildPythonPackage (finalAttrs: { + pname = "amd-aiter"; + version = "0.1.11.post1"; + pyproject = true; + + src = fetchFromGitHub { + owner = "ROCm"; + repo = "aiter"; + tag = "v${finalAttrs.version}"; + hash = "sha256-9CCSmEw0kIoxERhtkKhBkaAGx42kCssH7IPjTgbg9LA="; + }; + + postPatch = '' + rmdir 3rdparty/composable_kernel + ln -sf ${rocmPackages.composable_kernel.src} 3rdparty/composable_kernel + + # python relax deps hook? + substituteInPlace pyproject.toml \ + --replace-fail '"flydsl==0.0.1.dev95158637"' "" + + # TODO: upstream fix. get_git_commit_id_short() assumes a git clone; + # fails in hermetic builds, should fall back gracefully. + substituteInPlace csrc/cpp_itfs/utils.py \ + --replace-fail \ + 'commit_id = get_git_commit_id_short()' \ + 'commit_id = "${finalAttrs.version}"' + + # setuptools runs setup.py twice (metadata + wheel). prepare_packaging() + # copies 3rdparty/ (with nix store read-only files) into aiter_meta/, then + # the second run can't rmtree or overwrite them. + substituteInPlace setup.py \ + --replace-fail \ + $'prepare_packaging()\n\n\nclass' \ + $'if not os.path.exists("aiter_meta"): prepare_packaging()\n\n\nclass' \ + --replace-fail \ + 'if os.path.exists("aiter_meta") and os.path.isdir("aiter_meta"):' \ + 'if False:' + ''; + + env = { + SETUPTOOLS_SCM_PRETEND_VERSION = finalAttrs.version; + PREBUILD_KERNELS = "0"; + BUILD_TARGET = "rocm"; + ROCM_PATH = "${rocmPackages.clr}"; + # FIXME does it even matter when PREBUILD_KERNELS is off? + GPU_ARCHS = lib.concatStringsSep ";" (lib.lists.intersectLists gpuTargets supportedGpuTargets); + }; + + build-system = [ + setuptools + setuptools-scm + packaging + psutil + pybind11 + ninja + pandas + ]; + + buildInputs = [ rocmPackages.clr ]; + + nativeBuildInputs = [ + rocmPackages.hipcc + writableTmpDirAsHomeHook + ]; + + pythonImportsCheck = [ "aiter" ]; + + dependencies = [ + einops + ninja + numpy + packaging + pandas + psutil + pybind11 + torch + ]; + + # Import requires writable $HOME (JIT cache) and ROCm GPU + # pythonImportsCheck = [ "aiter" ]; + + # Most of the tests require gpu + doCheck = false; + + meta = { + description = "AI Tensor Engine for ROCm"; + homepage = "https://github.com/ROCm/aiter"; + license = lib.licenses.mit; + sourceProvenance = with lib.sourceTypes; [ + fromSource + binaryNativeCode + ]; + maintainers = with lib.maintainers; [ lach ]; + teams = [ lib.teams.rocm ]; + platforms = lib.platforms.linux; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index c14b91280b61..c16a93a2e78d 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -722,6 +722,8 @@ self: super: with self; { amcrest = callPackage ../development/python-modules/amcrest { }; + amd-aiter = callPackage ../development/python-modules/amd-aiter { }; + amd-quark = callPackage ../development/python-modules/amd-quark { }; amdsmi = toPythonModule ( From c842253ddda374a3c967c12b1c9f54ec9d57b146 Mon Sep 17 00:00:00 2001 From: Luna Nova Date: Fri, 10 Apr 2026 09:36:37 -0700 Subject: [PATCH 2/2] python3Packages.amd-aiter: fix JIT include path, add passthru tests --- .../python-modules/amd-aiter/default.nix | 128 ++++++++++++------ 1 file changed, 87 insertions(+), 41 deletions(-) diff --git a/pkgs/development/python-modules/amd-aiter/default.nix b/pkgs/development/python-modules/amd-aiter/default.nix index c3f1f3bd9db4..021d39ae5836 100644 --- a/pkgs/development/python-modules/amd-aiter/default.nix +++ b/pkgs/development/python-modules/amd-aiter/default.nix @@ -1,40 +1,46 @@ { lib, + amd-aiter, buildPythonPackage, - fetchFromGitHub, - setuptools, - setuptools-scm, - packaging, - psutil, - pybind11, einops, + fetchFromGitHub, ninja, numpy, + packaging, pandas, + psutil, + pybind11, + python, rocmPackages, + runCommand, + setuptools, + setuptools-scm, torch, writableTmpDirAsHomeHook, - - gpuTargets ? rocmPackages.clr.localGpuTargets or rocmPackages.clr.gpuTargets, - # From validate_and_update_archs at https://github.com/ROCm/aiter/blob/main/aiter/jit/core.py - supportedGpuTargets ? [ - "gfx90a" - "gfx940" - "gfx941" - "gfx942" - "gfx1100" - "gfx1101" - "gfx1102" - "gfx1103" - "gfx1150" - "gfx1151" - "gfx1152" - "gfx1153" - "gfx1200" - "gfx1201" - "gfx950" - ], }: +let + # Provide a default set of include paths needed by aiter at runtime for JIT modules + defaultRocmIncl = lib.makeIncludePath ( + (with rocmPackages; [ + clr + hipblas + hipblas-common + hipblaslt + hipcub + hipfft + hipsolver + hipsparse + rocblas + rocprim + rocsolver + rocsparse + rocthrust + ]) + ++ [ + pybind11 + ] + ); +in buildPythonPackage (finalAttrs: { pname = "amd-aiter"; version = "0.1.11.post1"; @@ -51,7 +57,6 @@ buildPythonPackage (finalAttrs: { rmdir 3rdparty/composable_kernel ln -sf ${rocmPackages.composable_kernel.src} 3rdparty/composable_kernel - # python relax deps hook? substituteInPlace pyproject.toml \ --replace-fail '"flydsl==0.0.1.dev95158637"' "" @@ -62,6 +67,13 @@ buildPythonPackage (finalAttrs: { 'commit_id = get_git_commit_id_short()' \ 'commit_id = "${finalAttrs.version}"' + # NIX_AITER_ROCM_INCL is a colon-separated list of include dirs + # defaults to packages known to be needed by aiter + substituteInPlace aiter/jit/utils/cpp_extension.py \ + --replace-fail \ + 'paths.append(_join_rocm_home("include"))' \ + 'paths.append(_join_rocm_home("include")); paths.extend(os.environ.get("NIX_AITER_ROCM_INCL", "${defaultRocmIncl}").split(":"))' + # setuptools runs setup.py twice (metadata + wheel). prepare_packaging() # copies 3rdparty/ (with nix store read-only files) into aiter_meta/, then # the second run can't rmtree or overwrite them. @@ -75,22 +87,20 @@ buildPythonPackage (finalAttrs: { ''; env = { - SETUPTOOLS_SCM_PRETEND_VERSION = finalAttrs.version; - PREBUILD_KERNELS = "0"; BUILD_TARGET = "rocm"; + PREBUILD_KERNELS = "0"; ROCM_PATH = "${rocmPackages.clr}"; - # FIXME does it even matter when PREBUILD_KERNELS is off? - GPU_ARCHS = lib.concatStringsSep ";" (lib.lists.intersectLists gpuTargets supportedGpuTargets); + SETUPTOOLS_SCM_PRETEND_VERSION = finalAttrs.version; }; build-system = [ - setuptools - setuptools-scm + ninja packaging + pandas psutil pybind11 - ninja - pandas + setuptools + setuptools-scm ]; buildInputs = [ rocmPackages.clr ]; @@ -100,8 +110,6 @@ buildPythonPackage (finalAttrs: { writableTmpDirAsHomeHook ]; - pythonImportsCheck = [ "aiter" ]; - dependencies = [ einops ninja @@ -113,12 +121,50 @@ buildPythonPackage (finalAttrs: { torch ]; - # Import requires writable $HOME (JIT cache) and ROCm GPU - # pythonImportsCheck = [ "aiter" ]; - - # Most of the tests require gpu + # Most tests and imports require a GPU and writable $HOME for JIT cache doCheck = false; + # Test JIT module builds for CDNA3 iff rocm enabled for torch + passthru.tests = lib.optionalAttrs torch.rocmSupport ( + let + mkJitTest = + name: moduleName: + runCommand "amd-aiter-jit-${name}" + { + nativeBuildInputs = [ + (python.withPackages (_: [ amd-aiter ])) + rocmPackages.clr + writableTmpDirAsHomeHook + ]; + env = { + CXX = "amdclang++"; + GPU_ARCHS = "gfx942"; + PYTORCH_ROCM_ARCH = "gfx942"; + }; + } + '' + export AITER_JIT_DIR=$(mktemp -d) + python3 -c " + from aiter.jit.core import get_args_of_build, build_module + args = get_args_of_build('${moduleName}') + build_module( + '${moduleName}', + args['srcs'], args['flags_extra_cc'], args['flags_extra_hip'], + args['blob_gen_cmd'], args['extra_include'], args['extra_ldflags'], + args['verbose'], args['is_python_module'], args['is_standalone'], + args['torch_exclude'], + ) + print('JIT compile of ${moduleName} finished') + " + touch $out + ''; + in + { + jit-module-opus-sort = mkJitTest "opus-sort" "module_moe_sorting_opus"; + jit-module-mhc = mkJitTest "mhc" "module_mhc"; + } + ); + meta = { description = "AI Tensor Engine for ROCm"; homepage = "https://github.com/ROCm/aiter";