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..bfc9a9e5975e 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -514,6 +514,26 @@ 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) +{ + 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); +} + +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 +848,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 +925,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)