Skip to content

fix(attention): Scope FMHA cubin loading by mask type #148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
suharvest wants to merge 1 commit into from
Closed

fix(attention): Scope FMHA cubin loading by mask type #148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 21 additions & 13 deletions cpp/kernels/contextAttentionKernels/contextFMHARunner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,18 +126,19 @@ 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;
}
};

struct FMHAKernelLoadHasher
{
size_t operator()(FMHAKernelLoadHashKey const& s) const noexcept
{
size_t key = s.data_type;
size_t key = static_cast<size_t>(s.data_type) ^ (static_cast<size_t>(s.attention_mask_type) << 8);
key <<= 16;
key ^= s.sm;
return key;
Expand Down Expand Up @@ -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]);
Expand All @@ -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;
}
Expand Down Expand Up @@ -269,6 +272,7 @@ class FMHAKernelList
int32_t mKernelMetaCount;
FMHADataType mDataType;
uint32_t mSMVersion;
int32_t mAttentionMaskType;
std::unordered_map<unsigned char const*, CUmodule> mModules;

std::unordered_map<FMHAKernelHashKey, FMHAKernelFuncInfo, FMHAKernelHasher> mFunctions;
Expand All @@ -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<std::mutex> 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<FMHAKernelList> newKernel = std::make_unique<FMHAKernelList>(type, sm);
std::unique_ptr<FMHAKernelList> newKernel
= std::make_unique<FMHAKernelList>(type, sm, attentionMaskType);
newKernel->loadFMHAKernels();
mKernels.insert(std::make_pair(hash_key, std::move(newKernel)));
findIter = mKernels.find(hash_key);
Expand All @@ -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
Expand Down Expand Up @@ -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;
}

Expand All @@ -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&)
Expand All @@ -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");

Expand Down
4 changes: 3 additions & 1 deletion cpp/kernels/contextAttentionKernels/contextFMHARunner.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 9 additions & 6 deletions cpp/plugins/attentionPlugin/attentionPlugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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();
}
Expand Down
5 changes: 3 additions & 2 deletions cpp/plugins/vitAttentionPlugin/vitAttentionPlugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -147,7 +148,7 @@ ViTAttentionPlugin::ViTAttentionPlugin(std::string const& name, PluginFieldColle
if (!mUseCuteDslFMHA)
#endif
{
ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType);
ContextFMHARunner::loadContextFMHAKernels(mSMVersion, mDataType, ContextAttentionMaskType::PADDING);
}
}

Expand Down
2 changes: 1 addition & 1 deletion unittests/contextAttentionTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ void TestContextAttentionAccuracy(std::vector<int32_t> 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,
Expand Down
3 changes: 2 additions & 1 deletion unittests/visionPackedMaskFMHATest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down