[Bug]: Async prompt-embedding H2D outlives its pooled-pinned source; the pinned pool can overwrite the buffer mid-copy

System Info

• TensorRT-LLM: Both 1.1.0 and 1.2.1. Reproduced with the stock libtensorrt_llm.so shipped in the container, and with a locally built .so from the same tag.
• GPU: NVIDIA H100 80GB HBM3, sm90.
• CUDA / driver: 13.1
• Container / backend: Triton Inference Server, TensorRT-LLM C++ backend (inflight fused batching); classic TensorRT backend 26.02 and 26.07-trtllm-python-py3
• Process topology (relevant): the decoder runs in trtllmExecutorWorker, i.e. a separate process and therefore a separate CUDA context from the other Triton models and Python backend stubs on the same GPU.
• Affected path: request-supplied prompt embedding table (prompt tuning), non-offloaded, where the host source is held in the reusable pinned-memory pool (cpp/tensorrt_llm/runtime/tllmBuffers.h).

This is a host-side object-lifetime defect, not a kernel or architecture issue. It should be hardware-independent. High concurrency and a competing CUDA context on the same device are only needed to widen the window enough to observe it.

Who can help?

Owners of:

cpp/tensorrt_llm/batch_manager/llmRequest.cpp — LlmRequest::movePromptEmbeddingTableToGpu()
cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp — PromptTuningBuffers::fill()
cpp/tensorrt_llm/runtime/bufferManager.cpp — BufferManager::copyFrom() / copy()
cpp/tensorrt_llm/runtime/tllmBuffers.h — pinned pool

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

There are two ways to hit this. The first is deterministic and does not require the race to fire.

A. Deterministic demonstration (recommended for review)

The regression test added in PR #18488 gates the CUDA stream with cudaLaunchHostFunc before the H2D, then asserts on source lifetime:

1. Insert a host callback that blocks the stream.
2. Create a pooled-pinned prompt embedding table; keep a weak_ptr to it.
3. Call PromptTuningBuffers::fill(), which queues the move and the event behind the gate.
4. Destroy or cancel the request.
5. On unpatched main, the weak_ptr has already expired — the source is back in the pool while the H2D is still pending on the stream.

This isolates the bug from scheduling luck entirely.

B. Statistical reproduction (how it was originally found)

1. Serve a model with prompt tuning enabled so that each request supplies its own prompt embedding table from pinned/pooled host memory (ours was 94,208 bytes per request).
2. Replay a fixed corpus of captured requests at concurrency 16 against the C++ backend with inflight batching. We used 3,108 requests per draw.
3. Add a competing CUDA process on the same GPU. It needs no access to TRT-LLM memory at all — eight Python backend instances running:

python

x = torch.ones(N, device='cuda')
y = x + 1

4. Classify outputs against a concurrency-1 golden capture.

Observed rates on our workload (each row 3,108 requests unless noted):

Configuration | Corrupted outputs -- | -- Decoder only, C=16 | 0 / 3,108 Decoder only, C=48 | 0 / 3,108 + CPU-only sidecar (same work, no CUDA) | 0 / 3,108 + GPU sidecar process with an initialised CUDA context but no kernels | 0 / 3,108 + dummy ones+1 CUDA sidecar, no synchronize | 10 / 3,108 + dummy ones+1 CUDA sidecar with synchronize | 15 / 3,108 + a real GPU neighbour model (×5 draws) | 252 / 15,540

Amplifier, if you need a higher yield for tracing: a 2 ms host sleep at the start of decoderSync() — after the launches, before the event wait — raised the rate to 28 / 3,108. A stall placed inside the stream (a fill-kernel nanosleep) does not amplify it, which is itself diagnostic: the vulnerable interval is between enqueue and the host's synchronisation point.

Note that CUDA_LAUNCH_BLOCKING=1 applied process-wide masks the bug (0 / 3,108, at ~24× the wall time), and so does MPS. Neither is a fix; both just close the window.

Expected behavior

BufferManager::copyFrom() issues cudaMemcpyAsync from a pinned host source. CUDA requires that source to remain allocated, unchanged, and at the same address until the copy completes on the stream.

The request-owned GPU prompt embedding table should therefore be a byte-exact copy of the host source, and greedy decoding over identical inputs should be deterministic regardless of what other processes are doing on the GPU.

actual behavior

The GPU copy of the prompt embedding table is torn: it contains a mixture of the original request's bytes and a later request's bytes.

Evidence from the run that localised it (one trapped request out of 3,108, greedy decoding, single repeated utterance so that every stage is deterministic and comparable to a C=1 golden):

Stage | Observation | Status -- | -- | -- CPU source, after the Triton handoff | golden hash 0xf890b624ba27d041 | correct Request-owned GPU table (first 16-byte canary) | 0x4465c937bd3dc1**00** vs golden 0x4465c937bd3dc1**ca** | torn Shared GPU prompt buffer after the D2D | identical torn value | faithful copy of a bad input Layer-0 KV after context | repeatable wrong hash 0x70837f2fce526439 vs golden 0xef0898cb6c2e1c4f | downstream Engine logits | argmax is already the stop token | downstream Sampled token | stop token, emitted after 1 generated token | user-visible failure

The Python/Triton input was correct and the host source was correct. The first divergence is the H2D itself, or the lifetime of its source during the H2D.

Two independent causal checks:

• Retain the source: keeping the original pooled-pinned tensor alive on LlmRequest until the request ends — with no added stream synchronisation and no change to kernel math, sampler logic, or scheduling — removed both the torn canary and the output failure (0 / 3,108).
• Remove the copy: supplying the prompt embedding table already on the GPU (via DLPack from an upstream model), so movePromptEmbeddingTableToGpu() is never called for that input, also gave 0 / 3,108 — on the completely unmodified stock library, with the sidecar still running at ~4,100 inferences/s.

Together these bracket the defect to the pooled-host → request-GPU transfer boundary.

Why this is easy to miss: the pinned block is returned to a pool, not unmapped, so the address stays valid. ASAN and Compute Sanitizer see a legal DMA from a mapped pinned region — no use-after-free, no illegal access. The result is corrupted-but-structurally-valid float data, so the symptom reads as model nondeterminism or sampling variance rather than memory corruption.

additional notes

Root cause — three individually reasonable contracts composing into a violation:

BufferManager::copyFrom() queues cudaMemcpyAsync on mStream and returns the destination. The caller must keep the source alive until the stream reaches the copy — but that obligation is implicit, and the method is named copyFrom, not copyFromAsync.
The pinned pool treats a released allocation as immediately reusable, and may hand the same address to the next request.
LlmRequest::movePromptEmbeddingTableToGpu() replaces mPromptEmbeddingTable with the GPU tensor as soon as copyFrom() returns, dropping the last host reference.

The third assumption violates the first whenever the source came from the second.

Failing sequence:

1 Request owns mPromptEmbeddingTable -> host pinned-pool tensor
2 BufferManager allocates a GPU destination
3 BufferManager queues cudaMemcpyAsync(dst, src) on its stream
4 copyFrom() returns to the CPU <-- the GPU may not have read a byte yet
5 LlmRequest replaces the member with the GPU TensorPtr
6 The last host TensorPtr reference disappears
7 The pinned pool marks the segment free; the address becomes reusable
8 Another request receives that address and writes new bytes
9 The original H2D reads a mixture of old and new bytes
10 The GPU destination stays allocated, containing corrupted data

Steps 5–8 all occur inside the interval in which CUDA requires the source to be stable. Nothing in the C++ type system, the allocator, or CUDA itself represents that interval.

Repository archaeology (in case this looks older than it is): BufferManager::copyFrom is present in the initial commit (2023-09-20) and pinnedPool was added 2024-01-09, but the exact copy-then-replace-the-source-member shape in movePromptEmbeddingTableToGpu() dates from 2025-03-11. All three pieces are required, which is one reason this survived in a heavily used stack.

Why most deployments never hit it — six conditions must hold at once:

request-supplied prompt embedding tables in use at all;
the source arriving from the reusable pinned pool;
the last source owner dropped immediately after queueing the H2D;
pool reuse landing before the DMA finishes reading;
enough GPU scheduling pressure to stretch that interval;
the torn bytes changing the output detectably.

Proposed fix — PR #18488. PromptTuningBuffers owns a list of PendingPromptEmbeddingTableH2DCopy holders. Each owns the pooled-pinned source, a CudaEvent, and a flag recording whether completion was successfully recorded. Hot path: link the holder → queue the H2D → record the event on the BufferManager stream → continue without blocking → poll with cudaEventQuery on later fill() calls → release the source only when the event reports complete. The 1–3 ordering is deliberate so that no potentially-throwing ownership insertion remains after the copy is queued. Teardown waits on the exact copy events rather than the whole runtime stream. Where completion cannot be proven, the holder is quarantined (released without destroying the source) — a bounded leak on an exceptional path is preferable to returning possibly-in-flight pinned memory to the pool.

An earlier prototype that simply called cudaEventSynchronize() inline in movePromptEmbeddingTableToGpu() did fix the race on a targeted reproducer, but hung the full pipeline (GPU utilisation 0%, client timeouts) when integrated. The CUDA event was never the problem — the placement and consumption policy were. Request cancellation and destruction must stay nonblocking, and a request destructor must never block an executor thread, which is why the guard lives in PromptTuningBuffers rather than on LlmRequest.

Related audit candidates (not claims of defect): LlmRequest::moveLoraWeightsToGpu() has the identical copy-then-replace shape; whether it is exploitable depends on source memory type, pooling, and later stream ordering, and was not exercised by these runs. Other copyFrom call sites worth the same four questions — is the transfer async for this source/destination pair, who owns the source after the call returns, can that storage be mutated or pooled before the stream reaches the copy, and is completion guaranteed indirectly — include temporary vectors, request word lists, Medusa tables, encoder inputs, and executor tensor copies.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start with LlmRequest::movePromptEmbeddingTableToGpu(), PromptTuningBuffers::fill(), BufferManager::copyFrom()/copy(), and the pinned pool in tllmBuffers.h. Review the regression test added in PR #18488 and run its deterministic stream-gated reproduction. Done means the pooled-pinned source remains valid through the asynchronous H2D and the test preserves byte-exact, deterministic output.