Skip to content

Added AITER V3 API check mechanism #662

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

Merged
wenchenvincent merged 4 commits into from
Jul 15, 2026
+88 −23
Merged

Added AITER V3 API check mechanism #662

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6e10d18
Added V3 API check mechanism
Micky774 Jul 7, 2026
8ca1c99
Updated to add QoLA-side AITER patch for V3 API checks
Micky774 Jul 7, 2026
0698c9a
Updated qola commit to main
Micky774 Jul 7, 2026
cd1ac69
Removed static, updated arg structure for build_bwd_fmha_args
Micky774 Jul 10, 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
2 changes: 1 addition & 1 deletion 3rdparty/QoLA
Open in desktop
Submodule QoLA updated 1 files
+56 −0 patches/aiter/0004-mha-bwd-v3-api-check-dry-run.patch
6 changes: 6 additions & 0 deletions transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,12 @@ struct CkAttnBwdArgs : CKAttnCommonArgs {
hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream);
hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream);

// Probe whether AITER's v3 (asm) path will run for the given config, without
// launching a kernel (backed by AITER's v3_api_check dry-run). Returns true iff
// the v3 path is selected; false means the CK v2 path (or no support) would run.
bool ck_attn_fwd_uses_v3(const CKAttnFwdArgs& args);
bool ck_attn_bwd_uses_v3(const CkAttnBwdArgs& args);
Comment on lines +146 to +150

@github-actions github-actions Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor doc-clarity nit for future callers: the probe's result depends on the caller's uses_fwd_v3 / uses_bwd_v3 flags, because build_{fwd,bwd}_fmha_args copies them into fmha_args.use_asm_v3 before flipping v3_api_check. So this reports "would v3 dispatch given this config as-configured", not "is v3 available for this config regardless of the opt-in flag".

If future callers use the probe to decide whether to enable v3 (rather than confirm an already-decided choice), they'll need to pre-set args.uses_{fwd,bwd}_v3 = true before calling — worth mentioning in this comment block so that's not a surprise. Not blocking.


}//namespace ck_fused_attn
#endif // CK_FUSED_ATTN_H

56 changes: 46 additions & 10 deletions transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -441,21 +441,24 @@ void dump_bwd_timings(const char* dump_path, float average_runtime){
file << average_runtime << "\n";
}

hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){
struct BwdFmhaArgs {
aiter::mha_bwd_args fmha_args;
bool is_mqa_gqa;
bool has_dbias;
BiasShape bias_shape;
};

// Populate the AITER mha_bwd_args from TE's CkAttnBwdArgs. Shared by ck_attn_bwd
// (real launch) and ck_attn_bwd_uses_v3 (v3 availability probe) so the probe can
// never disagree with the launch. v3_api_check is left false here; callers flip it.
// The stream-dependent max_seqlen override (NVTE_CK_RUNTIME_MAX_SEQLEN) is applied
// by ck_attn_bwd after this returns; it does not affect v3 kernel selection.
BwdFmhaArgs build_bwd_fmha_args(const CkAttnBwdArgs& args){

bool has_dropout = (args.dropout_probability > 0.f);
bool has_dbias = args.dbias_ptr != nullptr;
bool is_mqa_gqa = (args.h > args.hg);

bool ck_log_config = false;
if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) {
if (env_p != nullptr && std::string(env_p) == "1")
ck_log_config = true;
}
const char* dump_path = std::getenv("NVTE_DUMP_AITER_RT");
// print kernel name on verbose mode
ck_tile::stream_config stream_config{stream, dump_path!=nullptr, get_ck_log_stream() != nullptr};

bias_enum bias_type = bias_enum::no_bias;
BiasShape bias_shape = BiasShape::k11SS;
if (!args.is_group_mode()) {
Expand Down Expand Up @@ -578,6 +581,39 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){
fmha_args.p_undrop = 1.0 - args.dropout_probability;
fmha_args.drop_seed_offset = std::pair<const void*, const void*>{args.philox_seed_ptr, args.philox_offset_ptr};

return {fmha_args, is_mqa_gqa, has_dbias, bias_shape};
}

// Probe whether AITER's v3 (asm) backward path will run for this config, without
// launching any kernel. Builds the same args as ck_attn_bwd and relies on AITER's
// v3_api_check dry-run (returns 1 when v3 is available, -1 otherwise).
bool ck_attn_bwd_uses_v3(const CkAttnBwdArgs& args){
aiter::mha_bwd_args fmha_args = build_bwd_fmha_args(args).fmha_args;
fmha_args.v3_api_check = true;
// No kernel is launched in check mode, so the stream/log flags are irrelevant.
ck_tile::stream_config stream_config{nullptr, false, false};
return QOLA_NS(mha_bwd)(fmha_args, stream_config) == 1;
}

hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){

// build_bwd_fmha_args also derives the flags the post-dispatch MQA/GQA and dbias
// reductions below need, so we unpack them here rather than recomputing.
BwdFmhaArgs built = build_bwd_fmha_args(args);
aiter::mha_bwd_args& fmha_args = built.fmha_args;
const bool is_mqa_gqa = built.is_mqa_gqa;
const bool has_dbias = built.has_dbias;
const BiasShape bias_shape = built.bias_shape;

bool ck_log_config = false;
if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) {
if (env_p != nullptr && std::string(env_p) == "1")
ck_log_config = true;
}
const char* dump_path = std::getenv("NVTE_DUMP_AITER_RT");
// print kernel name on verbose mode
ck_tile::stream_config stream_config{stream, dump_path!=nullptr, get_ck_log_stream() != nullptr};

// modify the max_seqlen_q for better performance in 0-length cases
// lse_workspace_ptr used as buffer
if(const char* env_p = std::getenv("NVTE_CK_RUNTIME_MAX_SEQLEN")) {
Expand Down
47 changes: 35 additions & 12 deletions transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,18 +99,12 @@ void dump_fwd_timings(const char* dump_path, float average_runtime){
file << average_runtime << "\n";
}

hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){

bool has_dropout = (args.is_training && args.dropout_probability > 0.f);

bool ck_log_config = false;
if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) {
if (env_p != nullptr && std::string(env_p) == "1")
ck_log_config = true;
}
const char* dump_path = std::getenv("NVTE_DUMP_AITER_RT");
// print kernel name on verbose mode
ck_tile::stream_config stream_config{stream, dump_path!=nullptr, get_ck_log_stream() != nullptr};
// Populate the AITER mha_fwd_args from TE's CKAttnFwdArgs. Shared by ck_attn_fwd
// (real launch) and ck_attn_fwd_uses_v3 (v3 availability probe) so the probe can
// never disagree with the launch. v3_api_check is left false here; callers flip it.
// The stream-dependent max_seqlen override (NVTE_CK_RUNTIME_MAX_SEQLEN) is applied
// by ck_attn_fwd after this returns; it does not affect v3 kernel selection.
aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){

bias_enum bias_type = bias_enum::no_bias;
BiasShape bias_shape = BiasShape::k11SS;
Expand Down Expand Up @@ -216,6 +210,35 @@ hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){
fmha_args.block_scale_size_q = 0;
fmha_args.block_scale_size_kv = 0;

return fmha_args;
}

// Probe whether AITER's v3 (asm) forward path will run for this config, without
// launching any kernel. Builds the same args as ck_attn_fwd and relies on AITER's
// v3_api_check dry-run (returns 1 when v3 is available, -1 otherwise).
bool ck_attn_fwd_uses_v3(const CKAttnFwdArgs& args){
aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);
fmha_args.v3_api_check = true;
// No kernel is launched in check mode, so the stream/log flags are irrelevant.
ck_tile::stream_config stream_config{nullptr, false, false};
return QOLA_NS(mha_fwd)(fmha_args, stream_config) == 1;
}

hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){

bool has_dropout = (args.is_training && args.dropout_probability > 0.f);

bool ck_log_config = false;
if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) {
if (env_p != nullptr && std::string(env_p) == "1")
ck_log_config = true;
}
const char* dump_path = std::getenv("NVTE_DUMP_AITER_RT");
// print kernel name on verbose mode
ck_tile::stream_config stream_config{stream, dump_path!=nullptr, get_ck_log_stream() != nullptr};

aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);

if(const char* env_p = std::getenv("NVTE_CK_RUNTIME_MAX_SEQLEN")){
if(args.is_group_mode() && std::string(env_p) == "1"){
if(ck_log_config){
Expand Down
Loading