gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.triton: 3.6.0 -> 3.7.0; python3Packages.torch: 2.11.0 -> 2.12.0 (#519925)

Browse Source
This commit is contained in:
Nick Cao
2026-06-03 23:33:06 +00:00
committed by GitHub
parent 6f41fa8bd1 676ea585a7
commit dda0131117
14 changed files with 178 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/by-name/tr/triton-llvm/package.nix
+5 -3
View File
@@ -68,7 +68,9 @@ let
in
stdenv.mkDerivation (finalAttrs: {
pname = "triton-llvm";
version = "22.0.0-unstable-2025-09-26"; # See https://github.com/llvm/llvm-project/blob/main/cmake/Modules/LLVMVersion.cmake
version = "23.0.0-unstable-2026-01-29"; # See https://github.com/llvm/llvm-project/blob/main/cmake/Modules/LLVMVersion.cmake
__structuredAttrs = true;
outputs = [
"out"
@@ -84,8 +86,8 @@ stdenv.mkDerivation (finalAttrs: {
src = fetchFromGitHub {
owner = "llvm";
repo = "llvm-project";
rev = "f6ded0be897e2878612dd903f7e8bb85448269e5";
hash = "sha256-T76zHZZ2bp3Ye9GTV+MgbKqMbtmMGElMFsWuCkiWqrM=";
rev = "ac5dc54d509169d387fcfd495d71853d81c46484";
hash = "sha256-tA1KcZqyPsgfxQs9tbNhX11oFcNGJefxWmfCpYqdI9M=";
};
nativeBuildInputs = [
pkgs/development/python-modules/bitsandbytes/default.nix
+2
View File
@@ -9,6 +9,7 @@
cmake,
# build-system
ninja,
scikit-build-core,
setuptools,
@@ -121,6 +122,7 @@ buildPythonPackage (finalAttrs: {
];
build-system = [
ninja
scikit-build-core
setuptools
];
pkgs/development/python-modules/executorch/default.nix
-6
View File
@@ -102,12 +102,6 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
sed -i "1i #include <cstdint>" backends/apple/coreml/runtime/inmemoryfs/memory_buffer.hpp
sed -i "1i #include <cstdint>" extension/llm/tokenizers/third-party/sentencepiece/src/sentencepiece_processor.h
''
+ ''
substituteInPlace extension/llm/tokenizers/test/test_python_bindings.py \
--replace-fail \
'self.assertEqual(pytorch_tokenizers.__version__, "0.1.0")' \
'self.assertEqual(pytorch_tokenizers.__version__, "${pytorch-tokenizers.version}")'
'';
env = {
pkgs/development/python-modules/torch/source/default.nix
+6 -36
View File
@@ -3,7 +3,6 @@
lib,
fetchFromGitHub,
fetchFromGitLab,
fetchpatch,
git-unroll,
buildPythonPackage,
python,
@@ -48,7 +47,6 @@
llvmPackages,
# dependencies
astunparse,
binutils,
expecttest,
filelock,
@@ -229,6 +227,9 @@ let
rocm-smi
clr.icd
hipify
rocprofiler-sdk
rocprofiler-sdk.dev
amdsmi
]
++ lib.optionals (!vendorComposableKernel) [
composable_kernel
@@ -282,8 +283,9 @@ in
buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
pname = "torch";
# Don't forget to update torch-bin to the same version.
version = "2.11.0";
version = "2.12.0";
pyproject = true;
__structuredAttrs = true;
outputs = [
"out" # output standard python package
@@ -306,29 +308,6 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
patches = [
./clang19-template-warning.patch
# The GCC version upperbounds were wrong for cuda 12.8 and 12.9, which led downstream builds to
# illegitimately fail with:
# RuntimeError: The current installed version of g++ (14.3.0) is greater than the maximum
# required version by CUDA 12.9. Please make sure to use an adequate version of g++
# (>=6.0.0, <14.0).
# TODO: remove at the next release
(fetchpatch {
name = "allow-gcc-14-with-cuda-12.8-9";
url = "https://github.com/pytorch/pytorch/commit/39565a7dcf8f93ea22cedeaa20088b24ff6d2634.patch";
hash = "sha256-Au5fVbs7i33d9c4Xj8koiBP7lGnsTGTaX4VlE2gAfy8=";
})
# pybind11 3.0.3 changes led to ambiguous deduction in some return types
# that used `py::make_tuple`, so the type is explicitly specified where
# needed.
# Merged pull request: https://github.com/pytorch/pytorch/pull/179277
# TODO: remove at the next release
(fetchpatch {
name = "pybind11-3.0.3-ambiguous-return-type.patch";
url = "https://github.com/pytorch/pytorch/commit/b248ebc17075c0c3ad2b2532970d2ada32b2cf94.patch";
hash = "sha256-HY5JFGNoroFsfuUOO5j6WNP6gMHWUcIJFmWLqV8PV94=";
})
]
++ lib.optionals cudaSupport [
./fix-cmake-cuda-toolkit.patch
@@ -400,13 +379,6 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
--replace-fail "list(APPEND ATen_HIP_INCLUDE \''${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/include)" "" \
--replace-fail "list(APPEND ATen_HIP_INCLUDE \''${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/library/include)" ""
''
# Detection of NCCL version doesn't work particularly well when using the static binary.
+ lib.optionalString cudaSupport ''
substituteInPlace cmake/Modules/FindNCCL.cmake \
--replace-fail \
'message(FATAL_ERROR "Found NCCL header version and library version' \
'message(WARNING "Found NCCL header version and library version'
''
# Remove PyTorch's FindCUDAToolkit.cmake and use CMake's default.
# NOTE: Parts of pytorch rely on unmaintained FindCUDA.cmake with custom patches to support e.g.
# newer architectures (sm_90a). We do want to delete vendored patches, but have to keep them
@@ -521,9 +493,9 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
cmakeFlags = [
(lib.cmakeFeature "PYTHON_SIX_SOURCE_DIR" "${six.src}")
# (lib.cmakeBool "CMAKE_FIND_DEBUG_MODE" true)
(lib.cmakeFeature "CUDAToolkit_VERSION" cudaPackages.cudaMajorMinorVersion)
]
++ lib.optionals cudaSupport [
(lib.cmakeFeature "CUDAToolkit_VERSION" cudaPackages.cudaMajorMinorVersion)
# Unbreaks version discovery in enable_language(CUDA) when wrapping nvcc with ccache
# Cf. https://gitlab.kitware.com/cmake/cmake/-/issues/26363
(lib.cmakeFeature "CMAKE_CUDA_COMPILER_TOOLKIT_VERSION" cudaPackages.cudaMajorMinorVersion)
@@ -620,14 +592,12 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
"sympy"
];
dependencies = [
astunparse
expecttest
filelock
fsspec
hypothesis
jinja2
networkx
ninja
packaging
psutil
pyyaml
pkgs/development/python-modules/torch/source/src.nix
+24 -17
View File
@@ -4,7 +4,7 @@
fetchFromGitHub,
runCommand,
}:
assert version == "2.11.0";
assert version == "2.12.0";
rec {
src_aiter = fetchFromGitHub {
owner = "ROCm";
@@ -43,6 +43,12 @@ rec {
hash = "sha256-IDUIuAvgCzWaHoTJUZrH15bqoVcP8bZk+Gs1Ae6/CpY=";
};
src_composable_kernel = fetchFromGitHub {
owner = "ROCm";
repo = "composable_kernel";
rev = "f1746955fdaf80a3414de814bf32437686dac347";
hash = "sha256-B/xNuBPUdjL1b+0IzRnaSXT2FKUo5cYwYcKqfKqJ8Eg=";
};
src_composable_kernel_aiter = fetchFromGitHub {
owner = "ROCm";
repo = "composable_kernel";
rev = "fcc9372c009c8e0a23fece77b582da83b04a654f";
@@ -63,8 +69,8 @@ rec {
src_cpp-httplib = fetchFromGitHub {
owner = "yhirose";
repo = "cpp-httplib";
rev = "bd95e67c234930cd6d6bb11309588c5462c63cec";
hash = "sha256-5q77ersAJnPPpVChvntnqEly1/ek2KfX2iukTPUbKHc=";
rev = "4d7c9a788de136071ccf0dd4e96239151e2adadb";
hash = "sha256-VXEhoxoQjGEuA2g/y6fDTA4LrPd4SggrS3aOjznDSvc=";
};
src_cpr = fetchFromGitHub {
owner = "libcpr";
@@ -87,14 +93,14 @@ rec {
src_cudnn-frontend = fetchFromGitHub {
owner = "NVIDIA";
repo = "cudnn-frontend";
rev = "b8c0656e6f6c84fc194f4d57329b55d609eff596";
hash = "sha256-G1WYxRCsg67umzOZ9W+JwXV6hfl5n2wtsH9KxVUccTU=";
rev = "a91f0e04dcea10515f0f776fc5a89535e316a9c8";
hash = "sha256-OOKdkjsVnWgrtcI7IMPSRi2YxtqF2YNV4Fd2rD9I1K8=";
};
src_cutlass = fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
rev = "0d2b201e8c1c4a03efa6e9c468161916e2334725";
hash = "sha256-2xgdS2P2tEI/4qcZv9qjCYHFbAcVayMdCDeJfIiQN4U=";
rev = "da5e086dab31d63815acafdac9a9c5893b1c69e2";
hash = "sha256-0q9Ad0Z6E/rO2PdM4uQc8H0E0qs9uKc3reHepiHhjEc=";
};
src_cutlass_fbgemm = fetchFromGitHub {
owner = "jwfromm";
@@ -141,8 +147,8 @@ rec {
src_flash-attention = fetchFromGitHub {
owner = "Dao-AILab";
repo = "flash-attention";
rev = "e2743ab5b3803bb672b16437ba98a3b1d4576c50";
hash = "sha256-ft3jPiKZDHzZvkGPI34l8/Hq9Rf2f6UjDPKGU2IYz+E=";
rev = "fec3a6a18460c1b40f097208d4c16fe8964a679d";
hash = "sha256-7yEFNM2lslkBA/9slblAbiK1PHKqKmo1MCFJYz2BOLk=";
};
src_flatbuffers = fetchFromGitHub {
owner = "google";
@@ -243,7 +249,7 @@ rec {
src_ideep = fetchFromGitHub {
owner = "intel";
repo = "ideep";
rev = "8e7ddd65df95f13e41f0a40c820c5f35ae4a0ea3";
rev = "e539e0f9774e2018f0d56fe865da66581f692e3d";
hash = "sha256-AVSsugGYiQ4QOWMVaHj1hzlPTZmg65yrGMmrWytvUuM=";
};
src_ittapi = fetchFromGitHub {
@@ -273,8 +279,8 @@ rec {
src_kineto = fetchFromGitHub {
owner = "pytorch";
repo = "kineto";
rev = "7a731b6ae01cfc2b1fc75d83a91f84e682e43fd7";
hash = "sha256-kKASzILEvFhXDWDqBiNh21VxhtdT506NYoEDVrtVcGU=";
rev = "b2103f78d13fde4937af010c0ef8e24313568bc5";
hash = "sha256-Ix5zulGaUPbLeVrrTm/EzcVWT4TkDYcBsQADAl4N7TA=";
};
src_kleidiai = fetchFromGitHub {
owner = "ARM-software";
@@ -303,8 +309,8 @@ rec {
src_mkl-dnn = fetchFromGitHub {
owner = "intel";
repo = "mkl-dnn";
rev = "f1d471933dc852f956fd05389f9313c7148783d5";
hash = "sha256-/e57voLBNun/2koTF3sEb0Z/nDjCwq9NJVk7TaTSvMY=";
rev = "03c022d3ffdcee958cfacbe720048e725fdf644c";
hash = "sha256-xJTllrKs6mPNM85ZqyHTHWKpVOtOghmg4ZRFAvQZ4WU=";
};
src_MSLK = fetchFromGitHub {
owner = "meta-pytorch";
@@ -393,8 +399,8 @@ rec {
src_pytorch = fetchFromGitHub {
owner = "pytorch";
repo = "pytorch";
rev = "v2.11.0";
hash = "sha256-eA/pRQiibLJCKDUkMvmGq4suhrz37i0x0Lc/sH3Ag8E=";
rev = "v2.12.0";
hash = "sha256-IyQs9CQbbpZYpd+8YhIj/ULjsIWu6gjkGrGSeMWqKvw=";
};
src_sleef = fetchFromGitHub {
owner = "shibatch";
@@ -423,7 +429,7 @@ rec {
src_aiter_recursive = runCommand "aiter" { } ''
cp -r ${src_aiter} $out
chmod u+w $out/3rdparty/composable_kernel
cp -r ${src_composable_kernel_recursive}/* $out/3rdparty/composable_kernel
cp -r ${src_composable_kernel_aiter_recursive}/* $out/3rdparty/composable_kernel
'';
src_asmjit_recursive = src_asmjit;
src_benchmark_recursive = src_benchmark;
@@ -431,6 +437,7 @@ rec {
src_civetweb_recursive = src_civetweb;
src_clang-cindex-python3_recursive = src_clang-cindex-python3;
src_composable_kernel_recursive = src_composable_kernel;
src_composable_kernel_aiter_recursive = src_composable_kernel_aiter;
src_composable_kernel_fbgemm_MSLK_recursive = src_composable_kernel_fbgemm_MSLK;
src_composable_kernel_flash-attention_recursive = src_composable_kernel_flash-attention;
src_cpp-httplib_recursive = src_cpp-httplib;
pkgs/development/python-modules/torchcodec/default.nix
+15 -3
View File
@@ -26,14 +26,15 @@
buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
pname = "torchcodec";
version = "0.11.1";
version = "0.14.0";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "meta-pytorch";
repo = "torchcodec";
tag = "v${finalAttrs.version}";
hash = "sha256-aYQp9vEVQJgF1n/KsfnDvLQf5nD0/gsG+RAgVlhk7t8=";
hash = "sha256-eGof2Rk/dGYPlKVRSuJ+ZeeMh2u4K6/qXmROo187HTA=";
};
postPatch = ''
@@ -60,6 +61,9 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
]
++ lib.optionals cudaSupport [
cudaPackages.cuda_nvcc
]
++ lib.optionals rocmSupport [
torch.rocmPackages.clr
];
buildInputs = [
@@ -110,7 +114,15 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
];
disabledTests =
lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
lib.optionals rocmSupport [
# HSA runtime logs topology error in sandbox breaking test that asserts no output
"test_python_logger"
]
++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
# Fails in the sandbox:
# Error in cpuinfo: failed to parse the list of possible processors in /sys/devices/system/cpu/possible
"test_python_logger"
# AssertionError: index 0
"test_get_frames_played_at"
pkgs/development/python-modules/torchdata/default.nix
+15 -11
View File
@@ -56,18 +56,22 @@ buildPythonPackage (finalAttrs: {
pytestCheckHook
];
disabledTests =
lib.optionals (pythonAtLeast "3.14") [
# _pickle.PicklingError: Can't pickle local object <...worker_set_affinity at 0x7ffbcfa1ba00>
"test_set_affinity_in_worker_init"
disabledTests = [
# Failing since torch 2.12.0 update
# AssertionError: Scalars are not equal! Expected 0 but got 1.
"test_get_worker_info"
]
++ lib.optionals (pythonAtLeast "3.14") [
# _pickle.PicklingError: Can't pickle local object <...worker_set_affinity at 0x7ffbcfa1ba00>
"test_set_affinity_in_worker_init"
# RuntimeError: DataLoader timed out after 5 seconds
"test_ind_worker_queue"
]
++ lib.optionals stdenv.hostPlatform.isDarwin [
# RuntimeError: DataLoader timed out after 5 seconds
"test_ind_worker_queue"
];
# RuntimeError: DataLoader timed out after 5 seconds
"test_ind_worker_queue"
]
++ lib.optionals stdenv.hostPlatform.isDarwin [
# RuntimeError: DataLoader timed out after 5 seconds
"test_ind_worker_queue"
];
meta = {
description = "Iterative enhancement to the PyTorch torch.utils.data.DataLoader and torch.utils.data.Dataset/IterableDataset";
pkgs/development/python-modules/torchvision/default.nix
+3 -2
View File
@@ -28,14 +28,15 @@ let
in
buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
pname = "torchvision";
version = "0.26.0";
version = "0.27.0";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "pytorch";
repo = "vision";
tag = "v${finalAttrs.version}";
hash = "sha256-FOdDGY3v8yWBhtNo9tZP79/xwrc7AoIY5Y1ZABzWe6g=";
hash = "sha256-HOTD45xY7Gye1GI1+AsF3KmMUTAp1QlzHOUeBHvzv0A=";
};
nativeBuildInputs = [
pkgs/development/python-modules/triton/0003-nvidia-cudart-a-systempath.patch
+5 -14
View File
@@ -1,18 +1,18 @@
diff --git a/third_party/nvidia/backend/driver.c b/third_party/nvidia/backend/driver.c
index bff09d8c1..a5c341711 100644
index f51d7b37c1..8971588282 100644
--- a/third_party/nvidia/backend/driver.c
+++ b/third_party/nvidia/backend/driver.c
@@ -1,4 +1,4 @@
-#include "cuda.h"
+#include <cuda.h>
#include <dlfcn.h>
#include <stdalign.h>
#include <stdbool.h>
#include <stdlib.h>
diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py
index 2b39fea29..3346eb954 100644
index 21e9045b25..3660e32546 100644
--- a/third_party/nvidia/backend/driver.py
+++ b/third_party/nvidia/backend/driver.py
@@ -12,7 +12,8 @@ from triton.backends.compiler import GPUTarget
@@ -11,7 +11,8 @@ from triton.backends.compiler import GPUTarget
from triton.backends.driver import GPUDriver
dirname = os.path.dirname(os.path.realpath(__file__))
@@ -20,14 +20,5 @@ index 2b39fea29..3346eb954 100644
+import shlex
+include_dirs = [*shlex.split("@cudaToolkitIncludeDirs@"), os.path.join(dirname, "include")]
libdevice_dir = os.path.join(dirname, "lib")
libraries = ['cuda']
libraries = ['libcuda.so.1']
PyCUtensorMap = None
@@ -265,7 +266,7 @@ def make_launcher(constants, signature, tensordesc_meta):
params.append("&global_scratch")
params.append("&profile_scratch")
src = f"""
-#include \"cuda.h\"
+#include <cuda.h>
#include <dlfcn.h>
#include <stdbool.h>
#include <stdlib.h>
pkgs/development/python-modules/triton/0005-add-gcn5-gfx906-target.patch
-131
View File
@@ -1,131 +0,0 @@
diff --git a/test/Conversion/amd/tritongpu_to_llvm.mlir b/test/Conversion/amd/tritongpu_to_llvm.mlir
index 068e931a6..e651f7b0a 100644
--- a/test/Conversion/amd/tritongpu_to_llvm.mlir
+++ b/test/Conversion/amd/tritongpu_to_llvm.mlir
@@ -1,5 +1,6 @@
// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx942 --convert-builtin-func-to-llvm | FileCheck %s
// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx950 | FileCheck %s --check-prefix=GFX950
+// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx906 | FileCheck %s --check-prefix=GFX906
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
// CHECK-LABEL: atomic_add_f32_scalar
@@ -633,3 +634,27 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
tt.return
}
}
+
+// -----
+
+// GFX906-LABEL: v_dot_fp16_gfx906
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [2, 2], order = [1, 0]}>
+module attributes {"ttg.target" = "hip:gfx906", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
+ tt.func @v_dot_fp16_gfx906(%arg0: tensor<16x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>, %arg1: tensor<16x16xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>, %arg2: tensor<16x16xf32, #blocked>) {
+ // GFX906-COUNT-8: llvm.call_intrinsic "llvm.amdgcn.fdot2"
+ %0 = tt.dot %arg0, %arg1, %arg2, inputPrecision = ieee : tensor<16x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<16x16xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xf32, #blocked>
+ tt.return
+ }
+}
+
+// -----
+
+// GFX906-LABEL: v_dot_i8_gfx906
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [2, 2], order = [1, 0]}>
+module attributes {"ttg.target" = "hip:gfx906", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
+ tt.func @v_dot_i8_gfx906(%arg0: tensor<16x16xi8, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>, %arg1: tensor<16x16xi8, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>, %arg2: tensor<16x16xi32, #blocked>) {
+ // GFX906-COUNT-4: llvm.call_intrinsic "llvm.amdgcn.sdot4"
+ %0 = tt.dot %arg0, %arg1, %arg2, inputPrecision = ieee : tensor<16x16xi8, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<16x16xi8, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xi32, #blocked>
+ tt.return
+ }
+}
diff --git a/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h b/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
index a689fe438..5d84110dc 100644
--- a/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
+++ b/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
@@ -8,6 +8,7 @@ namespace mlir::triton::AMD {
// A list of ISA families we care about.
enum class ISAFamily {
Unknown,
+ GCN5_1,
CDNA1,
CDNA2,
CDNA3,
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
index af7b4b1f7..01cacc9de 100644
--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
@@ -71,6 +71,7 @@ llvm::AMDGPU::GPUKind TargetInfo::getGPUKind() const {
int TargetInfo::getWarpSize() const {
switch (getISAFamily()) {
+ case ISAFamily::GCN5_1:
case ISAFamily::CDNA1:
case ISAFamily::CDNA2:
case ISAFamily::CDNA3:
@@ -335,10 +336,10 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
return true;
if (numLaneToReduce != getWarpSize())
return false;
- if (isCDNA(getISAFamily()) && getISAFamily() == ISAFamily::CDNA1)
- return false;
- if (isRDNA(getISAFamily()) &&
- llvm::is_contained({ISAFamily::RDNA1, ISAFamily::RDNA2}, getISAFamily()))
+ // DPP warp reduce requires gfx90a+ (CDNA2+) or gfx11+ (RDNA3+).
+ // Pre-CDNA2 GFX9 (gfx906/gfx908) and GFX10 (RDNA1/2) are excluded.
+ auto v = getIsaVersion();
+ if (!((v.Major == 9 && (v.Minor > 0 || v.Stepping >= 0xa)) || v.Major >= 11))
return false;
Operation *reduxOp = op.getSingleCombiner();
@@ -438,7 +439,7 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
buf = createDppReduxOpWithBoundCtrl(valType, buf, 1 + dppCtrlRowShr,
allRows, allBanks);
- if (isCDNA(getISAFamily())) {
+ if (isCDNA(getISAFamily()) || getISAFamily() == ISAFamily::GCN5_1) {
// row_bcast:15 row_mask:0xa
buf = createDppReduxOpWithBoundCtrl(
valType, buf, static_cast<uint32_t>(DppCtrl::BCAST15), 0xa, allBanks);
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej
new file mode 100644
index 000000000..b75ab8f11
--- /dev/null
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej
@@ -0,0 +1,16 @@
+diff a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp (rejected hunks)
+@@ -395,10 +396,10 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
+ return true;
+ if (reduceLaneIdMask != (getWarpSize() - 1))
+ return false;
+- if (isCDNA(getISAFamily()) && getISAFamily() == ISAFamily::CDNA1)
+- return false;
+- if (isRDNA(getISAFamily()) &&
+- llvm::is_contained({ISAFamily::RDNA1, ISAFamily::RDNA2}, getISAFamily()))
++ // DPP warp reduce requires gfx90a+ (CDNA2+) or gfx11+ (RDNA3+).
++ // Pre-CDNA2 GFX9 (gfx906/gfx908) and GFX10 (RDNA1/2) are excluded.
++ auto v = getIsaVersion();
++ if (!((v.Major == 9 && (v.Minor > 0 || v.Stepping >= 0xa)) || v.Major >= 11))
+ return false;
+
+ Operation *reduxOp = op.getSingleCombiner();
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
index b79841ab5..afb80d9fb 100644
--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
@@ -12,6 +12,9 @@ ISAFamily deduceISAFamily(llvm::StringRef arch) {
if (kind == llvm::AMDGPU::GK_GFX1250)
return ISAFamily::GFX1250;
+ if (kind == llvm::AMDGPU::GK_GFX906)
+ return ISAFamily::GCN5_1;
+
// CDNA ISA cases
switch (kind) {
case llvm::AMDGPU::GK_GFX950:
@@ -41,6 +44,7 @@ ISAFamily deduceISAFamily(llvm::StringRef arch) {
bool supportsVDot(llvm::StringRef arch) {
switch (deduceISAFamily(arch)) {
+ case AMD::ISAFamily::GCN5_1:
case AMD::ISAFamily::CDNA1:
case AMD::ISAFamily::CDNA2:
case AMD::ISAFamily::CDNA3:
pkgs/development/python-modules/triton/default.nix
+30 -18
View File
@@ -42,17 +42,21 @@
cudaSupport ? config.cudaSupport,
}:
buildPythonPackage (finalAttrs: {
let
effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv;
in
buildPythonPackage.override { stdenv = effectiveStdenv; } (finalAttrs: {
pname = "triton";
version = "3.6.0";
version = "3.7.0";
pyproject = true;
__structuredAttrs = true;
# Remember to bump triton-llvm as well!
src = fetchFromGitHub {
owner = "triton-lang";
repo = "triton";
tag = "v${finalAttrs.version}";
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
hash = "sha256-FxbBY1lPq7765MqAPR7UljzPsmjOhKKbYExlKgeudew=";
};
patches = [
@@ -63,8 +67,9 @@ buildPythonPackage (finalAttrs: {
libcudaStubsDir =
if cudaSupport then "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" else null;
})
# Backport of https://github.com/triton-lang/triton/pull/9628 (does not apply cleanly)
./0005-add-gcn5-gfx906-target.patch
# Use our `cmakeFlags` instead and avoid downloading dependencies.
./inject-nix-cmakeFlags.patch
]
++ lib.optionals cudaSupport [
(replaceVars ./0003-nvidia-cudart-a-systempath.patch {
@@ -90,13 +95,6 @@ buildPythonPackage (finalAttrs: {
--replace-fail 'yield ("triton.profiler", "third_party/proton/proton")' 'pass' \
--replace-fail "curr_version.group(1) != version" "False"
''
# Use our `cmakeFlags` instead and avoid downloading dependencies
+ ''
substituteInPlace setup.py \
--replace-fail \
"cmake_args.extend(thirdparty_cmake_args)" \
"cmake_args.extend(thirdparty_cmake_args + os.environ.get('cmakeFlags', \"\").split())"
''
# Don't fetch googletest
+ ''
substituteInPlace cmake/AddTritonUnitTest.cmake \
@@ -104,6 +102,15 @@ buildPythonPackage (finalAttrs: {
--replace-fail "include(GoogleTest)" "find_package(GTest REQUIRED)"
''
# Hardcode the CC path so Triton's runtime JIT compilation doesn't break
# in environments without a compiler in PATH.
+ ''
substituteInPlace python/triton/runtime/build.py \
--replace-fail \
'cc = os.environ.get("CC")' \
'cc = os.environ.get("CC", "${lib.getExe' effectiveStdenv.cc "cc"}")'
''
# triton will try dlopening libcublas.so at runtime
+ lib.optionalString cudaSupport ''
substituteInPlace third_party/nvidia/include/cublas_instance.h \
@@ -135,7 +142,7 @@ buildPythonPackage (finalAttrs: {
# `find_package` is called with `NO_DEFAULT_PATH`
# https://cmake.org/cmake/help/latest/command/find_package.html
# https://github.com/triton-lang/triton/blob/c3c476f357f1e9768ea4e45aa5c17528449ab9ef/third_party/amd/CMakeLists.txt#L6
(lib.cmakeFeature "LLD_DIR" "${lib.getLib llvm}")
(lib.cmakeFeature "LLD_DIR" "${lib.getLib llvm}/lib/cmake/lld")
];
buildInputs = [
@@ -159,14 +166,19 @@ buildPythonPackage (finalAttrs: {
export MAX_JOBS="$NIX_BUILD_CORES"
'';
# `examples/plugins` (an MLIR example dialect plugin and a unit-test helper lib) is built
# unconditionally with the Python module and shipped into `triton/plugins/`.
# It is unused at runtime and keeps a forbidden RPATH reference to the build directory, which
# fails the fixup phase.
postInstall = ''
rm -rf "$out/${python.sitePackages}/triton/plugins"
'';
env = {
TRITON_BUILD_PROTON = "OFF";
TRITON_OFFLINE_BUILD = true;
}
// lib.optionalAttrs cudaSupport {
CC = lib.getExe' cudaPackages.backendStdenv.cc "cc";
CXX = lib.getExe' cudaPackages.backendStdenv.cc "c++";
NIX_CFLAGS_COMPILE = toString [
# Pybind11 started generating strange errors since python 3.12. Observed only in the CUDA branch.
# https://gist.github.com/SomeoneSerge/7d390b2b1313957c378e99ed57168219#file-gistfile0-txt-L1042
@@ -206,7 +218,7 @@ buildPythonPackage (finalAttrs: {
];
passthru = {
gpuCheck = stdenv.mkDerivation {
gpuCheck = effectiveStdenv.mkDerivation {
pname = "triton-pytest";
inherit (triton) version src;
@@ -370,7 +382,7 @@ buildPythonPackage (finalAttrs: {
if os.environ.get("HOME", None) == "/homeless-shelter":
os.environ["HOME"] = os.environ.get("TMPDIR", "/tmp")
if "CC" not in os.environ:
os.environ["CC"] = "${lib.getExe' cudaPackages.backendStdenv.cc "cc"}"
os.environ["CC"] = "${lib.getExe' effectiveStdenv.cc "cc"}"
torch.manual_seed(0)
size = 12345
x = torch.rand(size, device='cuda')
pkgs/development/python-modules/triton/inject-nix-cmakeFlags.patch
+16
View File
@@ -0,0 +1,16 @@
diff --git a/setup.py b/setup.py
index d8e0860291..bc09dada5a 100644
--- a/setup.py
+++ b/setup.py
@@ -476,6 +476,11 @@ class CMakeBuild(build_ext):
cmake_args.append("-DLLVM_EXTERNAL_LIT=" + lit_dir)
cmake_args.extend(thirdparty_cmake_args)
+ import json
+
+ nix_attrs_json = json.load(open(os.environ["NIX_ATTRS_JSON_FILE"]))
+ cmake_args.extend(thirdparty_cmake_args + nix_attrs_json.get("cmakeFlags", []))
+
# configuration
cfg = get_build_type()
build_args = ["--config", cfg]
pkgs/development/rocm-modules/rocprofiler-sdk/0009-rocprofiler-sdk-guard-ompt-auto-start.patch
+50
View File
@@ -0,0 +1,50 @@
From 502d2df0ddd84f0f2fa718b9a02ee0f01f36dac1 Mon Sep 17 00:00:00 2001
From: Luna Nova <git@lunnova.dev>
Date: Tue, 2 Jun 2026 01:31:05 -0700
Subject: [PATCH] [rocprofiler-sdk] don't auto-start OMPT tool before glog is
initialized
The OpenMP runtime resolves its weak ompt_start_tool to the strong symbol we
export and calls it during OpenMP init. That can fire from an unrelated
library's constructor (observed: OpenBLAS calling omp_get_num_places) while
librocprofiler-sdk is still being dlopen'd. init_logging() then dereferences
glog's gflags FLAGS_vmodule, whose backing pointer is still null until glog's
static initializers run, and the process segfaults on `import torch`.
The init order is not reliably fixable from here: a dependency cycle among the
ROCm libraries makes the loader run our constructors before libglog's, so a
flag set by our own constructor does not imply glog is ready. Instead gate the
tool start on glog itself, by testing whether gflags has populated FLAGS_vmodule
yet; reading the raw pointer is safe even before construction. It is read via an
asm label because the gflags header only exposes it as a reference, which the
compiler assumes is never null. Declining (nullptr) when glog is not yet ready
is a valid OMPT "no tool" response; deliberate profiling (rocprofv3 preloads the
tool early, before such constructors) is unaffected.
---
source/lib/rocprofiler-sdk/ompt.cpp | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/source/lib/rocprofiler-sdk/ompt.cpp b/source/lib/rocprofiler-sdk/ompt.cpp
index a9ccc57..3231087 100644
--- a/source/lib/rocprofiler-sdk/ompt.cpp
+++ b/source/lib/rocprofiler-sdk/ompt.cpp
@@ -199,6 +199,8 @@ finalize_ompt()
} // namespace ompt
} // namespace rocprofiler
+extern void* const rocprofiler_glog_vmodule_flag asm("_ZN3fLS13FLAGS_vmoduleB5cxx11E");
+
extern "C" {
rocprofiler_status_t
@@ -245,6 +247,8 @@ ompt_start_tool(unsigned int omp_version, const char* runtime_version) ROCPROFIL
ompt_start_tool_result_t*
ompt_start_tool(unsigned int omp_version, const char* runtime_version)
{
+ if(rocprofiler_glog_vmodule_flag == nullptr) return nullptr;
+
::rocprofiler::registration::init_logging();
::rocprofiler::registration::initialize();
return rocprofiler_ompt_start_tool(omp_version, runtime_version);
--
2.54.0
pkgs/development/rocm-modules/rocprofiler-sdk/default.nix
+7
View File
@@ -135,6 +135,13 @@ stdenv.mkDerivation (finalAttrs: {
./0006-rocprofiler-sdk-allow-using-system-otf2-dependency.patch
./0007-rocprofiler-sdk-Allow-using-system-json-dependency.patch
./0008-rocprofiler-sdk-stop-manually-setting-warning-flags-.patch
# Prevent a segfault on `import torch` (and other consumers that link
# librocprofiler-sdk for tracing) when OpenMP auto-discovers our
# ompt_start_tool before glog is initialized. Upstream incidentally fixes
# this by dropping glog in rocm-systems d80117889b3f
# "[rocprofiler-sdk] Migrate from glog to Abseil Logging (#4668)", which
# is too invasive (40 files + new abseil submodule) to backport.
./0009-rocprofiler-sdk-guard-ompt-auto-start.patch
];
postPatch = ''