[None][fix] Add mutex to avoid potentially concurrent modifications to std::unordered_map

cpp/tensorrt_llm/common/opUtils.cpp

@@ -31,6 +31,8 @@
 #include <mutex>
 #include <thread>
 
+using tensorrt_llm::common::op::OpCustomHash;
+
 TRTLLM_NAMESPACE_BEGIN
 #if ENABLE_MULTI_DEVICE
 
@@ -194,7 +196,6 @@ void const* tensorrt_llm::common::op::getCommSessionHandle()
 
 namespace
 {
-using tensorrt_llm::common::op::hash;
 
 // Get current cuda context, a default context will be created if there is no context.
 inline CUcontext getCurrentCudaCtx()
@@ -229,15 +230,14 @@ class PerCudaCtxPerThreadSingletonCreator
     PerCudaCtxPerThreadSingletonCreator(CreatorFunc creator, DeleterFunc deleter)
         : mCreator{std::move(creator)}
         , mDeleter{std::move(deleter)}
-        , mObservers{new std::unordered_map<CacheKey, std::weak_ptr<T>, hash<CacheKey>>()}
+        , mObservers{std::make_unique<CacheTy>()}
     {
     }
 
     ~PerCudaCtxPerThreadSingletonCreator()
     {
         std::lock_guard<std::mutex> lk{mMutex};
-        delete mObservers;
-        mObservers = nullptr;
+        mObservers.reset();
     }
 
     std::shared_ptr<T> operator()()
@@ -299,7 +299,8 @@ class PerCudaCtxPerThreadSingletonCreator
     mutable std::mutex mMutex;
     // CUDA resources are per-context and per-thread.
     using CacheKey = std::tuple<CUcontext, std::thread::id>;
-    std::unordered_map<CacheKey, std::weak_ptr<T>, hash<CacheKey>>* mObservers;
+    using CacheTy = std::unordered_map<CacheKey, std::weak_ptr<T>, OpCustomHash<CacheKey>>;
+    std::unique_ptr<CacheTy> mObservers;
 };
 
 // Structure to hold memory information

cpp/tensorrt_llm/common/opUtils.h

@@ -108,79 +108,51 @@ namespace
 {
 
 template <typename T>
-struct hash_helper;
-
-// Base case: use std::hash for basic types
-template <typename T>
-struct hash_helper
+struct OpCustomHash : public std::hash<T>
 {
-    size_t operator()(T const& v) const
-    {
-        return std::hash<T>{}(v);
-    }
 };
 
-// Specialization for std::set
+template <class T>
+inline size_t hash_combine(size_t hash, T const& value)
+{
+    static constexpr size_t seed = 0x9e3779b9ULL;
+    using RemoveCVRefT = std::remove_cv_t<std::remove_reference_t<T>>;
+    return OpCustomHash<RemoveCVRefT>{}(value) + seed + (hash << 6) + (hash >> 2);
+}
+
 template <typename T>
-struct hash_helper<std::set<T>>
+struct OpCustomHash<std::set<T>>
 {
     size_t operator()(std::set<T> const& s) const
     {
         size_t hash_value = 0;
         for (auto const& item : s)
         {
             // Recursively hash each element
-            hash_value ^= hash_helper<T>{}(item) + 0x9e3779b9 + (hash_value << 6) + (hash_value >> 2);
+            hash_value ^= hash_combine(hash_value, item);
         }
         return hash_value;
     }
 };
 
-// Helper for tuple hashing
-template <typename Tuple, size_t Index = std::tuple_size<Tuple>::value - 1>
-struct tuple_hash_helper
-{
-    static size_t hash(Tuple const& tuple)
-    {
-        size_t hash_value = tuple_hash_helper<Tuple, Index - 1>::hash(tuple);
-        return hash_value
-            ^ (hash_helper<typename std::tuple_element<Index, Tuple>::type>{}(std::get<Index>(tuple)) + 0x9e3779b9
-                + (hash_value << 6) + (hash_value >> 2));
-    }
-};
-
-// Base case for tuple hashing
-template <typename Tuple>
-struct tuple_hash_helper<Tuple, 0>
+template <class... Args>
+class OpCustomHash<std::tuple<Args...>>
 {
-    static size_t hash(Tuple const& tuple)
+    template <std::size_t... Idx>
+    static size_t hash_impl(std::tuple<Args...> const& t, std::integer_sequence<std::size_t, Idx...>)
     {
-        return hash_helper<typename std::tuple_element<0, Tuple>::type>{}(std::get<0>(tuple));
+        size_t value = 0;
+        return ((value ^= hash_combine(value, std::get<Idx>(t))), ...);
     }
-};
 
-// Specialization for std::tuple
-template <typename... Args>
-struct hash_helper<std::tuple<Args...>>
-{
+public:
     size_t operator()(std::tuple<Args...> const& t) const
     {
-        return tuple_hash_helper<std::tuple<Args...>>::hash(t);
+        return hash_impl(t, std::make_index_sequence<sizeof...(Args)>{});
     }
 };
-
 } // namespace
 
-// Main hash struct to be used
-template <typename T>
-struct hash
-{
-    size_t operator()(T const& v) const
-    {
-        return hash_helper<T>{}(v);
-    }
-};
-
 // for testing only
 void const* getCommSessionHandle();
 } // namespace common::op

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -39,7 +39,7 @@ namespace torch_ext
 {
 using tensorrt_llm::common::op::AttentionOp;
 using tensorrt_llm::common::op::AttentionWorkspaceManager;
-using tensorrt_llm::common::op::hash;
+using tensorrt_llm::common::op::OpCustomHash;
 using tensorrt_llm::runtime::RequestType;
 
 namespace
@@ -1307,8 +1307,10 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
 
     auto cache_key = std::make_tuple(op->data(), runner->data());
     using CacheKey = decltype(cache_key);
-    static std::unordered_map<CacheKey, std::shared_ptr<AttentionOp>, hash<CacheKey>> op_cache;
-    if (auto it = op_cache.find(cache_key); it != op_cache.end())
+    static std::unordered_map<CacheKey, std::shared_ptr<AttentionOp>, OpCustomHash<CacheKey>> op_cache;
+    static std::shared_mutex op_cache_mutex;
+    if (auto it = (static_cast<void>(std::shared_lock<std::shared_mutex>{op_cache_mutex}), op_cache.find(cache_key));
+        it != op_cache.end())
     {
         TLLM_LOG_TRACE("Attention op for layer %d is cached", local_layer_idx);
         op = it->second;
@@ -1319,7 +1321,9 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
             to_string(cache_key).c_str());
         op->initialize();
         runner->prepare(*op);
-        op_cache[cache_key] = op;
+        std::unique_lock<std::shared_mutex> lock{op_cache_mutex};
+        auto [iter, _] = op_cache.try_emplace(cache_key, op);
+        op = iter->second;
     }
 
     int32_t const num_seqs = host_context_lengths.size(0);
@@ -1449,7 +1453,7 @@ bool attention_supports_nvfp4_output(int64_t const num_heads, int64_t const num_
 
     auto cache_key = op->data();
     using CacheKey = decltype(cache_key);
-    static std::unordered_map<CacheKey, bool, hash<CacheKey>> op_cache;
+    static std::unordered_map<CacheKey, bool, OpCustomHash<CacheKey>> op_cache;
     if (auto it = op_cache.find(cache_key); it != op_cache.end())
     {
         TLLM_LOG_TRACE("Attention op runtime check is cached");