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

Merged
TI
4mo ago
+632-33

Summary

Phase 2 of the CUDA plugin execution provider "no-cuDNN" work. It lets single last-axis ArgMax/ArgMin run through a small custom CUDA kernel instead of cuDNN, fixes LogSoftmax classification in the plugin adapter, and adds a non-throwing cuDNN handle accessor so reduction kernels can fall back gracefully when cuDNN is disabled.

Key Changes

AreaChange
reduction_functions.cu / .hNew arg_min_max_last_axis<TIn, IsArgMax> kernel (instantiated for half, float, double) that computes ArgMax/ArgMin indices over the last dimension of a row-major matrix without cuDNN.
reduction_ops.ccIn ReduceComputeCore, route a single last-axis ArgMax/ArgMin (CUDNN_REDUCE_TENSOR_FLATTENED_INDICES) to the custom kernel when shapes fit int; otherwise fall through to the existing cuDNN path. ReduceKernel::ComputeImpl now uses TryGetCudnnHandle.
cuda_kernel.h (native) / cuda_kernel_adapter.h (plugin)Add TryGetCudnnHandle, which returns the cuDNN handle when available and nullptr otherwise (instead of throwing at handle acquisition).
softmax.hDetect LogSoftmax from node.OpType() instead of info.GetKernelDef().OpName(), so the plugin EP adapter classifies it correctly.
test_cuda_plugin_ep.pyAdd LogSoftmax and ArgMin tests; drop the @requires_cudnn gate from ArgMax, ReduceMean, ReduceSum; reduce over the last axis to exercise the cuDNN-free paths.
docs/cuda_plugin_ep/QUICK_START.mdDrop ArgMax and reductions from the list of ops that still require cuDNN.

Correctness Notes

  • select_last_index == 1 is already rejected on the CUDA EP, so the kernel keeping the first matching index (strict > / <) is spec-correct for the supported case.
  • The custom path guards n > 0, returns early for m == 0, computes the row offset in int64_t, and only engages when m and n fit in int (gsl::narrow_cast); larger tensors fall back to cuDNN.

Testing

  • python -m pytest onnxruntime/test/python/transformers/test_cuda_plugin_ep.py -k "log_softmax or argmax or argmin or reduce_mean or reduce_sum"
  • Plugin no-cuDNN validation: bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin
  • onnxruntime_provider_test --gtest_filter='*Reduce*:*ArgM*'