microsoft / microsoft/onnxruntime
CUDA EP: intermediate tensors are allocated without a stream, so the arena can hand them to another stream while work is still queued
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
Concurrent `Run()` calls on a single CUDA `InferenceSession` intermittently fail with a sticky
`cudaErrorIllegalAddress (700)`. The reported call is always innocent — every pointer and length in
it is valid — because the error is sticky and context-wide, so it surfaces on whatever API call
happens next rather than at the operation that caused it.
The cause is an allocation that never names its stream:
1. CUDA EP kernels build per-run intermediates with `Tensor::Create(type, shape, allocator)`, which
reaches the `Tensor` constructor and calls `allocator->Alloc(len)` — no stream
(`onnxruntime/core/framework/tensor.cc:94`).
2. `StreamAwareBFCArena` only associates a chunk with a stream when it is told which one, so these
chunks carry none.
3. `FindChunkPtr`'s `safe_to_use` accepts a free chunk when `!chunk->stream`
(`onnxruntime/core/framework/bfc_arena.cc:414`). That test conflates two different situations: a
chunk released *after its stream was synchronized* (safe for anyone) and a chunk that *never had
a stream* (not safe at all).
4. The intermediates are written by queued work — Einsum's `DataCopy` issues `cudaMemcpyAsync` and
`ZeroBuffer` issues `cudaMemsetAsync` on the run's stream — but the `unique_ptr` releases them as
soon as it leaves scope inside `Compute()`, while that work is still in flight. Another request's
thread then takes the same memory on a different stream, with nothing synchronizing the two.
The corrupted contents reach a kernel as an out-of-range address, and the process sees a 700 some
milliseconds later on an unrelated copy.
### Relationship to #29589
#29589 ("Fix usage of stream for memory pattern allocation with arena", commit `7e2f3b930`) is the
same class of bug in the same code, fixed for one caller: the memory-pattern buffer was allocated on
a fabricated `DummyStream` rather than the device's real stream, so the arena's association did not
match the work touching the memory. That one surfaced as unbounded memory growth (#29351) — chunks
never becoming reusable — rather than as corruption; ours is the opposite end of the same invariant,
memory becoming reusable too early.
The comment that #29589 removed is worth repeating, because it is the assumption that fails here:
> any memory pattern buffer would be in use for the entire inference, so there's no point at which
> another stream (as streams are per-inference) would be able to use it
That holds for a buffer that lives for the whole inference. It does not hold for an operator's
intermediates, which are released mid-inference, while another inference is running.
### Proposed fixes
Three independent changes; happy to open them as separate PRs.
**1. Give the CUDA EP a way to allocate a tensor on a stream, and use it in Einsum.**
Add a `Tensor` constructor `(elt_type, shape, allocator, Stream*)` that allocates through
`AllocOnStream` when the allocator is stream aware, falling back to the plain path otherwise, and a
matching `Tensor__construct` on the provider bridge. Then give the `CreateTensor` device helper the
`void* einsum_cuda_assets` parameter its siblings `DataCopy`, `ZeroBuffer` and `Transpose` already
carry, so the CUDA implementation can pass the run's stream — `einsum.cc` already builds
`EinsumCudaAssets` with `GetOrtStream(context)`, the stream simply never reached the allocation.
`Diagonal`, which allocates its output directly, gets the same treatment; the CPU implementation
ignores the argument.
**2. Sweep the remaining call sites.**
Einsum is not special. Sixteen further `Tensor::Create(type, shape, allocator)` calls build per-run
intermediates the same way and each is followed by an async copy or kernel on a stream that is
already in scope: the cross-device `Memcpy` for tensor sequences, Softmax's transpose scratch and
intermediate output, `ReduceCompute`'s output (which already takes the stream as a parameter),
Expand, Reshape, `IdentityOp` and the four Sequence element copies, beam search `TopK`, and the NCCL
all-gather / sharding / distributed-slice buffers. Conv and ConvTranspose `PrePack` are left alone —
they run at session initialization on the default stream and their buffers live as long as the
kernel.
**3. Make the arena report the situation instead of hiding it.**
The two meanings of a null `chunk->stream` are indistinguishable today, which is why this survived
so long. Track which one applies, count the unsafe case, warn once per arena, and expose it as
`AllocatorStats::num_unknown_stream_reuses` so it shows up in the arena dumps. Diagnostics only —
the flag takes no part in the reuse decision and a plain `BFCArena` never sets it. In our logs the
count is 230 per run before the fix and, importantly, zero from session initialization, so it goes
to zero for a model whose kernels all pass their stream.
Actually *refusing* the unsafe reuse is the natural follow-up, but it needs an answer for how such a
chunk becomes reusable again: it has no stream to match, so it would strand without a device-wide
sync point or a per-chunk event. That seems worth discussing separately.
### To reproduce
A 12-layer BERT encoder containing `Einsum` in the attention block; 148 nodes fall back to
CPU, which is what produces the small per-layer cross-device tensors involved; one CUDA
`InferenceSession`, 12–43 concurrent request threads. It is a race, so it reproduces
intermittently — in our runs, within a minute of load.
We instrumented the arena, the per-run stream cleanup and the CUDA memory API onto one timeline, and
attached a caller backtrace to every allocation made without a stream. Across 40 runs on 12 threads:
* **962 untagged device allocations per session**, from exactly **two call sites**, 481 of 128 bytes
and 481 of 2304 bytes — both resolving to
`EinsumOp::DeviceHelpers::CudaDeviceHelpers::CreateTensor`
(`onnxruntime/core/providers/cuda/math/einsum_utils/einsum_auxiliary_ops.cc:35`), reached through
the `Transpose` and `MatMul` wrappers in `einsum_auxiliary_ops.h`.
* **49 cross-thread handoffs** of a chunk that had never been associated with a stream, 39 of them
within 1 ms of the free, minimum gap **16 µs**.
* **Zero overlapping live allocations** across 82,658 alloc/free events — the arena's own
bookkeeping is correct. The hazard is entirely in *when* reuse is permitted.
A smaller run (12 threads, 9 runs) produced only 8 such handoffs, and all 8 were between the two
threads that then reported 700. One 150 µs window from that run:
```
23.138333 tid10859 D2D 128B -> 0x14e33d2be900 stream=...d400 (queued)
23.138343 tid10859 alloc 1664B stream=(nil) -> 0x14e33d2c8600
23.138386 tid10859 D2D 1664B <- 0x14e33d2c8600 stream=...d400 (queued)
23.138394 tid10859 free 0x14e33d2be900 tagged_stream=(nil)
23.138399 tid10859 free 0x14e33d2c8600 tagged_stream=(nil) <- 13us after queueing the read
23.138466 tid10860 alloc 16B -> 0x14e33d2be900 reuse=untagged <- 72us later, other stream
23.138477 tid10860 H2D 16B -> 0x14e33d2be900 stream=...e160
23.138482 tid10860 alloc 512B -> 0x14e33d2c8600 reuse=untagged
23.138493 tid10860 H2D 512B -> 0x14e33d2c8600 stream=...e160 <- overwrites the read source
23.237653 tid10859 D2H 39936B <- ... FAIL cudaErrorIllegalAddress(700)
29.900371 tid10860 H2D 144B -> ... FAIL cudaErrorIllegalAddress(700)
```
`tid10860` writes 512 bytes over the exact buffer `tid10859` had queued a 1664-byte read from, 107 µs
earlier, on a different stream.
Running a single request at a time makes the failure disappear.
### Urgency
Silent memory corruption between concurrent requests, leading to a sticky error soon after. It is not
detectable from the reported error, and the workaround is to serialize inference.
### Platform
Linux
### OS Version
AlmaLinux 9
### ONNX Runtime Installation
Built from Source
### ONNX Runtime Version or Commit ID
v1.29.0
### ONNX Runtime API
C++
### Architecture
X64
### Execution Provider
CUDA
### Execution Provider Library Version
CUDA 13.3
Contributor guide
Research direction
Start at the unstreamed allocation in onnxruntime/core/framework/tensor.cc:94 and the CUDA Einsum entry points in einsum_auxiliary_ops.cc and einsum_auxiliary_ops.h. Read bfc_arena.cc:414 and trace the listed remaining Tensor::Create call sites to understand their stream ownership. Done means per-run asynchronous intermediates carry the correct stream, remaining call sites are covered, and the requested arena diagnostics are implemented without changing plain BFCArena behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100