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};