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:
__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:
#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.
Summary
softmax_warp_forwardindexes its input with 32-bit arithmetic. Once the score tensorpasses 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.
Transposein the same codebasealready guards this class of overflow with a
narrowing_error;SoftmaxandErfdo not.The defect
onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh:first_batchandstrideare bothint, sofirst_batch * strideis evaluated in 32-bitand wraps at
INT32_MAX. The same directory'scu_inc/common.cuhstates the assumptionexplicitly:
Evidence that the bound is 2^31 elements
Bracketed independently at two sequence lengths differing 80x:
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 isbatch_count * stride = 2,158,755,840, exactly the first-faulting score tensor. Andwhich 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:
SoftmaxErfTransposenarrowing_errorMatMul,Add,Mul,LayerNormBoth reproduce inside a real model. At an
Erfgraph cut the output tail is the identical-3.896484375the precedingAddproduced, i.e. GELU never touched it. At aSoftmaxcut,an all-
-Infinput row yields0.0where 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]; withheads=16, chunk=128, overlap=16the boundary is:giving B_max = 2730 at S=256, 431 at S=2048, 44 at S=20480. A minimal repro needs only a
single
Softmaxnode over a[8194, 262144]fp16 tensor.Environment
main; the arithmetic is EP-independent, so the CUDA EP isaffected identically — we have not run the CUDA case ourselves.
Suggested fix
Widen the index arithmetic in
softmax_warp_forwardto 64-bit (int64_t, or castfirst_batchbefore the multiply), and — for the ops that cannot be widened cheaply — addthe same
narrowing_errorguardTransposealready carries, so an over-limit call failsloudly rather than returning a partially-written tensor.