Skip to content

mha: make impl_ptr_map caceh in fmha_v3_bwd and fmha_fwd_v3 thread-safe. #2201

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

Open
xinyazhang wants to merge 5 commits into
base: main
Choose a base branch
from
Open

mha: make impl_ptr_map caceh in fmha_v3_bwd and fmha_fwd_v3 thread-safe. #2201

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
461d730
mha: make impl_ptr_map caceh in fmha_v3_bwd and fmha_fwd_v3 thread-safe.
xinyazhang Mar 6, 2026
a1dea07
fmha_fwd_v3: share impl_ptr_map cache object to avoid duplicated loading
xinyazhang Mar 6, 2026
87ec181
mha: undef unused macros
xinyazhang Mar 6, 2026
c99e718
copilot doesn't like single macro usage so the coding practice is div…
xinyazhang Mar 6, 2026
a52d98b
restores csrc/cpp_itfs/mha_fwd.cu because it breaks XLA
xinyazhang Mar 6, 2026
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
7 changes: 7 additions & 0 deletions csrc/cpp_itfs/mha_bwd.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include "asm_fmha_v3_bwd_configs.hpp"
#include <memory>
#include <string>
#include <mutex>

namespace aiter {
std::tuple<int, int> get_padded_hdim(int hdim_q, int hdim_v, std::string arch_id)
Expand Down Expand Up @@ -369,7 +370,9 @@ float fmha_v3_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
AiterAsmKernel* impl_ptr_pre = nullptr;
AiterAsmKernel* impl_ptr_dqdkdv = nullptr;
AiterAsmKernel* impl_ptr_post = nullptr;
static std::mutex impl_ptr_mutex;
static std::unordered_map<std::string, std::unique_ptr<AiterAsmKernel>> impl_ptr_map;
#define LOCK_IMPL_PTR_MAP std::lock_guard<std::mutex> lock(impl_ptr_mutex)

auto it_pre = pre_cfgs->find(pre_kernel);
if(it_pre != pre_cfgs->end())
Expand All @@ -379,6 +382,7 @@ float fmha_v3_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
const char* co_name = cfg.co_name.c_str();
ts_odo = cfg.ts;

LOCK_IMPL_PTR_MAP;
auto result = impl_ptr_map.emplace(name, nullptr);
if(result.second)
{
Expand All @@ -400,6 +404,7 @@ float fmha_v3_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
const char* co_name = cfg.co_name.c_str();
ts_kv = cfg.ts;

LOCK_IMPL_PTR_MAP;
auto result = impl_ptr_map.emplace(name, nullptr);
if(result.second)
{
Expand All @@ -423,6 +428,7 @@ float fmha_v3_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
const char* co_name = cfg.co_name.c_str();
ts_dq = cfg.ts;

LOCK_IMPL_PTR_MAP;
auto result = impl_ptr_map.emplace(name, nullptr);
if(result.second)
{
Expand All @@ -436,6 +442,7 @@ float fmha_v3_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
return -1;
}
}
#undef LOCK_IMPL_PTR_MAP

if(a.v3_api_check)
return 1;
Expand Down
Loading