[CUDA] Add cuDNN-free ArgMax/ArgMin/ReduceSum and fix LogSoftmax on plugin EP#29620
Open
tianleiwu wants to merge 8 commits into
Open
[CUDA] Add cuDNN-free ArgMax/ArgMin/ReduceSum and fix LogSoftmax on plugin EP#29620tianleiwu wants to merge 8 commits into
tianleiwu wants to merge 8 commits into
Conversation
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.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR is phase 2 of the CUDA Plugin EP “no-cuDNN” effort: it adds a cuDNN-free last-axis ArgMax/ArgMin path, fixes LogSoftmax detection for the plugin adapter, and introduces a non-throwing cuDNN handle accessor so kernels can run/fall back when cuDNN is disabled.
Changes:
- Added a custom CUDA kernel (
arg_min_max_last_axis) for last-axis ArgMax/ArgMin indices (half/float/double) and routed eligible cases to it. - Switched reduction kernels to use
TryGetCudnnHandle(nullable) instead of throwing at handle acquisition, and fixed LogSoftmax classification usingnode.OpType(). - Expanded CUDA plugin EP Python tests (LogSoftmax, ArgMin) and adjusted reduction tests to exercise last-axis non-cuDNN paths; updated plugin quick-start docs accordingly.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_cuda_plugin_ep.py | Adds LogSoftmax/ArgMin tests and updates ArgMax/Reduce* tests to run without cuDNN and exercise last-axis paths. |
| onnxruntime/core/providers/cuda/reduction/reduction_ops.cc | Routes last-axis ArgMax/ArgMin to the new kernel and switches reduction execution to use TryGetCudnnHandle. |
| onnxruntime/core/providers/cuda/reduction/reduction_functions.h | Declares the new last-axis ArgMin/ArgMax helper. |
| onnxruntime/core/providers/cuda/reduction/reduction_functions.cu | Implements and instantiates the cuDNN-free last-axis ArgMin/ArgMax kernel. |
| onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h | Adds TryGetCudnnHandle for the plugin adapter. |
| onnxruntime/core/providers/cuda/math/softmax.h | Fixes LogSoftmax detection using node.OpType() for correct plugin classification. |
| onnxruntime/core/providers/cuda/cuda_kernel.h | Adds TryGetCudnnHandle (nullable) to native CUDA kernels. |
| docs/cuda_plugin_ep/QUICK_START.md | Updates the list of ops requiring cuDNN for plugin no-cuDNN validation. |
Expand the Linux no-cuDNN CUDA EP smoke test from a single Add to also run LogSoftmax, ReduceMean, ReduceSum, ArgMax and ArgMin over the last axis with enable_cudnn=0, verifying the phase-2 cuDNN-free op paths execute on the regular CUDA EP. The Windows plugin no-cuDNN CI already exercises these ops via test_cuda_plugin_ep.py.
…in_max helper - TryGetCudnnHandle now treats a cudnnSetStream failure as "no handle" and returns nullptr (via non-throwing CUDNN_CALL) instead of throwing through CUDNN_CALL_THROW, so callers can fall back to a cuDNN-free path. - arg_min_max_last_axis guards n <= 0 in addition to m == 0, keeping the helper safe against out-of-bounds reads if reused without the caller's n > 0 check.
tianleiwu
changed the title
The manylinux build image sets USER onnxruntimedev (uid 1001), so the disposable smoke-test container ran as non-root. The new cuDNN-removal step deletes system libcudnn*.so under /usr/lib64 and rebuilds the dynamic linker cache with ldconfig, both of which require root, causing 'ldconfig: Can't create temporary cache file /etc/ld.so.cache~: Permission denied' and aborting under set -e. Run this throwaway container with --user root so the removal and ldconfig succeed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase 2 of the CUDA plugin execution provider "no-cuDNN" work. It lets single last-axis
ArgMax/ArgMinrun through a small custom CUDA kernel instead of cuDNN, fixesLogSoftmaxclassification 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
reduction_functions.cu/.harg_min_max_last_axis<TIn, IsArgMax>kernel (instantiated forhalf,float,double) that computes ArgMax/ArgMin indices over the last dimension of a row-major matrix without cuDNN.reduction_ops.ccReduceComputeCore, route a single last-axis ArgMax/ArgMin (CUDNN_REDUCE_TENSOR_FLATTENED_INDICES) to the custom kernel when shapes fitint; otherwise fall through to the existing cuDNN path.ReduceKernel::ComputeImplnow usesTryGetCudnnHandle.cuda_kernel.h(native) /cuda_kernel_adapter.h(plugin)TryGetCudnnHandle, which returns the cuDNN handle when available andnullptrotherwise (instead of throwing at handle acquisition).softmax.hLogSoftmaxfromnode.OpType()instead ofinfo.GetKernelDef().OpName(), so the plugin EP adapter classifies it correctly.test_cuda_plugin_ep.pyLogSoftmaxandArgMintests; drop the@requires_cudnngate fromArgMax,ReduceMean,ReduceSum; reduce over the last axis to exercise the cuDNN-free paths.docs/cuda_plugin_ep/QUICK_START.mdArgMaxand 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.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"bash .env/cuda_130_plugin_no_cudnn.sh --build --test_pluginonnxruntime_provider_test --gtest_filter='*Reduce*:*ArgM*'