microsoft / microsoft/onnxruntime

softmax_warp_forward overflows int32 above 2^31 elements and can return silently wrong results (CUDA and ROCm EPs)

Open
#32,299 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ep:CUDA
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.