From 3be47f82b63773f91c3e5f728bb937d6eb9c8f6f Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 23 Jul 2025 07:34:40 +0000 Subject: [PATCH] [fix] include rope style in XQA kernel cache key Signed-off-by: Yao Yao --- .../decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp | 3 ++- .../decoderMaskedMultiheadAttention/decoderXQAImplCommon.h | 5 ++++- .../decoderXQAImplJIT/decoderXQAImplJIT.cpp | 4 ++-- .../decoderXQAImplPrecompiled.cpp | 5 +++-- 4 files changed, 11 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp index c3bd96ab49d7..45b54165e803 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp @@ -55,7 +55,8 @@ XQAKernelRuntimeHashKey getRuntimeHashKeyFromXQAParams(XQAParams const& xqaParam // precompiled XQA does not use is_fp8_output as hashing key return {xqaParams.kv_cache_data_type, head_size, beam_width, kernel_num_q_heads_over_kv, kernel_m_tilesize, xqaParams.paged_kv_cache ? static_cast(xqaParams.tokens_per_block) : 0, xqaParams.paged_kv_cache, - xqaParams.multi_query_tokens, isXqaJit ? xqaParams.is_fp8_output : false}; + xqaParams.multi_query_tokens, isXqaJit ? xqaParams.is_fp8_output : false, + isXqaJit ? std::optional(xqaParams.position_embedding_type) : std::nullopt}; } } // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h index 6fe53d1f4783..884f3c0f64db 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h @@ -67,6 +67,7 @@ struct XQAKernelRuntimeHashKey bool paged_kv_cache; bool multi_query_tokens; bool is_fp8_output; + std::optional position_embedding_type; bool operator==(XQAKernelRuntimeHashKey const& other) const { @@ -74,7 +75,7 @@ struct XQAKernelRuntimeHashKey && num_q_heads_per_kv == other.num_q_heads_per_kv && beam_size == other.beam_size && multi_query_tokens == other.multi_query_tokens && m_tilesize == other.m_tilesize && tokens_per_page == other.tokens_per_page && paged_kv_cache == other.paged_kv_cache - && is_fp8_output == other.is_fp8_output; + && is_fp8_output == other.is_fp8_output && position_embedding_type == other.position_embedding_type; } }; @@ -103,6 +104,8 @@ struct XQAKernelRuntimeHasher key ^= s.multi_query_tokens; key <<= 1; key ^= s.is_fp8_output; + key <<= 8; + key ^= static_cast(s.position_embedding_type.value_or(static_cast(-1))); return key; } }; diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp index a0e116dda99c..802ffe31c5cc 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp @@ -37,8 +37,8 @@ using ::tensorrt_llm::kernels::XQAKernelMetaInfo; XQAKernelRuntimeHashKey getRuntimeHashKeyFromKernelMeta(XQAKernelMetaInfo const& kernelMeta) { return {kernelMeta.mKVDataType, kernelMeta.mHeadDim, kernelMeta.mBeamWidth, kernelMeta.mNumQHeadsOverKV, - kernelMeta.mMTileSize, kernelMeta.mTokensPerPage, kernelMeta.mPagedKVCache, kernelMeta.mMultiQueryTokens, - 0 /* xqa jit param is_fp8_output */}; + kernelMeta.mMTileSize, kernelMeta.mTokensPerPage, kernelMeta.mPagedKVCache, kernelMeta.mMultiQueryTokens, false, + std::nullopt}; } } // anonymous namespace diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp index ca0b36502794..6c6d4cd0b27e 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp @@ -97,7 +97,7 @@ class XQAKernelList } XQAKernelRuntimeHashKey hash_key{kernelMeta.mKVDataType, kernelMeta.mHeadDim, kernelMeta.mBeamWidth, kernelMeta.mNumQHeadsOverKV, kernelMeta.mMTileSize, kernelMeta.mTokensPerPage, kernelMeta.mPagedKVCache, - kernelMeta.mMultiQueryTokens, 0 /* xqa jit param is_fp8_output */}; + kernelMeta.mMultiQueryTokens, false, std::nullopt}; mFunctions.insert(std::make_pair(hash_key, funcInfo)); } @@ -128,7 +128,8 @@ class XQAKernelList XQAKernelRuntimeHashKey hash_key = {xqaParams.kv_cache_data_type, head_size, beam_width, kernel_num_q_heads_over_kv, m_tilesize, xqaParams.paged_kv_cache ? static_cast(xqaParams.tokens_per_block) : 0, - xqaParams.paged_kv_cache, xqaParams.multi_query_tokens, 0 /* xqa jit param is_fp8_output */}; + xqaParams.paged_kv_cache, xqaParams.multi_query_tokens, 0, /* xqa jit param is_fp8_output */ + std::nullopt}; auto const findIter = mFunctions.find(hash_key); return findIter != mFunctions.end(); }