diff --git a/cpp/kernels/contextAttentionKernels/contextFMHARunner.cpp b/cpp/kernels/contextAttentionKernels/contextFMHARunner.cpp index ee98381d..e236cb80 100644 --- a/cpp/kernels/contextAttentionKernels/contextFMHARunner.cpp +++ b/cpp/kernels/contextAttentionKernels/contextFMHARunner.cpp @@ -126,10 +126,11 @@ struct FMHAKernelLoadHashKey { FMHADataType data_type; int32_t sm; + int32_t attention_mask_type; bool operator==(FMHAKernelLoadHashKey const& other) const noexcept { - return data_type == other.data_type && sm == other.sm; + return data_type == other.data_type && sm == other.sm && attention_mask_type == other.attention_mask_type; } }; @@ -137,7 +138,7 @@ struct FMHAKernelLoadHasher { size_t operator()(FMHAKernelLoadHashKey const& s) const noexcept { - size_t key = s.data_type; + size_t key = static_cast(s.data_type) ^ (static_cast(s.attention_mask_type) << 8); key <<= 16; key ^= s.sm; return key; @@ -195,9 +196,10 @@ class FMHAKernelList using TKernelMetaInfo = fmha_v2::FusedMultiHeadAttentionKernelMetaInfoV2; public: - FMHAKernelList(FMHADataType type, int32_t sm) noexcept + FMHAKernelList(FMHADataType type, int32_t sm, int32_t attentionMaskType) noexcept : mDataType(type) , mSMVersion(sm) + , mAttentionMaskType(attentionMaskType) { mKernelMeta = &(fmha_v2::sMhaKernelMetaInfosV2[0]); mKernelMetaCount = sizeof(fmha_v2::sMhaKernelMetaInfosV2) / sizeof(fmha_v2::sMhaKernelMetaInfosV2[0]); @@ -214,7 +216,8 @@ class FMHAKernelList { auto const& kernelMeta = mKernelMeta[i]; if (kernelMeta.mDataTypeIn != mDataType || kernelMeta.mDataTypeOut != mDataType - || kernelMeta.mSM != mSMVersion || kernelMeta.mCubin == nullptr) + || kernelMeta.mSM != mSMVersion || kernelMeta.mCubin == nullptr + || kernelMeta.mAttentionMaskType != mAttentionMaskType) { continue; } @@ -269,6 +272,7 @@ class FMHAKernelList int32_t mKernelMetaCount; FMHADataType mDataType; uint32_t mSMVersion; + int32_t mAttentionMaskType; std::unordered_map mModules; std::unordered_map mFunctions; @@ -279,17 +283,18 @@ class FMHAKernelLoader public: //! @throws std::runtime_error if a CUDA driver error occurs - FMHAKernelList* getFMHAKernelList(FMHADataType type, int32_t sm) + FMHAKernelList* getFMHAKernelList(FMHADataType type, int32_t sm, int32_t attentionMaskType) { static std::mutex s_mutex; std::lock_guard lg(s_mutex); - FMHAKernelLoadHashKey hash_key{type, sm}; + FMHAKernelLoadHashKey hash_key{type, sm, attentionMaskType}; auto findIter = mKernels.find(hash_key); if (findIter == mKernels.end()) { - std::unique_ptr newKernel = std::make_unique(type, sm); + std::unique_ptr newKernel + = std::make_unique(type, sm, attentionMaskType); newKernel->loadFMHAKernels(); mKernels.insert(std::make_pair(hash_key, std::move(newKernel))); findIter = mKernels.find(hash_key); @@ -310,9 +315,9 @@ class FMHAKernelLoader }; //! @throws std::runtime_error if a CUDA driver error occurs -inline FMHAKernelList* getFMHAKernels(FMHADataType type, int32_t sm) +inline FMHAKernelList* getFMHAKernels(FMHADataType type, int32_t sm, ContextAttentionMaskType maskType) { - return FMHAKernelLoader::Get().getFMHAKernelList(type, sm); + return FMHAKernelLoader::Get().getFMHAKernelList(type, sm, attentionMaskTypeToInt(maskType)); } }; // namespace @@ -491,9 +496,10 @@ bool ContextFMHARunner::canImplement(int32_t headSize, [[maybe_unused]] int32_t return false; } -bool ContextFMHARunner::loadContextFMHAKernels(int32_t smVersion, nvinfer1::DataType dataType) +bool ContextFMHARunner::loadContextFMHAKernels( + int32_t smVersion, nvinfer1::DataType dataType, ContextAttentionMaskType maskType) { - FMHAKernelList* fmhaKernelList = getFMHAKernels(trtToFMHADataType(dataType), smVersion); + FMHAKernelList* fmhaKernelList = getFMHAKernels(trtToFMHADataType(dataType), smVersion, maskType); return fmhaKernelList != nullptr; } @@ -505,7 +511,8 @@ bool ContextFMHARunner::isKernelAvailable() const noexcept mLaunchParams.force_unroll, mLaunchParams.force_fp32_acc, mLaunchParams.flash_attention, attentionMaskTypeToInt(mLaunchParams.attention_mask_type), mLaunchParams.use_granular_tiling, attentionInputLayoutToInt(mLaunchParams.attention_input_layout)}; - FMHAKernelList const* fmhaKernelList = getFMHAKernels(trtToFMHADataType(mDataType), mSmVersion); + FMHAKernelList const* fmhaKernelList + = getFMHAKernels(trtToFMHADataType(mDataType), mSmVersion, mLaunchParams.attention_mask_type); return fmhaKernelList != nullptr && fmhaKernelList->findKernelFunction(hashKey).mSharedMemBytes != 0; } catch (std::exception const&) @@ -528,7 +535,8 @@ void ContextFMHARunner::dispatchFMHAKernel(FusedMultiheadAttentionParamsV2& para mLaunchParams.force_fp32_acc, mLaunchParams.flash_attention, attentionMaskTypeToInt(mLaunchParams.attention_mask_type), mLaunchParams.use_granular_tiling, attentionInputLayoutToInt(mLaunchParams.attention_input_layout)}; - FMHAKernelList* fmhaKernelList = getFMHAKernels(trtToFMHADataType(mDataType), mSmVersion); + FMHAKernelList* fmhaKernelList + = getFMHAKernels(trtToFMHADataType(mDataType), mSmVersion, mLaunchParams.attention_mask_type); FMHAKernelFuncInfo kernelInfo = fmhaKernelList->findKernelFunction(hashKey); check::check(kernelInfo.mSharedMemBytes != 0, "There must be one kernel to implement the MHA"); diff --git a/cpp/kernels/contextAttentionKernels/contextFMHARunner.h b/cpp/kernels/contextAttentionKernels/contextFMHARunner.h index 59c21a50..01f7d15e 100644 --- a/cpp/kernels/contextAttentionKernels/contextFMHARunner.h +++ b/cpp/kernels/contextAttentionKernels/contextFMHARunner.h @@ -131,10 +131,12 @@ class ContextFMHARunner * @brief Load FMHA kernel cubins into device * @param sm CUDA compute capability * @param dataType Data type + * @param maskType Attention mask type * @return True if successful * @throws std::runtime_error if a CUDA driver error occurs */ - static bool loadContextFMHAKernels(int32_t sm, nvinfer1::DataType dataType); + static bool loadContextFMHAKernels( + int32_t sm, nvinfer1::DataType dataType, ContextAttentionMaskType maskType); private: nvinfer1::DataType mDataType; //!< Data type diff --git a/cpp/plugins/attentionPlugin/attentionPlugin.cpp b/cpp/plugins/attentionPlugin/attentionPlugin.cpp index 8d5ae1ac..6060703a 100644 --- a/cpp/plugins/attentionPlugin/attentionPlugin.cpp +++ b/cpp/plugins/attentionPlugin/attentionPlugin.cpp @@ -172,12 +172,13 @@ bool loadFMHAKernels( if (!useCuteDslFMHA) #endif { - canImplementFMHA = ContextFMHARunner::canImplement(headSize, smVersion, dataType, - AttentionInputLayout::SEPARATE_Q_K_V, - useSlidingWindow ? ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL : ContextAttentionMaskType::CAUSAL); + auto const maskType = useSlidingWindow ? ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL + : ContextAttentionMaskType::CAUSAL; + canImplementFMHA = ContextFMHARunner::canImplement( + headSize, smVersion, dataType, AttentionInputLayout::SEPARATE_Q_K_V, maskType); if (canImplementFMHA) { - if (!ContextFMHARunner::loadContextFMHAKernels(smVersion, dataType)) + if (!ContextFMHARunner::loadContextFMHAKernels(smVersion, dataType, maskType)) { LOG_ERROR("Failed to load FMHA_v2 cubins for SM%d", smVersion); canImplementFMHA = false; @@ -407,7 +408,8 @@ AttentionPlugin::AttentionPlugin(std::string const& name, int32_t numQHeads, int // Availability is discovered from the cubin metadata table. mCanImplementCustomMaskFMHA = ContextFMHARunner::canImplement(mHeadSize, mSMVersion, mDataType, AttentionInputLayout::SEPARATE_Q_K_V, ContextAttentionMaskType::CUSTOM_MASK) - && ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType); + && ContextFMHARunner::loadContextFMHAKernels( + mSMVersion, mDataType, ContextAttentionMaskType::CUSTOM_MASK); enforceVisionBlockKernelSupport(); @@ -511,7 +513,8 @@ AttentionPlugin::AttentionPlugin(std::string const& name, PluginFieldCollection { mCanImplementCustomMaskFMHA = ContextFMHARunner::canImplement(mHeadSize, mSMVersion, mDataType, AttentionInputLayout::SEPARATE_Q_K_V, ContextAttentionMaskType::CUSTOM_MASK) - && ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType); + && ContextFMHARunner::loadContextFMHAKernels( + mSMVersion, mDataType, ContextAttentionMaskType::CUSTOM_MASK); enforceVisionBlockKernelSupport(); } diff --git a/cpp/plugins/vitAttentionPlugin/vitAttentionPlugin.cpp b/cpp/plugins/vitAttentionPlugin/vitAttentionPlugin.cpp index 824de5c7..fb58d222 100644 --- a/cpp/plugins/vitAttentionPlugin/vitAttentionPlugin.cpp +++ b/cpp/plugins/vitAttentionPlugin/vitAttentionPlugin.cpp @@ -106,7 +106,8 @@ ViTAttentionPlugin::ViTAttentionPlugin( mHeadSize, mSMVersion, mDataType, AttentionInputLayout::SEPARATE_Q_K_V, ContextAttentionMaskType::PADDING); if (canImplementFMHA) { - ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType); + ContextFMHARunner::loadContextFMHAKernels( + mSMVersion, mDataType, ContextAttentionMaskType::PADDING); } } @@ -147,7 +148,7 @@ ViTAttentionPlugin::ViTAttentionPlugin(std::string const& name, PluginFieldColle if (!mUseCuteDslFMHA) #endif { - ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType); + ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType, ContextAttentionMaskType::PADDING); } } diff --git a/unittests/contextAttentionTest.cpp b/unittests/contextAttentionTest.cpp index 3251e2ce..906a84bf 100755 --- a/unittests/contextAttentionTest.cpp +++ b/unittests/contextAttentionTest.cpp @@ -114,7 +114,7 @@ void TestContextAttentionAccuracy(std::vector const& cuSeqlens, int32_t cudaMemcpy(outReference.data(), oTensorRef.rawPointer(), outSize * sizeof(half), cudaMemcpyDeviceToHost)); // Load context FMHA kernels - EXPECT_TRUE(ContextFMHARunner::loadContextFMHAKernels(smVersion, DataType::kHALF)); + EXPECT_TRUE(ContextFMHARunner::loadContextFMHAKernels(smVersion, DataType::kHALF, maskType)); // Create context FMHA runner ContextFMHARunner runner(DataType::kHALF, batchSize, maxSeqLen, numQHeads, numKVHeads, headSize, smVersion, diff --git a/unittests/visionPackedMaskFMHATest.cpp b/unittests/visionPackedMaskFMHATest.cpp index 8a556f88..21347be8 100644 --- a/unittests/visionPackedMaskFMHATest.cpp +++ b/unittests/visionPackedMaskFMHATest.cpp @@ -362,7 +362,8 @@ TEST_P(VisionFMHACustomMaskParityTest, MatchesVisionBlockReferenceOracle) // Establish the CUDA primary context before the driver-API cubin loads // below (needed when this test is the first CUDA user in the process). CUDA_CHECK(cudaFree(nullptr)); - ASSERT_TRUE(ContextFMHARunner::loadContextFMHAKernels(smVersion, DataType::kHALF)); + ASSERT_TRUE(ContextFMHARunner::loadContextFMHAKernels( + smVersion, DataType::kHALF, ContextAttentionMaskType::CUSTOM_MASK)); ContextFMHARunner runner(DataType::kHALF, /*batchSize=*/1, seqLen, numQHeads, numKVHeads, headDim, smVersion, inputLayout, ContextAttentionMaskType::CUSTOM_MASK); ASSERT_TRUE(runner.isKernelAvailable())