From 64e74511f6abff56cb5c5a27a4225567e0e52a37 Mon Sep 17 00:00:00 2001 From: rpengms Date: Mon, 31 Aug 2026 20:51:53 -0700 Subject: [PATCH] [None][fix] Fix prompt embedding table H2D source lifetime BufferManager::copyFrom queues host-to-device copies asynchronously on its CUDA stream. LlmRequest replaced its pooled-pinned prompt embedding table with the destination GPU tensor immediately, allowing the host tensor to be released before the copy completed. The pinned pool could then reuse and overwrite that source allocation while CUDA was still reading it. GPU contention or preemption from a sidecar made the lifetime window larger, causing corrupted prompt embeddings, nondeterministic decoder output, and random early EOS, including with greedy decoding. Retain each host prompt embedding table in PromptTuningBuffers before queueing the H2D copy. Record a CUDA event after the copy and reclaim the source during later nonblocking polls only after that event completes. During destruction, wait only for the exact copy events. If completion cannot be established on an exceptional path, quarantine the source rather than returning it to the reusable pinned pool. Add a deterministic regression test that blocks the CUDA stream, verifies that the source remains alive while the copy is pending, and verifies that it is released after completion. Tests: Not run locally; no configured C++ CUDA test build is available in this checkout. --- .../batch_manager/promptTuningBuffers.h | 13 +++ cpp/include/tensorrt_llm/runtime/cudaEvent.h | 16 +++ .../batch_manager/promptTuningBuffers.cpp | 106 +++++++++++++++++- .../batch_manager/llmRequestTest.cpp | 86 ++++++++++++++ 4 files changed, 219 insertions(+), 2 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h b/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h index a1d8849a8811..d9e3d187214d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h +++ b/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h @@ -24,6 +24,9 @@ #include "tensorrt_llm/runtime/promptTuningParams.h" #include "tensorrt_llm/runtime/worldConfig.h" +#include +#include + namespace tensorrt_llm::batch_manager { @@ -44,6 +47,8 @@ class PromptTuningBuffers PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, bool promptTableOffloading); + ~PromptTuningBuffers(); + void validate(std::optional const& optReqPromptEmbeddingTable, std::optional const& optReqPromptVocabSize); @@ -101,6 +106,14 @@ class PromptTuningBuffers void updateBufferStartPosition(size_t index, SizeType32 numRows); void clearBufferStartPositions(size_t index); + +private: + class PendingPromptEmbeddingTableH2DCopy; + + void releasePromptEmbeddingTableHostBuffersIfReady(); + PendingPromptEmbeddingTableH2DCopy& retainPromptEmbeddingTableHostBufferBeforeCopy(TensorPtr source); + + std::list> mPendingPromptEmbeddingTableH2DCopies; }; } // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/cudaEvent.h b/cpp/include/tensorrt_llm/runtime/cudaEvent.h index ad7b7fdb45cc..c4d429efc426 100644 --- a/cpp/include/tensorrt_llm/runtime/cudaEvent.h +++ b/cpp/include/tensorrt_llm/runtime/cudaEvent.h @@ -66,6 +66,22 @@ class CudaEvent TLLM_CUDA_CHECK(::cudaEventSynchronize(get())); } + //! \brief Returns true when all work captured by the event has completed. + [[nodiscard]] bool isComplete() const + { + auto const error = ::cudaEventQuery(get()); + if (error == cudaSuccess) + { + return true; + } + if (error == cudaErrorNotReady) + { + return false; + } + TLLM_CUDA_CHECK(error); + return false; + } + private: class Deleter { diff --git a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp b/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp index 1cf73a2c0d21..1b5962177691 100644 --- a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp @@ -18,13 +18,50 @@ #include "tensorrt_llm/batch_manager/promptTuningBuffers.h" #include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/runtime/cudaEvent.h" namespace tensorrt_llm::batch_manager { using SizeType32 = tensorrt_llm::runtime::SizeType32; using TensorPtr = runtime::ITensor::SharedPtr; +class PromptTuningBuffers::PendingPromptEmbeddingTableH2DCopy +{ +public: + explicit PendingPromptEmbeddingTableH2DCopy(TensorPtr source) + : mSource{std::move(source)} + { + } + + void recordCompletion(runtime::CudaStream const& stream) + { + stream.record(mCompletionEvent); + mCompletionEventRecorded = true; + } + + [[nodiscard]] bool isComplete() const + { + return mCompletionEventRecorded && mCompletionEvent.isComplete(); + } + + [[nodiscard]] bool hasRecordedCompletionEvent() const + { + return mCompletionEventRecorded; + } + + void synchronize() const + { + mCompletionEvent.synchronize(); + } + +private: + TensorPtr mSource; + runtime::CudaEvent mCompletionEvent; + bool mCompletionEventRecorded{false}; +}; + PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) { @@ -74,6 +111,58 @@ PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::Buffe mPromptTuningParams.tasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); } +PromptTuningBuffers::~PromptTuningBuffers() +{ + for (auto& pendingCopy : mPendingPromptEmbeddingTableH2DCopies) + { + if (!pendingCopy->hasRecordedCompletionEvent()) + { + // An exception occurred after ownership was retained but before completion could be recorded. Completion + // cannot be proven, so intentionally quarantine the pooled-pinned source instead of making it reusable. + TLLM_LOG_ERROR("Quarantining prompt embedding table source because its H2D completion was not recorded"); + static_cast(pendingCopy.release()); + continue; + } + + try + { + // Wait only for this source's H2D copy, not for unrelated work queued later on the runtime stream. + pendingCopy->synchronize(); + } + catch (std::exception const& e) + { + TLLM_LOG_EXCEPTION(e); + // Releasing this source without proving completion would recreate the pinned-pool reuse race. Leaking the + // small holder on this exceptional shutdown path is the safe failure mode. + static_cast(pendingCopy.release()); + } + } +} + +void PromptTuningBuffers::releasePromptEmbeddingTableHostBuffersIfReady() +{ + for (auto it = mPendingPromptEmbeddingTableH2DCopies.begin(); it != mPendingPromptEmbeddingTableH2DCopies.end();) + { + if ((*it)->isComplete()) + { + it = mPendingPromptEmbeddingTableH2DCopies.erase(it); + } + else + { + ++it; + } + } +} + +PromptTuningBuffers::PendingPromptEmbeddingTableH2DCopy& +PromptTuningBuffers::retainPromptEmbeddingTableHostBufferBeforeCopy(TensorPtr source) +{ + auto pendingCopy = std::make_unique(std::move(source)); + auto& result = *pendingCopy; + mPendingPromptEmbeddingTableH2DCopies.emplace_back(std::move(pendingCopy)); + return result; +} + void PromptTuningBuffers::validate( std::optional const& optReqPromptEmbeddingTable, std::optional const& optReqPromptVocabSize) { @@ -118,6 +207,8 @@ void PromptTuningBuffers::fill(RequestVector const& contextRequests, RequestVect { NVTX3_SCOPED_RANGE_WITH_NAME(range, "PromptTuningBuffers::fill"); + releasePromptEmbeddingTableHostBuffersIfReady(); + auto const numContextRequests = static_cast(contextRequests.size()); std::vector reqBeamWidths; @@ -172,8 +263,19 @@ void PromptTuningBuffers::fill(RequestVector const& contextRequests, RequestVect } else { - // Move to GPU - llmReq->movePromptEmbeddingTableToGpu(manager); + // Keep the pooled-pinned source alive outside the request until its asynchronous H2D copy has + // completed. Request cancellation must neither release the source nor wait for the copy. + auto hostPromptEmbeddingTable = optReqPromptEmbeddingTable.value(); + if (hostPromptEmbeddingTable->getMemoryType() != runtime::MemoryType::kGPU) + { + // Allocate and link the holder before queuing the asynchronous copy. Once copyFrom returns, + // no potentially-throwing ownership operation may occur before the completion event is + // recorded. + auto& pendingCopy + = retainPromptEmbeddingTableHostBufferBeforeCopy(std::move(hostPromptEmbeddingTable)); + llmReq->movePromptEmbeddingTableToGpu(manager); + pendingCopy.recordCompletion(manager.getStream()); + } optReqPromptEmbeddingTable = llmReq->getPromptEmbeddingTable(); } diff --git a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp index f1d39dc384dd..d189cfa33240 100644 --- a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp @@ -16,12 +16,17 @@ */ #include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/batch_manager/promptTuningBuffers.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/executor/types.h" +#include "tensorrt_llm/runtime/bufferManager.h" +#include #include #include +#include +#include #include #include #include @@ -36,6 +41,41 @@ using SizeType32 = tb::LlmRequest::SizeType32; using VecTokenExtraIds = tb::LlmRequest::VecTokenExtraIds; using VecUniqueTokens = tb::LlmRequest::VecUniqueTokens; +namespace +{ + +class CudaStreamHostGate +{ +public: + ~CudaStreamHostGate() + { + open(); + } + + static void CUDART_CB wait(void* opaque) + { + auto& gate = *static_cast(opaque); + std::unique_lock lock{gate.mMutex}; + gate.mCondition.wait(lock, [&gate] { return gate.mOpen; }); + } + + void open() + { + { + std::lock_guard lock{mMutex}; + mOpen = true; + } + mCondition.notify_all(); + } + +private: + std::mutex mMutex; + std::condition_variable mCondition; + bool mOpen{false}; +}; + +} // namespace + class LlmRequestTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) { protected: @@ -162,6 +202,52 @@ TEST_F(LlmRequestTest, fromExecutorRequest) } } +TEST_F(LlmRequestTest, promptEmbeddingTableH2DSourceLifetime) +{ + auto stream = std::make_shared(); + tr::BufferManager manager{stream}; + + SizeType32 constexpr promptVocabSize{23}; + SizeType32 constexpr hiddenSize{64}; + tr::ModelConfig modelConfig(256, 1, 1, 0, 1, hiddenSize, nvinfer1::DataType::kFLOAT); + modelConfig.setMaxPromptEmbeddingTableSize(promptVocabSize); + tr::WorldConfig worldConfig; + tb::PromptTuningBuffers promptTuningBuffers(1, manager, modelConfig, worldConfig, false); + + // Hold this stream before the H2D so the first cleanup poll deterministically sees an incomplete event. The gate's + // destructor opens it before promptTuningBuffers is destroyed, including on an early ASSERT return. + CudaStreamHostGate streamGate; + TLLM_CUDA_CHECK(cudaLaunchHostFunc(stream->get(), &CudaStreamHostGate::wait, &streamGate)); + + auto promptEmbeddingTable = tr::BufferManager::pinnedPool( + tr::ITensor::makeShape({1, promptVocabSize, hiddenSize}), nvinfer1::DataType::kFLOAT); + std::weak_ptr promptEmbeddingTableWeak = promptEmbeddingTable; + + std::optional promptEmbeddingTableOpt{std::move(promptEmbeddingTable)}; + auto request = std::make_shared(1, 1, VecTokens{1}, tr::SamplingConfig{}, false, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::move(promptEmbeddingTableOpt), + promptVocabSize); + tb::RequestVector contextRequests{request}; + + promptTuningBuffers.fill(contextRequests, {}, manager, false); + ASSERT_TRUE(request->getPromptEmbeddingTable().has_value()); + EXPECT_EQ(request->getPromptEmbeddingTable().value()->getMemoryType(), tr::MemoryType::kGPU); + + contextRequests.clear(); + request.reset(); + + // Request cancellation/destruction must not block and must not control the asynchronous copy's source lifetime. + EXPECT_FALSE(promptEmbeddingTableWeak.expired()); + + promptTuningBuffers.fill({}, {}, manager, false); + EXPECT_FALSE(promptEmbeddingTableWeak.expired()); + + streamGate.open(); + stream->synchronize(); + promptTuningBuffers.fill({}, {}, manager, false); + EXPECT_TRUE(promptEmbeddingTableWeak.expired()); +} + TEST_F(LlmRequestTest, invalidExecRequest) { VecTokens inputTokens{1, 2, 3, 4, 5};