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

.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'

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

@@ -1064,8 +1064,30 @@ 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;
+  }
+
+  cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const {
+    auto stream = Stream(ctx);
+    auto handle = GetCudnnHandle(stream);
+    if (handle != nullptr) {
+      return handle;
+    }
+
+    handle = DefaultCudnnHandle();
+    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()) {
+        return nullptr;
+      }
     }
     return handle;
   }

onnxruntime/core/providers/cuda/reduction/reduction_functions.cu

@@ -513,6 +513,52 @@ 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) {
+  // 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<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 @@
     }
   }
 
+  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 @@
   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")