[#19172][fix] Fix prompt embedding table H2D source lifetime

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 <list>
+#include <memory>
+
 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<TensorPtr> const& optReqPromptEmbeddingTable,
         std::optional<SizeType32> 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<std::unique_ptr<PendingPromptEmbeddingTableH2DCopy>> mPendingPromptEmbeddingTableH2DCopies;
 };
 
 } // namespace tensorrt_llm::batch_manager

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
     {

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<void>(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<void>(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<PendingPromptEmbeddingTableH2DCopy>(std::move(source));
+    auto& result = *pendingCopy;
+    mPendingPromptEmbeddingTableH2DCopies.emplace_back(std::move(pendingCopy));
+    return result;
+}
+
 void PromptTuningBuffers::validate(
     std::optional<TensorPtr> const& optReqPromptEmbeddingTable, std::optional<SizeType32> 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<SizeType32>(contextRequests.size());
 
     std::vector<SizeType32> 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();
                     }
 

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 <cuda_runtime_api.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include <condition_variable>
+#include <mutex>
 #include <numeric>
 #include <string>
 #include <vector>
@@ -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<CudaStreamHostGate*>(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::CudaStream>();
+    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<tr::ITensor> promptEmbeddingTableWeak = promptEmbeddingTable;
+
+    std::optional<tr::ITensor::SharedPtr> promptEmbeddingTableOpt{std::move(promptEmbeddingTable)};
+    auto request = std::make_shared<tb::LlmRequest>(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};