microsoft / microsoft/onnxruntime

CUDA AveragePool fails with CUDNN_STATUS_NOT_SUPPORTED when batch dimension reaches 65536

Open
#29,819 2 comments 0 reactions 0 assignees View on GitHub
ep:CUDA stale
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.