From 5f834b36b3144a7df8a6edb3eca24977b0c101d8 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:39:25 -0700 Subject: [PATCH 1/2] [https://nvbugs/6607487][fix] fix LoRA host/device cache dtype reconfiguration race PeftCacheManager::configureDataType reconfigured the host and device LoraCache dtypes with two independent calls to LoraCache::setDataType. Each call is individually locked, but there is no lock spanning both, so a concurrent LoraCache::copyTask could observe the host cache already switched to the new dtype while the device cache still held the old one (or vice versa), tripping the 'LoRA host and device cache dtypes must match' check. copyTask's dtype check also only locked the device cache's mCacheMutex before reading both mPageManagerConfig.getDataType() values, despite a comment claiming this made the check atomic with respect to cache reconfiguration -- the host side of the comparison was read without holding host's mCacheMutex, racing against a concurrent configureDataType call. Add LoraCache::setDataTypeCoordinated, which reconfigures a pair of caches (host + device) under one combined lock so neither observer nor writer can see a torn state, and use it from configureDataType. Extend copyTask's dtype-check lock to cover both caches' mCacheMutex so the check is actually atomic against setDataTypeCoordinated. Also fix an unrelated mismatched-mutex bug in copyTask where device-cache TaskValue state was guarded by the host cache's mCacheMutex instead of the device cache's. This likely explains the intermittent (~0.5%) hang in test_qwen3_fp8_lora: an exception thrown from the dtype-mismatch check inside PeftCacheManager's ensure-worker-pool thread was not translated into a terminal response for the affected request, so the client blocked in queue.get() until the test's timeout. Re-enable test_qwen3_fp8_lora by removing its waives.txt entry. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/loraCache.h | 17 +++++++++++++ .../batch_manager/peftCacheManager.cpp | 7 ++++-- cpp/tensorrt_llm/runtime/loraCache.cpp | 25 ++++++++++++++++--- tests/integration/test_lists/waives.txt | 1 - 4 files changed, 43 insertions(+), 7 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/loraCache.h b/cpp/include/tensorrt_llm/runtime/loraCache.h index d6f8ca964539..ad4e86d6a0cd 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCache.h +++ b/cpp/include/tensorrt_llm/runtime/loraCache.h @@ -194,6 +194,17 @@ class LoraCache */ void setDataType(tensorrt_llm::DataType dataType); + /** + * \brief Reinitialize this cache and another cache (e.g. a host/device pair) to store the + * given data type as a single atomic operation. + * + * Holds both caches' mutexes for the full duration so a concurrent copyTask (which locks + * the same mutexes to verify the two caches agree on dtype before copying) can never + * observe one cache already reconfigured to the new dtype while the other still holds the + * old one. + */ + void setDataTypeCoordinated(LoraCache& other, tensorrt_llm::DataType dataType); + [[nodiscard]] tensorrt_llm::DataType getDataType() const; /** @@ -444,6 +455,12 @@ class LoraCache void bumpTaskInProgress(TaskIdType taskId); [[nodiscard]] ValueStatus getStatus(TaskIdType taskId) const; + /** + * \brief Core of setDataType, assumes mPagesMutex and mCacheMutex are already held by the + * caller. Used by setDataTypeCoordinated to reconfigure two caches under one combined lock. + */ + void setDataTypeLocked(tensorrt_llm::DataType dataType); + /** * \brief claim numPages, evicting tasks if needed * \param[in] numPages: number of pages to claim diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index 3e51ddc4c518..51fb73cda853 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -405,8 +405,11 @@ void PeftCacheManager::configureDataType(tensorrt_llm::DataType dataType) return; } - mHostLoraCache->setDataType(dataType); - mDeviceLoraCache->setDataType(dataType); + // Reconfigure host and device caches to the new dtype atomically: LoraCache::copyTask + // locks both caches' mCacheMutex to verify they agree on dtype before copying a task, so + // reconfiguring them one at a time here would let a concurrent copyTask observe one cache + // already switched to the new dtype while the other still holds the old one. + mHostLoraCache->setDataTypeCoordinated(*mDeviceLoraCache, dataType); mDataType = dataType; } diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 6749377445b2..66e3fa07593f 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -514,6 +514,18 @@ LoraCache::LoraCache(LoraCachePageManagerConfig const& pageManagerConfig, ModelC void LoraCache::setDataType(tensorrt_llm::DataType dataType) { std::scoped_lock lock(mPagesMutex, mCacheMutex); + setDataTypeLocked(dataType); +} + +void LoraCache::setDataTypeCoordinated(LoraCache& other, tensorrt_llm::DataType dataType) +{ + std::scoped_lock lock(mPagesMutex, mCacheMutex, other.mPagesMutex, other.mCacheMutex); + setDataTypeLocked(dataType); + other.setDataTypeLocked(dataType); +} + +void LoraCache::setDataTypeLocked(tensorrt_llm::DataType dataType) +{ TLLM_CHECK_WITH_INFO(mCacheMap.empty(), "Cannot change LoRA cache dtype after a task has been inserted"); if (mPageManagerConfig.getDataType() == dataType) { @@ -828,9 +840,12 @@ void LoraCache::copyTask(TaskIdType taskId, LoraCache& deviceCache, bool markDon // TaskValue* otherTaskValuePtr = copyTaskGetOtherTaskValue(taskId, taskValue, deviceCache, markDone); std::optional optOtherTaskValuePtr = [&]() -> std::optional { - std::lock_guard deviceCacheLock(deviceCache.mCacheMutex); - // setDataType also holds mCacheMutex, so the check and target task - // insertion are atomic with respect to cache reconfiguration. + // Lock both caches' mCacheMutex (this=host, deviceCache=device): setDataTypeCoordinated + // holds both caches' mCacheMutex for the full duration of a host+device dtype swap, so + // locking both here too makes this dtype check and the following task-map insertion + // atomic with respect to cache reconfiguration. Locking deviceCache.mCacheMutex alone is + // not sufficient since mPageManagerConfig on the host side (this) is read here too. + std::scoped_lock cacheLock(mCacheMutex, deviceCache.mCacheMutex); TLLM_CHECK_WITH_INFO(mPageManagerConfig.getDataType() == deviceCache.mPageManagerConfig.getDataType(), "LoRA host and device cache dtypes must match"); auto otherStatus = deviceCache.getStatus(taskId); @@ -902,7 +917,9 @@ void LoraCache::copyTask(TaskIdType taskId, LoraCache& deviceCache, bool markDon bool otherIsDone; { - std::lock_guard lk(mCacheMutex); + // otherTaskValue belongs to deviceCache's mCacheMap, so it must be guarded by + // deviceCache's mCacheMutex, not this (host) cache's mCacheMutex. + std::lock_guard lk(deviceCache.mCacheMutex); otherIsDone = otherTaskValue->done; otherTaskValue->loadInProgress = false; otherTaskValue->loaded = true; diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c709cc7d3984..bd2b2f477423 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -365,7 +365,6 @@ unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169) unittest/_torch/modules/test_w4a16_nvfp4_linear.py::test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin SKIP (https://nvbugs/6581071) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) -unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py::TestQwen3LoRA::test_qwen3_fp8_lora SKIP (https://nvbugs/6607487) unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-balanced] SKIP (https://nvbugs/6507113) unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part4" SKIP (https://nvbugs/6437410) From 048e4315bc9a36cc7e781866504345494a5be60e Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:42:27 -0700 Subject: [PATCH 2/2] [https://nvbugs/6607487][fix] guard setDataTypeCoordinated against self-aliasing Passing the same LoraCache as both the receiver and other would lock mPagesMutex/mCacheMutex twice via the same std::scoped_lock call, which is undefined behavior (self-deadlock) for a non-recursive std::mutex. Route that case through the existing single-cache setDataType instead. Addresses CodeRabbit review feedback on PR #17678. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- cpp/tensorrt_llm/runtime/loraCache.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 66e3fa07593f..bfc9a9e5975e 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -519,6 +519,14 @@ void LoraCache::setDataType(tensorrt_llm::DataType dataType) void LoraCache::setDataTypeCoordinated(LoraCache& other, tensorrt_llm::DataType dataType) { + if (&other == this) + { + // Locking mPagesMutex/mCacheMutex twice via the same std::scoped_lock call below would + // be undefined behavior (self-deadlock) for a non-recursive std::mutex. + setDataType(dataType); + return; + } + std::scoped_lock lock(mPagesMutex, mCacheMutex, other.mPagesMutex, other.mCacheMutex); setDataTypeLocked(dataType); other.setDataTypeLocked(dataType);