[CUDA] Add cuDNN-free ArgMax/ArgMin/ReduceSum and fix LogSoftmax on plugin EP
Merged
TI
4mo ago+632+632+012345678901234567890123456789-33-33-01234567890123456789
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
| Area | Change |
|---|---|
reduction_functions.cu / .h | New 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.cc | In 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.h | Detect LogSoftmax from node.OpType() instead of info.GetKernelDef().OpName(), so the plugin EP adapter classifies it correctly. |
test_cuda_plugin_ep.py | Add 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.md | Drop ArgMax and reductions from the list of ops that still require cuDNN. |
Correctness Notes
select_last_index == 1is 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 form == 0, computes the row offset inint64_t, and only engages whenmandnfit inint(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*'