From d98d0b5987eaa6245e1f0a7a4af462cec51f1d6a Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 23 Jul 2026 01:26:18 -0700 Subject: [PATCH 1/3] [None][fix] Protect attention op cache from concurrent access Signed-off-by: Yihan Wang --- cpp/tensorrt_llm/common/opUtils.cpp | 5 +- cpp/tensorrt_llm/common/opUtils.h | 120 +++++++++----------------- cpp/tensorrt_llm/thop/attentionOp.cpp | 13 +-- 3 files changed, 53 insertions(+), 85 deletions(-) diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index 560750c2ba93..5b5ac76dea73 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -194,7 +194,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,7 +228,7 @@ class PerCudaCtxPerThreadSingletonCreator PerCudaCtxPerThreadSingletonCreator(CreatorFunc creator, DeleterFunc deleter) : mCreator{std::move(creator)} , mDeleter{std::move(deleter)} - , mObservers{new std::unordered_map, hash>()} + , mObservers{new std::unordered_map>()} { } @@ -299,7 +298,7 @@ class PerCudaCtxPerThreadSingletonCreator mutable std::mutex mMutex; // CUDA resources are per-context and per-thread. using CacheKey = std::tuple; - std::unordered_map, hash>* mObservers; + std::unordered_map>* mObservers; }; // Structure to hold memory information diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 22169843a9f5..8f4e0f8b0678 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -104,83 +104,6 @@ class UniqPtrWNullCopy : public std::unique_ptr } }; -namespace -{ - -template -struct hash_helper; - -// Base case: use std::hash for basic types -template -struct hash_helper -{ - size_t operator()(T const& v) const - { - return std::hash{}(v); - } -}; - -// Specialization for std::set -template -struct hash_helper> -{ - size_t operator()(std::set const& s) const - { - size_t hash_value = 0; - for (auto const& item : s) - { - // Recursively hash each element - hash_value ^= hash_helper{}(item) + 0x9e3779b9 + (hash_value << 6) + (hash_value >> 2); - } - 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 -{ - static size_t hash(Tuple const& tuple) - { - return hash_helper::type>{}(std::get<0>(tuple)); - } -}; - -// Specialization for std::tuple -template -struct hash_helper> -{ - size_t operator()(std::tuple const& t) const - { - return tuple_hash_helper>::hash(t); - } -}; - -} // 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 @@ -218,8 +141,51 @@ std::shared_ptr getComm(std::set const& group); std::shared_ptr getCublasHandle(); std::shared_ptr getCublasLtHandle(); +namespace +{ +template +inline size_t hash_combine(size_t hash, T const& value) +{ + static constexpr size_t seed = 0x9e3779b9ULL; + return std::hash>>{}(value) + seed + (hash << 6) + (hash >> 2); +} +} // namespace + TRTLLM_NAMESPACE_END +// Specialization for iterable containers. +template +struct std::hash> +{ + size_t operator()(std::set const& s) const + { + size_t hash_value = 0; + for (auto const& item : s) + { + // Recursively hash each element + hash_value ^= tensorrt_llm::hash_combine(hash_value, item); + } + return hash_value; + } +}; + +// Specialization for tuple-like containers. +template +struct std::hash> +{ + template + static size_t hash_impl(std::tuple const& t, std::integer_sequence) + { + size_t value = 0; + return ((value ^= tensorrt_llm::hash_combine(value, std::get(t))), ...); + } + + size_t operator()(std::tuple const& t) const + { + return hash_impl(t, std::make_index_sequence{}); + } +}; + #ifndef DEBUG #define PLUGIN_CHECK(status) \ diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 811810a8d5a2..09716178fe3e 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -39,7 +39,6 @@ namespace torch_ext { using tensorrt_llm::common::op::AttentionOp; using tensorrt_llm::common::op::AttentionWorkspaceManager; -using tensorrt_llm::common::op::hash; using tensorrt_llm::runtime::RequestType; namespace @@ -1307,8 +1306,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> 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 +1320,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 +1452,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"); From b45eb6d7bcfb80fd394e13845350ffd431d0b31c Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 23 Jul 2026 01:30:22 -0700 Subject: [PATCH 2/3] Update comments Signed-off-by: Yihan Wang --- cpp/tensorrt_llm/common/opUtils.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 8f4e0f8b0678..6b4e096c285e 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -153,7 +153,6 @@ inline size_t hash_combine(size_t hash, T const& value) TRTLLM_NAMESPACE_END -// Specialization for iterable containers. template struct std::hash> { @@ -169,7 +168,6 @@ struct std::hash> } }; -// Specialization for tuple-like containers. template struct std::hash> { From 40f2f31b965b78917de0a0c127860083e143b953 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Fri, 24 Jul 2026 00:05:13 -0700 Subject: [PATCH 3/3] Keep the custom hash class Signed-off-by: Yihan Wang --- cpp/tensorrt_llm/common/opUtils.cpp | 10 +-- cpp/tensorrt_llm/common/opUtils.h | 90 +++++++++++++++------------ cpp/tensorrt_llm/thop/attentionOp.cpp | 5 +- 3 files changed, 58 insertions(+), 47 deletions(-) diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index 5b5ac76dea73..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 @@ -228,15 +230,14 @@ class PerCudaCtxPerThreadSingletonCreator PerCudaCtxPerThreadSingletonCreator(CreatorFunc creator, DeleterFunc deleter) : mCreator{std::move(creator)} , mDeleter{std::move(deleter)} - , mObservers{new std::unordered_map>()} + , mObservers{std::make_unique()} { } ~PerCudaCtxPerThreadSingletonCreator() { std::lock_guard lk{mMutex}; - delete mObservers; - mObservers = nullptr; + mObservers.reset(); } std::shared_ptr operator()() @@ -298,7 +299,8 @@ class PerCudaCtxPerThreadSingletonCreator mutable std::mutex mMutex; // CUDA resources are per-context and per-thread. using CacheKey = std::tuple; - std::unordered_map>* 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 6b4e096c285e..fb0f2aca3d62 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -104,6 +104,55 @@ class UniqPtrWNullCopy : public std::unique_ptr } }; +namespace +{ + +template +struct OpCustomHash : public std::hash +{ +}; + +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 OpCustomHash> +{ + size_t operator()(std::set const& s) const + { + size_t hash_value = 0; + for (auto const& item : s) + { + // Recursively hash each element + hash_value ^= hash_combine(hash_value, item); + } + return hash_value; + } +}; + +template +class OpCustomHash> +{ + template + static size_t hash_impl(std::tuple const& t, std::integer_sequence) + { + size_t value = 0; + return ((value ^= hash_combine(value, std::get(t))), ...); + } + +public: + size_t operator()(std::tuple const& t) const + { + return hash_impl(t, std::make_index_sequence{}); + } +}; +} // namespace + // for testing only void const* getCommSessionHandle(); } // namespace common::op @@ -141,49 +190,8 @@ std::shared_ptr getComm(std::set const& group); std::shared_ptr getCublasHandle(); std::shared_ptr getCublasLtHandle(); -namespace -{ -template -inline size_t hash_combine(size_t hash, T const& value) -{ - static constexpr size_t seed = 0x9e3779b9ULL; - return std::hash>>{}(value) + seed + (hash << 6) + (hash >> 2); -} -} // namespace - TRTLLM_NAMESPACE_END -template -struct std::hash> -{ - size_t operator()(std::set const& s) const - { - size_t hash_value = 0; - for (auto const& item : s) - { - // Recursively hash each element - hash_value ^= tensorrt_llm::hash_combine(hash_value, item); - } - return hash_value; - } -}; - -template -struct std::hash> -{ - template - static size_t hash_impl(std::tuple const& t, std::integer_sequence) - { - size_t value = 0; - return ((value ^= tensorrt_llm::hash_combine(value, std::get(t))), ...); - } - - size_t operator()(std::tuple const& t) const - { - return hash_impl(t, std::make_index_sequence{}); - } -}; - #ifndef DEBUG #define PLUGIN_CHECK(status) \ diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 09716178fe3e..d17c1278e27b 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -39,6 +39,7 @@ namespace torch_ext { using tensorrt_llm::common::op::AttentionOp; using tensorrt_llm::common::op::AttentionWorkspaceManager; +using tensorrt_llm::common::op::OpCustomHash; using tensorrt_llm::runtime::RequestType; namespace @@ -1306,7 +1307,7 @@ void attention(torch::Tensor q, std::optional k, std::optionaldata(), runner->data()); using CacheKey = decltype(cache_key); - static std::unordered_map> op_cache; + 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()) @@ -1452,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");