diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index 560750c2ba93..cfb40fe66ead 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -31,6 +31,8 @@ #include #include +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, hash>()} + , mObservers{std::make_unique()} { } ~PerCudaCtxPerThreadSingletonCreator() { std::lock_guard lk{mMutex}; - delete mObservers; - mObservers = nullptr; + mObservers.reset(); } std::shared_ptr operator()() @@ -299,7 +299,8 @@ class PerCudaCtxPerThreadSingletonCreator mutable std::mutex mMutex; // CUDA resources are per-context and per-thread. using CacheKey = std::tuple; - std::unordered_map, hash>* mObservers; + using CacheTy = std::unordered_map, OpCustomHash>; + std::unique_ptr mObservers; }; // Structure to hold memory information diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 22169843a9f5..fb0f2aca3d62 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -108,21 +108,20 @@ namespace { template -struct hash_helper; - -// Base case: use std::hash for basic types -template -struct hash_helper +struct OpCustomHash : public std::hash { - size_t operator()(T const& v) const - { - return std::hash{}(v); - } }; -// Specialization for std::set +template +inline size_t hash_combine(size_t hash, T const& value) +{ + static constexpr size_t seed = 0x9e3779b9ULL; + using RemoveCVRefT = std::remove_cv_t>; + return OpCustomHash{}(value) + seed + (hash << 6) + (hash >> 2); +} + template -struct hash_helper> +struct OpCustomHash> { size_t operator()(std::set const& s) const { @@ -130,57 +129,30 @@ struct hash_helper> for (auto const& item : s) { // Recursively hash each element - hash_value ^= hash_helper{}(item) + 0x9e3779b9 + (hash_value << 6) + (hash_value >> 2); + hash_value ^= hash_combine(hash_value, item); } return hash_value; } }; -// Helper for tuple hashing -template ::value - 1> -struct tuple_hash_helper -{ - static size_t hash(Tuple const& tuple) - { - size_t hash_value = tuple_hash_helper::hash(tuple); - return hash_value - ^ (hash_helper::type>{}(std::get(tuple)) + 0x9e3779b9 - + (hash_value << 6) + (hash_value >> 2)); - } -}; - -// Base case for tuple hashing -template -struct tuple_hash_helper +template +class OpCustomHash> { - static size_t hash(Tuple const& tuple) + template + static size_t hash_impl(std::tuple const& t, std::integer_sequence) { - return hash_helper::type>{}(std::get<0>(tuple)); + size_t value = 0; + return ((value ^= hash_combine(value, std::get(t))), ...); } -}; -// Specialization for std::tuple -template -struct hash_helper> -{ +public: size_t operator()(std::tuple const& t) const { - return tuple_hash_helper>::hash(t); + return hash_impl(t, std::make_index_sequence{}); } }; - } // namespace -// Main hash struct to be used -template -struct hash -{ - size_t operator()(T const& v) const - { - return hash_helper{}(v); - } -}; - // for testing only void const* getCommSessionHandle(); } // namespace common::op diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 811810a8d5a2..d17c1278e27b 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/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 k, std::optionaldata(), runner->data()); using CacheKey = decltype(cache_key); - static std::unordered_map, hash> op_cache; - if (auto it = op_cache.find(cache_key); it != op_cache.end()) + static std::unordered_map, OpCustomHash> op_cache; + static std::shared_mutex op_cache_mutex; + if (auto it = (static_cast(std::shared_lock{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 k, std::optionalinitialize(); runner->prepare(*op); - op_cache[cache_key] = op; + std::unique_lock 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> op_cache; + static std::unordered_map> op_cache; if (auto it = op_cache.find(cache_key); it != op_cache.end()) { TLLM_LOG_TRACE("Attention op runtime check is cached");