microsoft / microsoft/onnxruntime
CUDA AveragePool fails with CUDNN_STATUS_NOT_SUPPORTED when batch dimension reaches 65536
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
### Description
`AveragePool` executed by `CUDAExecutionProvider` fails with
`CUDNN_STATUS_NOT_SUPPORTED` when the effective batch dimension reaches 65536.
The same operation succeeds:
- With `N=65535`
- In PyTorch CUDA with `N=196608`
- After replacing `AveragePool` with `Reshape + ReduceMean`
This appears to be specific to ONNX Runtime's cuDNN pooling path.
### Environment
- OS: Ubuntu 24.04
- GPU: NVIDIA GeForce RTX 5090
- NVIDIA driver: 575.64.03
- ONNX Runtime: 1.23.2
- CUDA: 12.9
- cuDNN: 9.10.2
- Python: 3.10
- ONNX opset: 18
The problem was also reproduced with newer ONNX Runtime/cuDNN combinations
and NVIDIA driver 610.43.02.
### Expected behavior
AveragePool should execute for this valid tensor shape, or ONNX Runtime
should provide a non-cuDNN CUDA fallback when cudnnPoolingForward returns
CUDNN_STATUS_NOT_SUPPORTED.
### To reproduce
### Minimal reproduction
```python
import numpy as np
import onnx
import onnxruntime as ort
from onnx import TensorProto, helper
def make_model(path, n):
input_shape = [n, 1, 1, 512]
output_shape = [n, 1, 1, 256]
graph = helper.make_graph(
[
helper.make_node(
"AveragePool",
["input"],
["output"],
name="node_avg_pool2d",
kernel_shape=[1, 2],
strides=[1, 2],
pads=[0, 0, 0, 0],
ceil_mode=0,
count_include_pad=1,
)
],
"large_batch_average_pool",
[
helper.make_tensor_value_info(
"input", TensorProto.FLOAT, input_shape
)
],
[
helper.make_tensor_value_info(
"output", TensorProto.FLOAT, output_shape
)
],
)
model = helper.make_model(
graph,
opset_imports=[helper.make_opsetid("", 18)],
ir_version=10,
)
onnx.checker.check_model(model)
onnx.save(model, path)
for n in [65535, 65536]:
path = f"avgpool_{n}.onnx"
make_model(path, n)
session = ort.InferenceSession(
path,
providers=["CUDAExecutionProvider"],
)
x = np.zeros((n, 1, 1, 512), dtype=np.float32)
try:
y = session.run(None, {"input": x})
print(n, "SUCCESS", y[0].shape)
except Exception as exc:
print(n, "FAILED", exc)
```
Observed result:
65535 SUCCESS (65535, 1, 1, 256)
65536 FAILED CUDNN_STATUS_NOT_SUPPORTED
### Error
CUDNN failure 3000: CUDNN_STATUS_NOT_SUPPORTED
file=/onnxruntime_src/onnxruntime/core/providers/cuda/nn/pool.cc
line=240
expr=PoolingForwardHelper(
GetCudnnHandle(context),
pooling_desc,
&alpha,
x_tensor,
x_data,
&beta,
y_tensor,
y_data
)
### Comparison with PyTorch
The equivalent PyTorch operation succeeds on the same GPU:
```python
import torch
import torch.nn.functional as F
x = torch.zeros((196608, 1, 1, 512), device="cuda")
y = F.avg_pool2d(x, kernel_size=(1, 2), stride=(1, 2))
print(y.shape)
```
Replacing the ONNX node with the following graph also succeeds using only
CUDAExecutionProvider:
Reshape [N, C, H, W] -> [N, C, H, W/2, 2]
ReduceMean(axis=-1)
The replacement produced bit-exact outputs in the complete model.
### Urgency
There is no fixed external deadline, but this is blocking validation and deployment of the high-resolution model.
### Platform
Linux
### OS Version
Ubuntu 24.04
### ONNX Runtime Installation
Released Package
### ONNX Runtime Version or Commit ID
1.23.2
### ONNX Runtime API
Python
### Architecture
X64
### Execution Provider
CUDA
### Execution Provider Library Version
_No response_
Contributor guide
Research direction
Start with onnxruntime/core/providers/cuda/nn/pool.cc at the reported failure around line 240, then run the supplied minimal reproduction with N=65535 and N=65536. Trace the CUDA AveragePool path and compare it with the successful Reshape + ReduceMean graph. Done means the valid N=65536 shape executes through CUDAExecutionProvider without CUDNN_STATUS_NOT_SUPPORTED and preserves the expected output shape.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100