microsoft / microsoft/onnxruntime
softmax_warp_forward overflows int32 above 2^31 elements and can return silently wrong results (CUDA and ROCm EPs)
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Summary
`softmax_warp_forward` indexes its input with 32-bit arithmetic. Once the score tensor
passes 2^31 elements the index multiply wraps and the kernel reads 2^32 bytes below its
base pointer. Depending on allocator layout this presents in one of three ways — a memory
fault, a **silently unwritten output tail**, or an apparent multi-minute hang — so a model
that exceeds the limit may return wrong logits rather than raising.
The silent-corruption case is the reason for this report. `Transpose` in the same codebase
already guards this class of overflow with a `narrowing_error`; `Softmax` and `Erf` do not.
### The defect
`onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh`:
```cpp
__global__ void softmax_warp_forward(output_t* dst, const input_t* src,
int batch_size, int stride, int element_count)
...
int first_batch = ...;
src += first_batch * stride + local_idx;
```
`first_batch` and `stride` are both `int`, so `first_batch * stride` is evaluated in 32-bit
and wraps at `INT32_MAX`. The same directory's `cu_inc/common.cuh` states the assumption
explicitly:
```cpp
#define CUDA_LONG int32_t // "We would like to use 64-bit integer to support large
// matrices. However, CUDA seems to support only 32-bit integer."
```
### Evidence that the bound is 2^31 *elements*
Bracketed independently at two sequence lengths differing 80x:
| S | last clean batch | score elements | / 2^31 | first faulting batch | score elements | / 2^31 |
|---|---|---|---|---|---|---|
| 20480 | 44 | 2,110,783,488 | 0.9829 | 45 | 2,158,755,840 | 1.0052 |
| 256 | 2730 | 2,146,959,360 | 0.9998 | 2731 | 2,147,745,792 | 1.0001 |
The S=256 bracket pins the boundary to **0.037%** of 2^31. The same boundary was reproduced
four times — two sequence lengths and two dtypes (fp16, fp32) — landing on the identical
*element* count each time. A byte-count or 4 GiB limit would have halved the fp32 boundary;
it did not move. The fp32 case is clean at a score tensor of 8.44 GB, which excludes any
byte-size explanation.
### The kernel, and that the offset is exactly one wrap
Under `AMD_SERIALIZE_KERNEL=3 AMD_LOG_LEVEL=4`, the last dispatch before the abort is
```
onnxruntime::rocm::softmax_warp_forward<__half, __half, float, 7>
batch_count = 16,865,280 stride = 128
```
`batch_count * stride = 2,158,755,840`, exactly the first-faulting score tensor. And
```
src_base - fault_address = 0x1_0000_0000 = 2^32 bytes = 2^31 fp16 elements
```
which is one exact int32 wrap, not an approximate overrun. The batch-44 control traces to a
clean `rc=0`.
### The part that matters most: it does not always crash
Single-op probes at 1.0376 x 2^31 elements:
| op | behaviour above the boundary |
|---|---|
| `Softmax` | **silently wrong** — output tail never written; exact threshold (correct at 8192 rows, garbage at 8194) |
| `Erf` | **silently wrong** — tail never written |
| `Transpose` | clean guarded `narrowing_error` |
| `MatMul`, `Add`, `Mul`, `LayerNorm` | correct |
Both reproduce inside a real model. At an `Erf` graph cut the output tail is the identical
`-3.896484375` the preceding `Add` produced, i.e. GELU never touched it. At a `Softmax` cut,
an all-`-Inf` input row yields `0.0` where a correct softmax must produce NaN.
Whether a given run faults or silently corrupts depends on allocator layout — the same
configuration produced both. **The crash is therefore not a guard.**
### Reproduction
Any model whose softmax input exceeds 2^31 elements. Ours is chunked attention producing a
score tensor of shape `[B*n_chunks, heads, chunk, chunk]`; with `heads=16, chunk=128,
overlap=16` the boundary is:
```
n_chunks(S) = ceil(max(S - chunk, 0) / (chunk - overlap)) + 1
B_max(S) = floor((2^31 - 1) / (n_chunks(S) * heads * chunk^2))
```
giving B_max = 2730 at S=256, 431 at S=2048, 44 at S=20480. A minimal repro needs only a
single `Softmax` node over a `[8194, 262144]` fp16 tensor.
### Environment
- onnxruntime-rocm 1.22.1, ROCm 7.0, AMD MI300A (gfx942)
- Source inspected on ORT `main`; the arithmetic is EP-independent, so the CUDA EP is
affected identically — we have not run the CUDA case ourselves.
### Suggested fix
Widen the index arithmetic in `softmax_warp_forward` to 64-bit (`int64_t`, or cast
`first_batch` before the multiply), and — for the ops that cannot be widened cheaply — add
the same `narrowing_error` guard `Transpose` already carries, so an over-limit call fails
loudly rather than returning a partially-written tensor.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh and inspect softmax_warp_forward, then compare its indexing and bounds handling with Transpose's narrowing_error guard. Use the documented single-op Softmax reproduction to check behavior above 2^31 elements; done means oversized inputs either produce correct output or fail explicitly rather than silently leaving an output tail unwritten.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100