From 6e10d18baec229a1a29c9ad9c78e4cd8a037630a Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 7 Jul 2026 17:26:02 +0000 Subject: [PATCH 1/4] Added V3 API check mechanism --- .../include/ck_fused_attn/ck_fused_attn.hpp | 6 +++ .../ck_fused_attn/src/ck_fused_attn_bwd.cpp | 52 +++++++++++++++---- .../ck_fused_attn/src/ck_fused_attn_fwd.cpp | 47 ++++++++++++----- 3 files changed, 83 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp index 127d75b4ca..150c8bfb32 100644 --- a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp +++ b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp @@ -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); + }//namespace ck_fused_attn #endif // CK_FUSED_ATTN_H diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp index b0ed9ae6e5..f3ea8d2721 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp @@ -441,21 +441,17 @@ 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){ +// 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. +static aiter::mha_bwd_args 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()) { @@ -578,6 +574,42 @@ 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{args.philox_seed_ptr, args.philox_offset_ptr}; + return fmha_args; +} + +// 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.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){ + + // Recomputed here (also computed inside build_bwd_fmha_args) for the post-dispatch + // MQA/GQA reduction and dbias handling below. + bool is_mqa_gqa = (args.h > args.hg); + bool has_dbias = args.dbias_ptr != nullptr; + BiasShape bias_shape = BiasShape::k11SS; + if (!args.is_group_mode()) { + std::tie(std::ignore, bias_shape) = get_ck_bias_type_shape(&args); + } + + 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_bwd_args fmha_args = build_bwd_fmha_args(args); + // 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")) { diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp index 0f407230c0..82955394d0 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp @@ -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. +static aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){ bias_enum bias_type = bias_enum::no_bias; BiasShape bias_shape = BiasShape::k11SS; @@ -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){ From 8ca1c99baec732a1d3d796f44ddff0174d1ae9e9 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 7 Jul 2026 17:30:37 +0000 Subject: [PATCH 2/4] Updated to add QoLA-side AITER patch for V3 API checks --- 3rdparty/QoLA | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/QoLA b/3rdparty/QoLA index 47e0f5501a..c4393ec60e 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit 47e0f5501af51521c8a911fbbefa4e24b92dd7b6 +Subproject commit c4393ec60e177362374106d1a39700c6d53e3215 From 0698c9a41b9e4fc1bb3136e6ccaa16640e4a2c69 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 7 Jul 2026 19:01:38 +0000 Subject: [PATCH 3/4] Updated qola commit to main --- 3rdparty/QoLA | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/QoLA b/3rdparty/QoLA index c4393ec60e..0d6665b6d8 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit c4393ec60e177362374106d1a39700c6d53e3215 +Subproject commit 0d6665b6d87feeafe4b782dcc736c5e755dc5b34 From cd1ac6970c0f6684910f8ad5cf47612752b7bc99 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Fri, 10 Jul 2026 17:40:40 +0000 Subject: [PATCH 4/4] Removed static, updated arg structure for build_bwd_fmha_args --- .../ck_fused_attn/src/ck_fused_attn_bwd.cpp | 30 +++++++++++-------- .../ck_fused_attn/src/ck_fused_attn_fwd.cpp | 2 +- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp index f3ea8d2721..a7425be988 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp @@ -441,12 +441,19 @@ void dump_bwd_timings(const char* dump_path, float average_runtime){ file << average_runtime << "\n"; } +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. -static aiter::mha_bwd_args build_bwd_fmha_args(const CkAttnBwdArgs& args){ +BwdFmhaArgs build_bwd_fmha_args(const CkAttnBwdArgs& args){ bool has_dropout = (args.dropout_probability > 0.f); bool has_dbias = args.dbias_ptr != nullptr; @@ -574,14 +581,14 @@ static aiter::mha_bwd_args build_bwd_fmha_args(const CkAttnBwdArgs& args){ fmha_args.p_undrop = 1.0 - args.dropout_probability; fmha_args.drop_seed_offset = std::pair{args.philox_seed_ptr, args.philox_offset_ptr}; - return fmha_args; + 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); + 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}; @@ -590,14 +597,13 @@ bool ck_attn_bwd_uses_v3(const CkAttnBwdArgs& args){ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ - // Recomputed here (also computed inside build_bwd_fmha_args) for the post-dispatch - // MQA/GQA reduction and dbias handling below. - bool is_mqa_gqa = (args.h > args.hg); - bool has_dbias = args.dbias_ptr != nullptr; - BiasShape bias_shape = BiasShape::k11SS; - if (!args.is_group_mode()) { - std::tie(std::ignore, bias_shape) = get_ck_bias_type_shape(&args); - } + // 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") ) { @@ -608,8 +614,6 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ // print kernel name on verbose mode ck_tile::stream_config stream_config{stream, dump_path!=nullptr, get_ck_log_stream() != nullptr}; - aiter::mha_bwd_args fmha_args = build_bwd_fmha_args(args); - // 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")) { diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp index 82955394d0..88e3bac91a 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp @@ -104,7 +104,7 @@ void dump_fwd_timings(const char* dump_path, float average_runtime){ // 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. -static aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){ +aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){ bias_enum bias_type = bias_enum::no_bias; BiasShape bias_shape = BiasShape::k11SS;