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

Author: tianleiwu

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