[CUDA] Add cuDNN-free ArgMax/ArgMin/ReduceSum and fix LogSoftmax on plugin EP

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
 

onnxruntime/core/providers/cuda/cuda_kernel.h

@@ -133,6 +133,10 @@ class CudaKernel : public OpKernel {
     return RequireCudnnHandle(GetCudnnHandle(static_cast<CudaStream*>(ctx->GetComputeStream())));
   }
 
+  inline cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const {
+    return GetCudnnHandle(static_cast<CudaStream*>(ctx->GetComputeStream()));
+  }
+
   static inline cudnnHandle_t GetCudnnHandle(onnxruntime::CudaStream* stream) {
     return stream ? stream->cudnn_handle_ : nullptr;
   }

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;

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;

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 <typename TIn, bool IsArgMax>
+__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<int64_t>(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 <typename TIn, bool IsArgMax>
+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<TIn, IsArgMax><<<grid_size, block_size, 0, stream>>>(input, output, m, n);
+  return CUDA_CALL(cudaGetLastError());
+}
+
+#define INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(T)                                                                          \
+  template Status arg_min_max_last_axis<T, true>(cudaStream_t stream, const T* input, int64_t* output, int m, int n); \
+  template Status arg_min_max_last_axis<T, false>(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
 

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 <typename TIn, typename TOut>
 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 <typename TIn, bool IsArgMax>
+Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n);
+
 /** Apply unary elementwise division. */
 template <typename T>
 void UnaryDiv(cudaStream_t stream, const T* input, T* output, T denominator, size_t count);

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<int>::max() && n <= std::numeric_limits<int>::max()) {
+          if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) {
+            return arg_min_max_last_axis<CudaT, true>(stream, reinterpret_cast<const CudaT*>(input.Data<T>()),
+                                                      output.MutableData<int64_t>(), gsl::narrow_cast<int>(m),
+                                                      gsl::narrow_cast<int>(n));
+          }
+          if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) {
+            return arg_min_max_last_axis<CudaT, false>(stream, reinterpret_cast<const CudaT*>(input.Data<T>()),
+                                                       output.MutableData<int64_t>(), gsl::narrow_cast<int>(m),
+                                                       gsl::narrow_cast<int>(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<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
   const bool fast_reduction = fast_reduction_ && !ctx->GetUseDeterministicCompute();
   return ReduceComputeCore<T, ReduceTensorIndices>(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)                                                                           \

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")