From cd5f511a8e6f1715869f98faa31cb72a22f20037 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 25 Jun 2026 16:56:48 +0000 Subject: [PATCH 1/7] [CUDA] Add cuDNN-free ArgMax/ArgMin and fix LogSoftmax on plugin EP Phase 2 of the CUDA plugin EP no-cuDNN work. - Add a custom arg_min_max_last_axis CUDA kernel and route single last-axis ArgMax/ArgMin through it in ReduceComputeCore, so these ops no longer require a cuDNN handle. select_last_index==1 is already rejected on CUDA, so keeping the first matching index is correct. - Add TryGetCudnnHandle (native and plugin adapter) that returns the cuDNN handle when available and nullptr otherwise, and use it in the reduction kernel so unsupported cuDNN paths degrade instead of throwing during handle acquisition. - Detect LogSoftmax from node.OpType() instead of the KernelDef op name so the plugin EP adapter classifies it correctly. - Enable and extend plugin EP tests: LogSoftmax and ArgMin coverage, drop the requires_cudnn gate from ArgMax/ReduceMean/ReduceSum, and reduce over the last axis to exercise the cuDNN-free paths. --- docs/cuda_plugin_ep/QUICK_START.md | 2 +- onnxruntime/core/providers/cuda/cuda_kernel.h | 4 ++ .../core/providers/cuda/math/softmax.h | 2 +- .../cuda/plugin/cuda_kernel_adapter.h | 14 ++++++ .../cuda/reduction/reduction_functions.cu | 45 +++++++++++++++++ .../cuda/reduction/reduction_functions.h | 4 ++ .../providers/cuda/reduction/reduction_ops.cc | 25 +++++++++- .../transformers/test_cuda_plugin_ep.py | 49 +++++++++++++++---- 8 files changed, 133 insertions(+), 12 deletions(-) diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index ab7b4308f13e5..d6593d8e22d18 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -27,7 +27,7 @@ For local Linux CUDA 13 validation, use the no-cuDNN helper script. It keeps `CU bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin ``` -The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, ArgMax, reductions, Einsum, and cuDNN-backed pooling paths. +The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, Einsum, and cuDNN-backed pooling paths. ## Minimum ONNX Runtime Version diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 85d629764da48..6152b5f888eb2 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -133,6 +133,10 @@ class CudaKernel : public OpKernel { return RequireCudnnHandle(GetCudnnHandle(static_cast(ctx->GetComputeStream()))); } + inline cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + return GetCudnnHandle(static_cast(ctx->GetComputeStream())); + } + static inline cudnnHandle_t GetCudnnHandle(onnxruntime::CudaStream* stream) { return stream ? stream->cudnn_handle_ : nullptr; } diff --git a/onnxruntime/core/providers/cuda/math/softmax.h b/onnxruntime/core/providers/cuda/math/softmax.h index c0c0818042c15..09b6eeec834b5 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.h +++ b/onnxruntime/core/providers/cuda/math/softmax.h @@ -47,7 +47,7 @@ class Softmax final : public CudaKernel { } } - log_softmax_ = info.GetKernelDef().OpName() == "LogSoftmax"; + log_softmax_ = node.OpType() == "LogSoftmax"; } Status ComputeInternal(OpKernelContext* context) const override; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index b1cc748be754a..eacefd468a40f 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -1070,6 +1070,20 @@ class CudaKernel : public OpKernel { return handle; } + cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + auto stream = Stream(ctx); + auto handle = GetCudnnHandle(stream); + if (handle != nullptr) { + return handle; + } + + handle = DefaultCudnnHandle(); + if (handle != nullptr && stream != nullptr) { + CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + } + return handle; + } + static cublasHandle_t GetCublasHandle(cudaStream_t s) { auto* sync = cuda_plugin::CudaSyncStream::FromCudaStream(s); return sync ? sync->GetCublasHandle() : nullptr; diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index ed97507f87641..514692dffe6ca 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -513,6 +513,51 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +namespace detail { +template +__global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, int m, int n) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= m) return; + + const int64_t row_offset = static_cast(row) * n; + TIn best_value = input[row_offset]; + int64_t best_index = 0; + for (int i = 1; i < n; ++i) { + const TIn value = input[row_offset + i]; + if constexpr (IsArgMax) { + if (value > best_value) { + best_value = value; + best_index = i; + } + } else { + if (value < best_value) { + best_value = value; + best_index = i; + } + } + } + + output[row] = best_index; +} +} // namespace detail + +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n) { + if (m == 0) return Status::OK(); + constexpr int block_size = 256; + const int grid_size = (m + block_size - 1) / block_size; + detail::arg_min_max_last_axis_kernel<<>>(input, output, m, n); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(T) \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n); \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n) +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(half); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(float); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(double); +#undef INSTANTIATE_ARG_MIN_MAX_LAST_AXIS + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 30b0afe4fe799..6a7cf98ca4970 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -103,6 +103,10 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); +/** Computes ArgMax/ArgMin indices over the last dimension in a row-major matrix. */ +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n); + /** Apply unary elementwise division. */ template void UnaryDiv(cudaStream_t stream, const T* input, T* output, T denominator, size_t count); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index a8019cda5c411..51de9f49eb71e 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -481,6 +481,29 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke } } + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES) { + if (axes.size() == 1) { + const int64_t rank = input_shape.NumDimensions(); + const int64_t axis = HandleNegativeAxis(axes[0], rank); + if (axis == rank - 1) { + const int64_t m = input_shape.SizeToDimension(axis); + const int64_t n = input_shape[axis]; + if (n > 0 && m <= std::numeric_limits::max() && n <= std::numeric_limits::max()) { + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + } + } + } + } + // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. // Therefore zeroing out the memory is required CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0, output.SizeInBytes(), stream)); @@ -785,7 +808,7 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const bool fast_reduction = fast_reduction_ && !ctx->GetUseDeterministicCompute(); return ReduceComputeCore(AllocatorPtr{}, this, *X, prepare_reduce_metadata, *Y, cudnn_reduce_op, axes, calculate_log_, calculate_sqt_, log_sum_exp_, fast_reduction, - Stream(ctx), GetComputeStream(ctx), GetCudnnHandle(ctx)); + Stream(ctx), GetComputeStream(ctx), TryGetCudnnHandle(ctx)); } #define SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(T) \ diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index caff34b1c79d4..60e91d5bb2546 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -1804,6 +1804,25 @@ def expected(f): result = _run_model_test(target_device, "Softmax", model, feed, expected) self.assertEqual(result, TEST_PASS, "Softmax test failed") + def test_op_log_softmax(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "LogSoftmax", + [("X", TensorProto.FLOAT, [2, 5])], + [("Y", TensorProto.FLOAT, [2, 5])], + attrs={"axis": 1}, + opset=13, + ) + feed = {"X": np.random.rand(2, 5).astype(np.float32)} + + def expected(f): + x = f["X"] + shifted = x - np.max(x, axis=1, keepdims=True) + return shifted - np.log(np.sum(np.exp(shifted), axis=1, keepdims=True)) + + result = _run_model_test(target_device, "LogSoftmax", model, feed, expected) + self.assertEqual(result, TEST_PASS, "LogSoftmax test failed") + def test_op_relu(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1900,7 +1919,6 @@ def test_op_flatten(self): result = _run_model_test(target_device, "Flatten", model, feed, lambda f: f["X"].reshape(2, 12)) self.assertEqual(result, TEST_PASS, "Flatten test failed") - @requires_cudnn def test_op_argmax(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1916,6 +1934,21 @@ def test_op_argmax(self): ) self.assertEqual(result, TEST_PASS, "ArgMax test failed") + def test_op_argmin(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ArgMin", + [("X", TensorProto.FLOAT, [3, 5])], + [("Y", TensorProto.INT64, [3, 1])], + attrs={"axis": 1, "keepdims": 1}, + opset=13, + ) + feed = {"X": np.random.rand(3, 5).astype(np.float32)} + result = _run_model_test( + target_device, "ArgMin", model, feed, lambda f: np.argmin(f["X"], axis=1).reshape(3, 1) + ) + self.assertEqual(result, TEST_PASS, "ArgMin test failed") + def test_op_topk(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2030,37 +2063,35 @@ def expected(f): result = _run_model_test(target_device, "ConvTranspose", model, feed, expected) self.assertEqual(result, TEST_PASS, "ConvTranspose test failed") - @requires_cudnn def test_op_reduce_mean(self): target_device = get_cuda_plugin_device() model = _make_simple_model( "ReduceMean", [("X", TensorProto.FLOAT, [3, 4, 5])], - [("Y", TensorProto.FLOAT, [3, 1, 5])], - attrs={"axes": [1], "keepdims": 1}, + [("Y", TensorProto.FLOAT, [3, 4, 1])], + attrs={"axes": [2], "keepdims": 1}, opset=13, ) feed = {"X": np.random.rand(3, 4, 5).astype(np.float32)} result = _run_model_test( - target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=1, keepdims=True) + target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=2, keepdims=True) ) self.assertEqual(result, TEST_PASS, "ReduceMean test failed") - @requires_cudnn def test_op_reduce_sum(self): target_device = get_cuda_plugin_device() model = _make_simple_model( "ReduceSum", [("X", TensorProto.FLOAT, [3, 4, 5]), ("axes", TensorProto.INT64, [1])], - [("Y", TensorProto.FLOAT, [3, 1, 5])], + [("Y", TensorProto.FLOAT, [3, 4, 1])], attrs={"keepdims": 1}, opset=13, ) - axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [1]) + axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [2]) model.graph.initializer.append(axes_init) feed = {"X": np.random.rand(3, 4, 5).astype(np.float32)} result = _run_model_test( - target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=1, keepdims=True) + target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=2, keepdims=True) ) self.assertEqual(result, TEST_PASS, "ReduceSum test failed") From 8a9e6e57a296037d275b9fe20c45142a8eecf7d2 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 18:27:04 +0000 Subject: [PATCH 2/7] [CUDA] Cover cuDNN-free ops in Linux no-cuDNN smoke test Expand the Linux no-cuDNN CUDA EP smoke test from a single Add to also run LogSoftmax, ReduceMean, ReduceSum, ArgMax and ArgMin over the last axis with enable_cudnn=0, verifying the phase-2 cuDNN-free op paths execute on the regular CUDA EP. The Windows plugin no-cuDNN CI already exercises these ops via test_cuda_plugin_ep.py. --- .github/workflows/linux_cuda_no_cudnn.yml | 118 +++++++++++++++++++--- 1 file changed, 104 insertions(+), 14 deletions(-) diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index abe92589aa705..a3bcb4e410cb0 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -93,7 +93,7 @@ jobs: ldd /build/Release/Release/libonnxruntime_providers_cuda.so | tee /tmp/ldd.txt ! grep -i cudnn /tmp/ldd.txt' - - name: Run no-cuDNN CUDA EP smoke test + - name: Run no-cuDNN CUDA EP op smoke test run: | docker run --rm --gpus all \ -v "${{ runner.temp }}/Release:/build/Release" \ @@ -111,21 +111,111 @@ jobs: python -m pip install --no-cache-dir --force-reinstall --no-deps numpy onnx "$WHEEL_PATH" python - <<"PY" import numpy as np - import onnx import onnxruntime as ort from onnx import TensorProto, helper - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]) - y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3]) - node = helper.make_node("Add", ["x", "x"], ["y"]) - graph = helper.make_graph([node], "cuda_no_cudnn_smoke", [x], [y]) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) - model.ir_version = 10 - providers = [("CUDAExecutionProvider", {"enable_cudnn": "0"})] - sess = ort.InferenceSession(model.SerializeToString(), providers=providers) - data = np.arange(6, dtype=np.float32).reshape(2, 3) - result = sess.run(None, {"x": data})[0] - np.testing.assert_allclose(result, data + data) - print("CUDA no-cuDNN smoke test passed") + + + def run_op(name, model, feeds, expected, rtol=1e-4, atol=1e-4): + model.ir_version = 10 + sess = ort.InferenceSession(model.SerializeToString(), providers=providers) + got = sess.run(None, feeds)[0] + np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol) + print("[no-cuDNN] " + name + " passed") + + + def make_model(nodes, inputs, outputs, opset, initializers=None): + graph = helper.make_graph(nodes, "no_cudnn_smoke", inputs, outputs, initializer=initializers or []) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + + def f32_in(shape): + return helper.make_tensor_value_info("x", TensorProto.FLOAT, shape) + + + data = np.random.rand(2, 3).astype(np.float32) + + # Elementwise baseline (does not depend on cuDNN). + run_op( + "Add", + make_model( + [helper.make_node("Add", ["x", "x"], ["y"])], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 21, + ), + {"x": data}, + data + data, + ) + + # LogSoftmax over the last axis (softmax kernel, no cuDNN). + shifted = data - data.max(axis=1, keepdims=True) + log_softmax = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True)) + run_op( + "LogSoftmax", + make_model( + [helper.make_node("LogSoftmax", ["x"], ["y"], axis=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 13, + ), + {"x": data}, + log_softmax, + ) + + # ReduceMean over the last axis (matrix-reduction path, no cuDNN). + run_op( + "ReduceMean", + make_model( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[1], keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1])], + 13, + ), + {"x": data}, + data.mean(axis=1, keepdims=True), + ) + + # ReduceSum over the last axis (axes supplied as an input in opset 13). + run_op( + "ReduceSum", + make_model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1])], + 13, + initializers=[helper.make_tensor("axes", TensorProto.INT64, [1], [1])], + ), + {"x": data}, + data.sum(axis=1, keepdims=True), + ) + + # ArgMax over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMax", + make_model( + [helper.make_node("ArgMax", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmax(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + # ArgMin over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMin", + make_model( + [helper.make_node("ArgMin", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmin(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + print("CUDA no-cuDNN op smoke tests passed") PY' From 4284ee13d1cc24e00ebd9cf3bb5a631f31ce44e7 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 18:44:28 +0000 Subject: [PATCH 3/7] [CUDA] Address review: non-throwing TryGetCudnnHandle and guard arg_min_max helper - TryGetCudnnHandle now treats a cudnnSetStream failure as "no handle" and returns nullptr (via non-throwing CUDNN_CALL) instead of throwing through CUDNN_CALL_THROW, so callers can fall back to a cuDNN-free path. - arg_min_max_last_axis guards n <= 0 in addition to m == 0, keeping the helper safe against out-of-bounds reads if reused without the caller's n > 0 check. --- .../core/providers/cuda/plugin/cuda_kernel_adapter.h | 6 +++++- .../core/providers/cuda/reduction/reduction_functions.cu | 3 ++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index eacefd468a40f..1c5cda879327e 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -1079,7 +1079,11 @@ class CudaKernel : public OpKernel { handle = DefaultCudnnHandle(); if (handle != nullptr && stream != nullptr) { - CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + // Keep this accessor non-throwing: if the stream cannot be bound, treat it as "no cuDNN handle" + // so callers can fall back to a cuDNN-free path instead of failing. + if (!CUDNN_CALL(cudnnSetStream(handle, stream)).IsOK()) { + return nullptr; + } } return handle; } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 514692dffe6ca..4ff9457bd078a 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -543,7 +543,8 @@ __global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, template Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n) { - if (m == 0) return Status::OK(); + // The kernel reads input[row_offset] unconditionally, so a non-empty reduction axis is required. + if (m == 0 || n <= 0) return Status::OK(); constexpr int block_size = 256; const int grid_size = (m + block_size - 1) / block_size; detail::arg_min_max_last_axis_kernel<<>>(input, output, m, n); From faf1e40db3ffbaad076b5ad19766ffef348740fb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 21:07:21 +0000 Subject: [PATCH 4/7] Bind cuDNN handle to default stream (0) unconditionally --- .../core/providers/cuda/plugin/cuda_kernel_adapter.h | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index 1c5cda879327e..7c6cdd0e4dc80 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -1064,9 +1064,10 @@ class CudaKernel : public OpKernel { std::string("cuDNN is unavailable or disabled for CUDA Plugin Execution Provider: ") + onnxruntime::cuda::CudnnLibrary::Get().Error())); } - if (handle != nullptr && stream != nullptr) { - CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); - } + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. + CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); return handle; } @@ -1078,7 +1079,10 @@ class CudaKernel : public OpKernel { } handle = DefaultCudnnHandle(); - if (handle != nullptr && stream != nullptr) { + if (handle != nullptr) { + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. // Keep this accessor non-throwing: if the stream cannot be bound, treat it as "no cuDNN handle" // so callers can fall back to a cuDNN-free path instead of failing. if (!CUDNN_CALL(cudnnSetStream(handle, stream)).IsOK()) { From e78fc34b1a79b86dac23163f7c4b98fd01a98cd2 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 21:31:18 +0000 Subject: [PATCH 5/7] Fix reduce sum --- .github/workflows/linux_cuda_no_cudnn.yml | 34 +++- .../cuda/reduction/reduction_functions.cu | 186 ++++++++++++++++++ .../cuda/reduction/reduction_functions.h | 9 + .../providers/cuda/reduction/reduction_ops.cc | 29 ++- .../cpu/reduction/reduction_ops_test.cc | 25 +++ .../test_cases/reduction_functions_test.cc | 93 +++++++++ .../transformers/test_cuda_plugin_ep.py | 6 +- 7 files changed, 366 insertions(+), 16 deletions(-) diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index a3bcb4e410cb0..5adcf016ea47d 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -102,6 +102,25 @@ jobs: PATH=/opt/python/cp312-cp312/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-} export PATH LD_LIBRARY_PATH + + # The build image contains cuDNN. Remove its runtime libraries from this disposable + # test container so the smoke test proves that the provider works when cuDNN is + # physically unavailable, rather than merely relying on enable_cudnn=0. + for root in /usr /opt /lib /lib64; do + if [ -e "$root" ]; then + find "$root" -name "libcudnn*.so*" -print -delete 2>/dev/null || true + fi + done + ldconfig + if ldconfig -p | grep -qi cudnn; then + echo "cuDNN remains available in the dynamic linker cache" >&2 + exit 1 + fi + if find /usr /opt /lib /lib64 -name "libcudnn*.so*" -print -quit 2>/dev/null | grep -q .; then + echo "cuDNN runtime libraries remain in the no-cuDNN test container" >&2 + exit 1 + fi + WHEEL_PATH=$(find /build/Release/Release/dist -type f -name "onnxruntime_gpu-*.whl" | head -n 1) if [ -z "$WHEEL_PATH" ]; then echo "No built onnxruntime GPU wheel found under /build/Release/Release/dist" >&2 @@ -119,7 +138,9 @@ jobs: def run_op(name, model, feeds, expected, rtol=1e-4, atol=1e-4): model.ir_version = 10 - sess = ort.InferenceSession(model.SerializeToString(), providers=providers) + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + sess = ort.InferenceSession(model.SerializeToString(), sess_options=options, providers=providers) got = sess.run(None, feeds)[0] np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol) print("[no-cuDNN] " + name + " passed") @@ -177,18 +198,19 @@ jobs: data.mean(axis=1, keepdims=True), ) - # ReduceSum over the last axis (axes supplied as an input in opset 13). + # ReduceSum over a middle axis, which requires the general no-cuDNN path. + reduce_sum_data = np.random.rand(2, 3, 4).astype(np.float32) run_op( "ReduceSum", make_model( [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=1)], - [f32_in([2, 3])], - [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1])], + [f32_in([2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1, 4])], 13, initializers=[helper.make_tensor("axes", TensorProto.INT64, [1], [1])], ), - {"x": data}, - data.sum(axis=1, keepdims=True), + {"x": reduce_sum_data}, + reduce_sum_data.sum(axis=1, keepdims=True), ) # ArgMax over the last axis (custom cuDNN-free kernel). diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 4ff9457bd078a..afc2d9b061a55 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -4,6 +4,8 @@ #include "core/providers/cuda/reduction/reduction_functions.h" #include +#include +#include #include #include @@ -513,6 +515,190 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +namespace detail { +constexpr int kMaxReduceRank = 16; + +struct ReduceSumNdMetadata { + int output_segment_count{}; + int reduction_segment_count{}; + int64_t output_segment_sizes[kMaxReduceRank]{}; + int64_t output_segment_strides[kMaxReduceRank]{}; + int64_t reduction_segment_sizes[kMaxReduceRank]{}; + int64_t reduction_segment_strides[kMaxReduceRank]{}; + int64_t output_count{}; + int64_t reduction_count{}; +}; + +template +struct SumState { + T sum{}; + + __device__ __forceinline__ void Add(T value) { sum += value; } + __device__ __forceinline__ T Result() const { return sum; } +}; + +template <> +struct SumState { + double sum{}; + double correction{}; + + // Neumaier compensation is used instead of Kahan so that independently accumulated + // thread partials can be merged without losing a small value between large values. + __device__ __forceinline__ void Add(double value) { + const double next = __dadd_rn(sum, value); + const double error = fabs(sum) >= fabs(value) + ? __dadd_rn(__dsub_rn(sum, next), value) + : __dadd_rn(__dsub_rn(value, next), sum); + correction = __dadd_rn(correction, error); + sum = next; + } + + __device__ __forceinline__ double Result() const { return __dadd_rn(sum, correction); } +}; + +template +struct MergeSumState { + __device__ __forceinline__ SumState operator()(SumState lhs, const SumState& rhs) const { + lhs.Add(rhs.sum); + if constexpr (std::is_same_v) { + lhs.Add(rhs.correction); + } + return lhs; + } +}; + +template +__device__ __forceinline__ T CastReduceSumResult(TAccum value) { + if constexpr (std::is_integral_v) { + const double value_as_double = static_cast(value); + const double max_value = static_cast(std::numeric_limits::max()); + const double min_value = static_cast(std::numeric_limits::min()); + if (value_as_double >= max_value) return std::numeric_limits::max(); + if (value_as_double <= min_value) return std::numeric_limits::min(); + } + return static_cast(value); +} + +template +__global__ void reduce_sum_nd_kernel(const T* input, T* output, ReduceSumNdMetadata metadata) { + using TAccum = std::conditional_t, double, AccumulationType_t>; + using BlockReduce = cub::BlockReduce, BlockSize>; + __shared__ typename BlockReduce::TempStorage reduce_storage; + __shared__ int64_t input_base; + + // One cooperative block reduces each output. Grid-striding keeps the launch bounded for + // large outputs while retaining enough blocks to saturate the device. + for (int64_t output_index = blockIdx.x; + output_index < metadata.output_count; + output_index += gridDim.x) { + if (threadIdx.x == 0) { + int64_t remaining = output_index; + int64_t base = 0; + for (int segment = metadata.output_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.output_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.output_segment_sizes[segment]; + base += coordinate * metadata.output_segment_strides[segment]; + } + input_base = base; + } + __syncthreads(); + + SumState thread_sum{}; + for (int64_t reduction_index = threadIdx.x; reduction_index < metadata.reduction_count; + reduction_index += BlockSize) { + int64_t remaining = reduction_index; + int64_t input_index = input_base; + // Adjacent reduced dimensions are collapsed on the host, so this loop performs + // divisions only at reduced/non-reduced boundaries rather than once per rank. + for (int segment = metadata.reduction_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.reduction_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.reduction_segment_sizes[segment]; + input_index += coordinate * metadata.reduction_segment_strides[segment]; + } + thread_sum.Add(static_cast(input[input_index])); + } + + const SumState block_sum = BlockReduce(reduce_storage).Reduce(thread_sum, MergeSumState{}); + if (threadIdx.x == 0) { + output[output_index] = CastReduceSumResult(block_sum.Result()); + } + __syncthreads(); // reduce_storage is reused by the next grid-stride iteration. + } +} +} // namespace detail + +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes) { + ORT_RETURN_IF_NOT(dims.size() <= detail::kMaxReduceRank, + "The general CUDA ReduceSum kernel supports ranks up to ", detail::kMaxReduceRank, "."); + + detail::ReduceSumNdMetadata metadata; + const int rank = gsl::narrow_cast(dims.size()); + std::array strides{}; + SafeInt stride = 1; + for (int axis = rank - 1; axis >= 0; --axis) { + ORT_RETURN_IF_NOT(dims[axis] > 0, "ReduceSum dimensions must be positive."); + strides[axis] = static_cast(stride); + stride *= dims[axis]; + } + + std::array reduced{}; + if (axes.empty()) { + for (int axis = 0; axis < rank; ++axis) reduced[axis] = true; + } else { + for (int64_t axis : axes) { + if (axis < 0) axis += rank; + ORT_RETURN_IF_NOT(axis >= 0 && axis < rank, "ReduceSum axis is out of range."); + ORT_RETURN_IF_NOT(!reduced[axis], "ReduceSum axes must not contain duplicates."); + reduced[axis] = true; + } + } + + SafeInt output_count = 1; + SafeInt reduction_count = 1; + for (int axis = 0; axis < rank;) { + const bool is_reduced = reduced[axis]; + SafeInt segment_size = 1; + int last_axis = axis; + do { + segment_size *= dims[axis]; + last_axis = axis++; + } while (axis < rank && reduced[axis] == is_reduced); + + if (is_reduced) { + const int segment = metadata.reduction_segment_count++; + metadata.reduction_segment_sizes[segment] = static_cast(segment_size); + metadata.reduction_segment_strides[segment] = strides[last_axis]; + reduction_count *= static_cast(segment_size); + } else { + const int segment = metadata.output_segment_count++; + metadata.output_segment_sizes[segment] = static_cast(segment_size); + metadata.output_segment_strides[segment] = strides[last_axis]; + output_count *= static_cast(segment_size); + } + } + metadata.output_count = static_cast(output_count); + metadata.reduction_count = static_cast(reduction_count); + + constexpr int block_size = 256; + constexpr int max_blocks = 65535; + const int grid_size = static_cast(std::min(max_blocks, metadata.output_count)); + detail::reduce_sum_nd_kernel<<>>(input, output, metadata); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_REDUCE_SUM_ND(T) \ + template Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, \ + gsl::span dims, gsl::span axes) +INSTANTIATE_REDUCE_SUM_ND(half); +INSTANTIATE_REDUCE_SUM_ND(float); +INSTANTIATE_REDUCE_SUM_ND(double); +INSTANTIATE_REDUCE_SUM_ND(BFloat16); +INSTANTIATE_REDUCE_SUM_ND(int32_t); +INSTANTIATE_REDUCE_SUM_ND(int64_t); +#undef INSTANTIATE_REDUCE_SUM_ND + namespace detail { template __global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, int m, int n) { diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 6a7cf98ca4970..2e73a73a353f7 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -103,6 +103,15 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); +/** + * Computes ReduceSum for an arbitrary set of axes. This is the general fallback + * for axis layouts that cannot be flattened into one of the optimized matrix + * reductions above. + */ +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes); + /** Computes ArgMax/ArgMin indices over the last dimension in a row-major matrix. */ template Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 51de9f49eb71e..ed6195c0bbb8b 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -299,10 +299,6 @@ Status PrepareForReduce(const Tensor* X, const int64_t rank = gsl::narrow(input_shape.NumDimensions()); prepare_reduce_metadata.input_count = input_shape.Size(); - if (rank > 8) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuDNN only supports up to 8-D tensors in reduction"); - } - const auto input_dims = input_shape.GetDims(); std::vector reduced(rank, false); if (axes.size() > 0) { @@ -504,6 +500,18 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke } } + // Preserve the optimized matrix reductions above and the existing cuDNN path when available. + // Without cuDNN, use a general CUDA kernel for plain ReduceSum axis layouts that cannot be + // represented as a contiguous matrix reduction. + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES) { + if ((cudnn_handle == nullptr || input_shape.NumDimensions() > 8) && + cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && + !calculate_log && !calculate_sqt && !log_sum_exp) { + return reduce_sum_nd(stream, reinterpret_cast(input.Data()), + reinterpret_cast(output.MutableData()), input_shape.GetDims(), axes); + } + } + // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. // Therefore zeroing out the memory is required CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0, output.SizeInBytes(), stream)); @@ -820,9 +828,8 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const Tensor* X = ctx->Input(0); \ TensorShapeVector axes; \ size_t num_inputs = ctx->InputCount(); \ - if (num_inputs == 2) { \ - const Tensor* axes_tensor = ctx->Input(1); \ - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); \ + const Tensor* axes_tensor = num_inputs == 2 ? ctx->Input(1) : nullptr; \ + if (axes_tensor != nullptr) { \ ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); \ auto nDims = static_cast(axes_tensor->Shape()[0]); \ const auto* data = axes_tensor->Data(); \ @@ -881,6 +888,14 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe return Status::OK(); \ } \ \ + if constexpr (std::is_same_v || std::is_same_v) { \ + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && \ + !calculate_log_ && !calculate_sqt_ && !log_sum_exp_) { \ + return reduce_sum_nd(Stream(ctx), reinterpret_cast(X->Data()), \ + reinterpret_cast(Y->MutableData()), X->Shape().GetDims(), axes); \ + } \ + } \ + \ CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream(ctx))); \ \ size_t indices_bytes = 0; \ diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 95f5ae30016e2..13d6016d09de9 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -2863,6 +2863,31 @@ TEST(ReductionOpTest, ReduceSum_int64) { test.Run(); } +#if defined(USE_CUDA) +TEST(ReductionOpTest, ReduceSum_int64_omitted_optional_axes) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3}, {1, 2, 3}); + test.AddOptionalInputEdge(); + test.AddOutput("reduced", {}, {6}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(ReductionOpTest, ReduceSum_int64_cancellation) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + const int64_t large = int64_t{1} << 53; + test.AddInput("data", {3}, {large, 1, -large}); + test.AddInput("axes", {1}, {0}); + test.AddOutput("reduced", {}, {1}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(ReductionOpTest, ReduceSum_default_axes_keepdims) { OpTester test("ReduceSum"); test.AddAttribute("keepdims", (int64_t)1); diff --git a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc index 593255b9e9c23..410527f0a6ea6 100644 --- a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc @@ -3,7 +3,9 @@ #include "gtest/gtest.h" +#include #include +#include #include "core/providers/cuda/shared_inc/cuda_utils.h" #include "core/common/optional.h" @@ -238,6 +240,97 @@ TEST(ReductionFunctionsTest, ReduceColumnsToColumnRepeated) { TestReduceColumnsToColumnRepeated(17, 8192, 100, 2e-4f); } +TEST(ReductionFunctionsTest, ReduceSumNdMiddleAndMultipleAxes) { + const std::vector dims{2, 3, 4, 2}; + const std::vector axes{1, 3}; + std::vector input(48); + std::iota(input.begin(), input.end(), 1.0f); + std::vector expected(8, 0.0f); + for (int64_t d0 = 0; d0 < dims[0]; ++d0) { + for (int64_t d1 = 0; d1 < dims[1]; ++d1) { + for (int64_t d2 = 0; d2 < dims[2]; ++d2) { + for (int64_t d3 = 0; d3 < dims[3]; ++d3) { + expected[d0 * dims[2] + d2] += input[((d0 * dims[1] + d1) * dims[2] + d2) * dims[3] + d3]; + } + } + } + } + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 1e-6f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdIntegerSaturation) { + const std::vector dims{2, 3, 2}; + const std::vector axes{1}; + const int32_t big = 1'100'000'000; + const std::vector input(12, big); + const std::vector expected(4, std::numeric_limits::max()); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdLargeReductionSmallOutput) { + const std::vector dims{2, 131072, 3}; + const std::vector axes{1}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(6, 131072.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdInt64Cancellation) { + const std::vector dims{1, 3, 1}; + const std::vector axes{1}; + const int64_t large = int64_t{1} << 53; + const std::vector input{large, 1, -large}; + const std::vector expected{1}; + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int64_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int64_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdRank9) { + const std::vector dims{2, 2, 2, 2, 2, 2, 2, 2, 2}; + const std::vector axes{1, 3, 5, 7}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(32, 16.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + TEST(ReductionFunctionsTest, BufferOffsets) { const int m = 2048; const int n = 1024; diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index 60e91d5bb2546..b8f5bbfc926ea 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -2083,15 +2083,15 @@ def test_op_reduce_sum(self): model = _make_simple_model( "ReduceSum", [("X", TensorProto.FLOAT, [3, 4, 5]), ("axes", TensorProto.INT64, [1])], - [("Y", TensorProto.FLOAT, [3, 4, 1])], + [("Y", TensorProto.FLOAT, [3, 1, 5])], attrs={"keepdims": 1}, opset=13, ) - axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [2]) + axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [1]) model.graph.initializer.append(axes_init) feed = {"X": np.random.rand(3, 4, 5).astype(np.float32)} result = _run_model_test( - target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=2, keepdims=True) + target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=1, keepdims=True) ) self.assertEqual(result, TEST_PASS, "ReduceSum test failed") From 614f12b44409d41f2dcfd3fe638b015ca2c288eb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 21:44:46 +0000 Subject: [PATCH 6/7] add int ReduceSum test cases --- .../transformers/test_cuda_plugin_ep.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index b8f5bbfc926ea..c95a03dc50d8e 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -2095,6 +2095,32 @@ def test_op_reduce_sum(self): ) self.assertEqual(result, TEST_PASS, "ReduceSum test failed") + def _run_reduce_sum_integer_last_axis(self, onnx_dtype, np_dtype): + # Mirrors the qwen attention_mask usage: a rank-2 integer ReduceSum over the last axis. + # Integer ReduceSum takes a specialized path that does not use the float matrix fast path, + # so without cuDNN it must fall back to the general native kernel. Covering it here keeps the + # no-cuDNN CI honest for the exact layout that broke real models. + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ReduceSum", + [("X", onnx_dtype, [2, 8]), ("axes", TensorProto.INT64, [1])], + [("Y", onnx_dtype, [2, 1])], + attrs={"keepdims": 1}, + opset=13, + ) + axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [1]) + model.graph.initializer.append(axes_init) + feed = {"X": np.arange(16, dtype=np_dtype).reshape(2, 8)} + return _run_model_test(target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=1, keepdims=True)) + + def test_op_reduce_sum_int64_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT64, np.int64) + self.assertEqual(result, TEST_PASS, "ReduceSum int64 test failed") + + def test_op_reduce_sum_int32_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT32, np.int32) + self.assertEqual(result, TEST_PASS, "ReduceSum int32 test failed") + def test_op_gather_nd(self): target_device = get_cuda_plugin_device() model = _make_simple_model( From fee704e80ee9a6c1dd877b5536a804ab46ff95d9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 23:27:29 +0000 Subject: [PATCH 7/7] ci: run no-cuDNN smoke test container as root The manylinux build image sets USER onnxruntimedev (uid 1001), so the disposable smoke-test container ran as non-root. The new cuDNN-removal step deletes system libcudnn*.so under /usr/lib64 and rebuilds the dynamic linker cache with ldconfig, both of which require root, causing 'ldconfig: Can't create temporary cache file /etc/ld.so.cache~: Permission denied' and aborting under set -e. Run this throwaway container with --user root so the removal and ldconfig succeed. --- .github/workflows/linux_cuda_no_cudnn.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index 5adcf016ea47d..47f8b33f15bbb 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -95,7 +95,10 @@ jobs: - name: Run no-cuDNN CUDA EP op smoke test run: | - docker run --rm --gpus all \ + # The build image runs as a non-root user by default, but this step + # deletes the system cuDNN libraries and rebuilds the dynamic linker + # cache (both require root). Run this disposable test container as root. + docker run --rm --gpus all --user root \ -v "${{ runner.temp }}/Release:/build/Release" \ "${{ steps.build_docker_image_step.outputs.full-image-name }}" \ bash -lc 'set -e