Split torch tool for cpu, cuda and rocm

aotriton/aotriton-cms.patch

@@ -0,0 +1,27 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 5d88a22..8d3cded 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -127,8 +127,7 @@ install(FILES ${CMAKE_CURRENT_BINARY_DIR}/include/aotriton/config.h
+
+ # Kernel Storage V2 uses xz/LZMA for compression
+ if(UNIX)
+-  include(FindPkgConfig)
+-  pkg_search_module(LZMA REQUIRED liblzma)
++  find_package(LibLZMA REQUIRED)
+   add_library(lzma_interface INTERFACE)
+   target_link_libraries(lzma_interface INTERFACE ${LZMA_LIBRARIES})
+   target_link_directories(lzma_interface INTERFACE ${LZMA_LIBRARY_DIRS})
+diff --git a/v3python/base/typed_choice.py b/v3python/base/typed_choice.py
+index 2028988..94a6358 100644
+--- a/v3python/base/typed_choice.py
++++ b/v3python/base/typed_choice.py
+@@ -305,7 +305,7 @@ class Guess(object):
+
+ class GuessNumpy(Guess):
+     FACTORY = {
+-        np.bool   : constexpr.bool_t,
++        np.bool_  : constexpr.bool_t,
+         np.int8   : constexpr.int8_t,
+         np.int16  : constexpr.int16_t,
+         np.int32  : constexpr.int32_t,

aotriton/spec

@@ -0,0 +1,37 @@
+### RPM external aotriton 0.11.2b
+%define aotriton_git_sha1 dd1b68b604b5258ee7a9f7b66ad95e7a82c18065
+## INCLUDE rocm-flags
+%define rocm_gpus_cmake %(for ARCH in %{rocm_archs}; do printf "%s;" "${ARCH%%:*}"; done | sed 's/;$//')
+
+Source: git+https://github.com/ROCm/aotriton?obj=main/%{realversion}&export=%{n}-%{realversion}&submodules=1&output=/%{n}-%{realversion}.tgz
+Patch0: aotriton-cms
+Requires: py3-filelock py3-iniconfig py3-packaging py3-pluggy py3-numpy py3-setuptools py3-wheel py3-pybind11 py3-pandas py3-PyYAML
+Requires: rocm xz py3-triton
+BuildRequires: ninja cmake
+
+%prep
+%setup -n %{n}-%{realversion}
+%patch0 -p1
+
+%build
+rm -rf ../build ; mkdir ../build; cd ../build
+AOTRITON_CI_SUPPLIED_SHA1=%{aotriton_git_sha1} \
+cmake ../%{n}-%{realversion} \
+  -G Ninja \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DCMAKE_INSTALL_PREFIX:STRING=%{i} \
+  -DAOTRITON_TARGET_ARCH="%{rocm_gpus_cmake}" \
+  -DAOTRITON_NO_PYTHON=OFF \
+  -DAOTRITON_USE_TORCH=OFF \
+  -DAOTRITON_INHERIT_SYSTEM_SITE_TRITON=ON \
+  -DLZMA_LIBRARY_DIRS="${XZ_ROOT}/ilb" \
+  -DLZMA_LIBRARIES="${XZ_ROOT}/lib/liblzma.so" \
+  -DCMAKE_CXX_FLAGS="-I${XZ_ROOT}/include" \
+  -DCMAKE_PREFIX_PATH="%{cmake_prefix_path}" \
+  -DCMAKE_VERBOSE_MAKEFILE=TRUE
+
+ninja -v %{makeprocesses}
+
+%install
+cd ../build
+ninja -v %{makeprocesses} install

cmssw-gpu-backend-specific-packages.file

@@ -0,0 +1,2 @@
+%define gpu_backend_specific_packages py3-torch py3-torch-sparse py3-torch-cluster py3-torch-scatter
+%define gpu_types cuda rocm

cmssw-patch-tools.spec

@@ -5,5 +5,6 @@
 Requires: cmssw
 
 ## INCLUDE cmssw-vectorization
+## INCLUDE cmssw-gpu-backend-specific-packages
 ## INCLUDE cmssw-drop-tools
 ## INCLUDE scram/tool-conf-src

cmssw-queue-override.file

@@ -2,6 +2,7 @@
 ## NO_AUTO_RUNPATH
 ## NO_VERSION_SUFFIX
 ## UPLOAD_DEPENDENCIES dqmgui
+## INCLUDE cmssw-gpu-backend-specific-packages
 %if "%(case %realversion in (*_COVERAGE_X*) echo true ;; (*) echo false ;; esac)" == "true"
 %define release_usercxxflags    -fprofile-arcs -ftest-coverage
 %endif

cmssw-tools.spec

@@ -196,9 +196,9 @@ Requires: oracle-fake
 Requires: xtensor
 Requires: xtl
 Requires: xgboost
-Requires: pytorch-custom-ops
 
 ## INCLUDE tfaot-models
 ## INCLUDE cmssw-vectorization
+## INCLUDE cmssw-gpu-backend-specific-packages
 ## INCLUDE cmssw-drop-tools
 ## INCLUDE scram/tool-conf-src

pip/requirements.txt

@@ -405,11 +405,27 @@ tomli==2.4.0
 tomli-w==1.2.0
 tomlkit==0.14.0
 toolz==1.1.0
+#Torch packages: Update the version of -cuda/-rocm to match the default package version
+#e.g. for torch==2.10.0, torch-cuda and toch-rocm should also be 2.10.0
+#Do the same for other torch-* packages e.g version of torch-scatter should be used for
+#torch-scatter-cuda and torch-scatter-rocm too.
 torch==2.10.0
+torch-cuda==2.10.0
+torch-rocm==2.10.0
+torch-cluster==1.6.3
+torch-cluster-cuda==1.6.3
+torch-cluster-rocm==1.6.3
+torch-scatter==2.1.2
+torch-scatter-cuda==2.1.2
+torch-scatter-rocm==2.1.2
+torch-sparse==0.6.18
+torch-sparse-cuda==0.6.18
+torch-sparse-rocm==0.6.18
 torchvision==0.25.0
 tornado==6.5.4
 tqdm==4.67.3
 traitlets==5.14.3
+triton==3.6.0
 trove-classifiers==2026.1.14.14
 typed-ast==1.5.5
 types-python-dateutil==2.9.0.20260124

pip/torch-cluster-cuda.file

@@ -0,0 +1,10 @@
+## INCLUDE cuda-flags
+%define cuda_arch_float $(echo %{cuda_arch} | tr ' ' '\\n' | sed -E 's|([0-9])$|.\\1|' | tr '\\n' ' ' | sed 's| *$||')
+Requires: py3-torch-cuda
+%define pip_name torch-cluster
+
+%define PipPreBuildPy \
+  export TORCH_CUDA_ARCH_LIST="%{cuda_arch_float}" \
+  export FORCE_CUDA=1
+
+## INCLUDE pip/torch-extension-build

pip/torch-cluster.file

@@ -0,0 +1,2 @@
+Requires: py3-torch
+## INCLUDE pip/torch-extension-build

pip/torch-cuda.file

@@ -0,0 +1,14 @@
+## INCLUDE  cuda-flags
+%define cuda_arch_float $(echo %{cuda_arch} | tr ' ' '\\n' | sed -E 's|([0-9])$|.\\1|' | tr '\\n' ' ' | sed -e's/ *$/+PTX/')
+Requires: cuda cudnn
+%define pip_name torch
+%define PipPreBuildPy \
+  export USE_CUDA=ON \
+  export USE_CUDNN=ON \
+  export TORCH_CUDA_ARCH_LIST="%{cuda_arch_float}" \
+  export nvtx3_dir=${CUDA_ROOT}/include \
+  export CUDNN_ROOT=${CUDNN_ROOT}
+
+%define PipPreInstall mv  %{i}/bin/torchrun %{i}/bin/torchrun-cuda
+
+## INCLUDE pip/torch

pip/torch-extension-build.file

@@ -0,0 +1,6 @@
+BuildRequires: cmake ninja
+%define patchsrc sed -i -e 's|use_ninja=False|use_ninja=True|' setup.py
+%define PipPreBuild \
+  export CPATH="${PY3_PYBIND11_ROOT}/include" \
+  export MAX_JOBS=%{compiling_processes} \
+  export FORCE_CUDA=0

pip/torch-rocm.file

@@ -0,0 +1,17 @@
+## INCLUDE  rocm-flags
+Requires: rocm
+
+%define pip_name torch
+%define PipPreBuildPy \
+  export USE_ROCM=ON; \
+  export ROCM_PATH=${ROCM_ROOT}; \
+  export ROCM_SOURCE_DIR=${ROCM_ROOT}; \
+  export AMDGPU_TARGETS="$(echo '%{rocm_archs}' | sed -e 's/,/ /g')"; \
+  export PYTORCH_ROCM_ARCH="%{rocm_archs}"; \
+  pushd $PIPFILE; \
+    python3 tools/amd_build/build_amd.py; \
+  popd
+
+%define PipPreInstall mv  %{i}/bin/torchrun %{i}/bin/torchrun-rocm
+
+## INCLUDE pip/torch

pip/torch-scatter-cuda.file

@@ -0,0 +1,10 @@
+## INCLUDE cuda-flags
+%define cuda_arch_float $(echo %{cuda_arch} | tr ' ' '\\n' | sed -E 's|([0-9])$|.\\1|' | tr '\\n' ' ' | sed 's| *$||')
+Requires: py3-torch-cuda
+%define pip_name torch-scatter
+
+%define PipPreBuildPy \
+  export TORCH_CUDA_ARCH_LIST="%{cuda_arch_float}" \
+  export FORCE_CUDA=1
+
+## INCLUDE pip/torch-extension-build

pip/torch-scatter.file

@@ -0,0 +1,2 @@
+Requires: py3-torch
+## INCLUDE pip/torch-extension-build

pip/torch-sparse-cuda.file

@@ -0,0 +1,10 @@
+## INCLUDE cuda-flags
+%define cuda_arch_float $(echo %{cuda_arch} | tr ' ' '\\n' | sed -E 's|([0-9])$|.\\1|' | tr '\\n' ' ' | sed 's| *$||')
+Requires: py3-torch-cuda
+%define pip_name torch-sparse
+
+%define PipPreBuildPy \
+  export TORCH_CUDA_ARCH_LIST="%{cuda_arch_float}" \
+  export FORCE_CUDA=1
+
+## INCLUDE pip/torch-extension-build

pip/torch-sparse.file

@@ -0,0 +1,2 @@
+Requires: py3-torch
+## INCLUDE pip/torch-extension-build

pip/torch.file

@@ -1,4 +1,17 @@
-## INCLUDE  torch-requires
+## INCLUDE microarch_flags
+
+%define source0 git+https://github.com/pytorch/pytorch.git?obj=main/v%{realversion}&export=%{n}-%{realversion}&submodules=1&output=/%{n}-%{realversion}.tgz
+Source97: torch-build-env
+Source98: FindEigen3.cmake
+Source99: scram-tools.file/tools/eigen/env
+Patch0: patches/pytorch-system-fmt
+
+BuildRequires: cmake ninja python-python3 py3-packaging
+Requires: fftw3 eigen fxdiv numactl openmpi protobuf psimd python3 py3-PyYAML
+Requires: OpenBLAS zlib protobuf fmt py3-pybind11 py3-typing-extensions
+Requires: py3-filelock py3-fsspec py3-Jinja2 py3-networkx py3-sympy
+
+%define patchsrc0 cp %{_sourcedir}/FindEigen3.cmake cmake/Modules/
 
 %define PipPreBuild %{expand: %(cat %{_sourcedir}/torch-build-env)}
 %define PipPostInstall \

pip/triton.file

@@ -0,0 +1,26 @@
+%define source0 https://github.com/triton-lang/triton/releases/download/v%{realversion}/triton-%{realversion}.tar.gz
+BuildRequires: py3-pybind11 ninja
+Requires: json cuda zlib zstd libxml2
+
+#Triton does not support llvm 21 so lets build it using its internal llvm
+#Requires: llvm
+#export TRITON_APPEND_CMAKE_ARGS="-DLLVM_DIR=${LLVM_ROOT}/lib64/cmake/llvm -DLLD_DIR=${LLVM_ROOT}/lib64/cmake/lld -DMLIR_DIR=${LLVM_ROOT}/lib64/cmake/mlir"
+#export LLVM_SYSPATH="${LLVM_ROOT}"
+
+%define PipPreBuild \
+  export TRITON_CACHE_DIR=${TMPDIR} \
+  export TRITON_HOME=${TMPDIR} \
+  export MAX_JOBS=%{compiling_processes} \
+  export TRITON_BUILD_PYTHON_MODULE=ON \
+  export TRITON_BUILD_UT=OFF \
+  export TRITON_BUILD_WITH_CLANG_LLD=OFF \
+  export TRITON_BUILD_WITH_ASAN=OFF \
+  export TRITON_BUILD_WITH_CCACHE=OFF \
+  export TRITON_BUILD_PROTON=OFF \
+  export TRITON_PARALLEL_LINK_JOBS=%{compiling_processes} \
+  export JSON_SYSPATH=${JSON_ROOT} \
+  export TRITON_APPEND_CMAKE_ARGS="-DCMAKE_PREFIX_PATH='%{cmake_prefix_path}'" \
+  export TRITON_NVDISASM_PATH=$CUDA_ROOT/bin/nvdisasm \
+  export TRITON_CUOBJDUMP_PATH=$CUDA_ROOT/bin/cuobjdump \
+  export TRITON_PTXAS_PATH=${CUDA_ROOT}/bin/ptxas
+

python_tools.spec

@@ -19,6 +19,11 @@ Requires: py3-law
 Requires: py3-protobuf
 
 Requires: py3-torch
+Requires: py3-torch-cluster
+Requires: py3-torch-scatter
+Requires: py3-torch-sparse
+%{!?without_cuda:Requires: py3-torch-cuda py3-torch-cluster-cuda py3-torch-scatter-cuda py3-torch-sparse-cuda}
+
 Requires: py3-tables
 Requires: py3-tabulate
 Requires: py3-numexpr