[https://nvbugs/6607487][fix] fix LoRA host/device cache dtype reconfiguration race

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

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

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<TaskValuePtr> optOtherTaskValuePtr = [&]() -> std::optional<TaskValuePtr>
     {
-        std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> lk(deviceCache.mCacheMutex);
         otherIsDone = otherTaskValue->done;
         otherTaskValue->loadInProgress = false;
         otherTaskValue->loaded = true;

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)