From 446a85be5063f20c049ec1b335c7dc750c1009f5 Mon Sep 17 00:00:00 2001 From: Yongshu Bai Date: Tue, 8 Sep 2026 10:53:11 +0800 Subject: [PATCH 1/2] fix(metax): support flash-attn 2.6.3 forward ABI --- .../adaptor/flash_attention_adaptor.hpp | 54 +++++++++++-------- .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 8 ++- .../ops/multi_head_attention/mha_flashattn.cc | 17 +++--- .../mha_varlen_flashattn.cc | 13 +++-- 4 files changed, 58 insertions(+), 34 deletions(-) diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index b1e778e55..7b62f3b96 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -14,40 +14,45 @@ namespace flash { #endif std::vector -mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) +mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + std::optional &alibi_slopes_, // num_heads or batch_size x num_heads #if defined(ENABLE_METAX_API) - std::optional &softmax_lse_, // MetaX flash-attn dense fwd ABI includes an optional preallocated LSE tensor + std::optional &attn_mask_, #endif - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads const float p_dropout, const float softmax_scale, bool is_causal, int window_size_left, int window_size_right, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) const float softcap, +#endif const bool return_softmax, std::optional gen_ #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_, + bool return_max_logit_ #endif ); std::vector -mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - std::optional &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &cu_seqlens_q, // b+1 - const at::Tensor &cu_seqlens_k, // b+1 - std::optional &seqused_k, // b. If given, only this many elements of each batch element's keys are used. +mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + std::optional &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &cu_seqlens_q, // b+1 + const at::Tensor &cu_seqlens_k, // b+1 + std::optional &seqused_k, // b. If given, only this many elements of each batch element's keys are used. +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) std::optional &leftpad_k_, // batch_size std::optional &block_table_, // batch_size x max_num_blocks_per_seq - std::optional &alibi_slopes_, // num_heads or b x num_heads +#endif + std::optional &alibi_slopes_, // num_heads or b x num_heads int max_seqlen_q, const int max_seqlen_k, const float p_dropout, @@ -56,13 +61,16 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea bool is_causal, int window_size_left, int window_size_right, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) const float softcap, +#endif const bool return_softmax, std::optional gen_ #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_, + bool return_max_logit_ #endif ); @@ -123,21 +131,25 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size std::optional &rotary_cos_, // seqlen_ro x (rotary_dim / 2) std::optional &rotary_sin_, // seqlen_ro x (rotary_dim / 2) std::optional &cache_batch_idx_, // indices to index into the KV cache - std::optional &leftpad_k_, // batch_size - std::optional &block_table_, // batch_size x max_num_blocks_per_seq - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - std::optional &out_, // batch_size x seqlen_q x num_heads x head_size +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) + std::optional &leftpad_k_, // batch_size +#endif + std::optional &block_table_, // batch_size x max_num_blocks_per_seq + std::optional &alibi_slopes_, // num_heads or batch_size x num_heads + std::optional &out_, // batch_size x seqlen_q x num_heads x head_size const float softmax_scale, bool is_causal, int window_size_left, int window_size_right, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) const float softcap, +#endif bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 int num_splits #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_ #endif ); diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 32c25719e..39f773e4b 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -234,7 +234,7 @@ void run(void *planned_meta) { : std::optional(out_tensor); #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; + std::optional s_aux = std::nullopt; #endif auto result = INFINICORE_FLASH_OP(mha_fwd_kvcache)( @@ -247,7 +247,9 @@ void run(void *planned_meta) { rotary_cos, rotary_sin, cache_batch_idx, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) leftpad_k, +#endif block_table, alibi_slopes, out, @@ -255,12 +257,14 @@ void run(void *planned_meta) { true, -1, -1, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) 0.0f, +#endif false, 0 #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) , - flash_attn_mars_ext + s_aux #endif ); diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc index 5852e9f69..9805f9208 100644 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc @@ -66,15 +66,15 @@ void run(void *planned_meta) { auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); auto out = std::optional(out_work); -#if defined(ENABLE_METAX_API) - std::optional softmax_lse = std::nullopt; -#endif auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; auto scale = p->scale; auto is_causal = p->is_causal; +#if defined(ENABLE_METAX_API) + std::optional attn_mask = std::nullopt; +#endif #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; + std::optional s_aux = std::nullopt; #endif INFINICORE_FLASH_OP(mha_fwd) @@ -83,21 +83,24 @@ void run(void *planned_meta) { k, v, out, + alibi_slopes, #if defined(ENABLE_METAX_API) - softmax_lse, + attn_mask, #endif - alibi_slopes, 0.0, scale, is_causal, -1, -1, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) 0.0, +#endif false, std::nullopt #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) , - flash_attn_mars_ext + s_aux, + false #endif ); diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index 16fa1eb8b..858009735 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -198,11 +198,11 @@ void run(void *planned_meta) { const std::optional no_tensor; const std::optional block_table = p->block_table ? std::optional{ - p->infiniops_block_table->tensor(*p->block_table)} + p->infiniops_block_table->tensor(*p->block_table)} : std::nullopt; const std::optional alibi_slopes = p->alibi_slopes ? std::optional{ - p->infiniops_alibi_slopes->tensor(*p->alibi_slopes)} + p->infiniops_alibi_slopes->tensor(*p->alibi_slopes)} : std::nullopt; infini::ops::FlashAttnVarlenFunc::Call( @@ -304,7 +304,7 @@ void run(void *planned_meta) { auto scale = p->scale; #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; + std::optional s_aux = std::nullopt; #endif INFINICORE_FLASH_OP(mha_varlen_fwd) @@ -316,8 +316,10 @@ void run(void *planned_meta) { cu_seqlens_q, cu_seqlens_kv, seqused_k, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) leftpad_k, block_table, +#endif alibi_slopes, max_seqlen_q, max_seqlen_k, @@ -327,12 +329,15 @@ void run(void *planned_meta) { true, -1, -1, +#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) 0.0, +#endif false, std::nullopt #if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) , - flash_attn_mars_ext + s_aux, + false #endif ); From 010ccd3172b336974a96a90d23b9c4ede8b4b8ed Mon Sep 17 00:00:00 2001 From: Yongshu Bai Date: Tue, 8 Sep 2026 16:04:22 +0800 Subject: [PATCH 2/2] feat(metax): select flash-attn ABI from wheel symbols MetaX flash-attn ships two incompatible forward ABIs: flash_attn 2.5.3 wheels (MACA/HPCC 2.x, mha_fwd/mha_varlen_fwd/mha_fwd_kvcache take 13/18/18 args) and flash_attn 2.6.3+metax wheels (MACA/HPCC 3.x, the same functions take 16/23/21 args with softcap, leftpad_k, varlen block_table, s_aux and return_max_logit appended). Select the ABI at configure time from the wheel that will actually be linked: xmake/metax.lua inspects its demangled dynamic symbols in on_load (same approach as Cambricon in xmake/bang.lua) and injects INFINICORE_METAX_FA_ABI=253|263 as a public define, so it also reaches infinicore-test and removes the previous ODR hazard between the two targets. The HPCC/MACA Version.txt probe moves from xmake.lua into the same hook as a fallback when the wheel cannot be inspected, and a new --metax-fa-abi=auto|253|263 option allows overriding the result. The shared header derives a single INFINICORE_METAX_FA263 gate macro from the injected ABI (falling back to the HPCC major version when it is absent), replacing the scattered INFINICORE_HPCC_VERSION_MAJOR conditions in the declarations and the three flashattn call sites. Optional parameters are respelled through INFINICORE_FA_OPTIONAL (c10::optional under MetaX when exists, std::optional otherwise) so declarations mangle exactly like the wheel on any torch version; non-MetaX platforms expand to std::optional and are unaffected (verified by symbol comparison against origin/main). Also refuse paged KV explicitly on the 2.5.3 varlen path (no block_table parameter) instead of silently dropping it, and scope the kvcache leftpad_k local to the gate that uses it. Co-Authored-By: Claude Code --- .../adaptor/flash_attention_adaptor.hpp | 187 +++++++++++------- .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 10 +- .../ops/multi_head_attention/mha_flashattn.cc | 6 +- .../mha_varlen_flashattn.cc | 15 +- xmake.lua | 34 ++-- xmake/metax.lua | 109 +++++++++- 6 files changed, 247 insertions(+), 114 deletions(-) diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index 7b62f3b96..43bb28bdb 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -2,6 +2,46 @@ #pragma once #include "aten_adaptor.hpp" +#if defined(ENABLE_METAX_API) + +// MetaX flash-attn wheels declare their API in terms of c10::optional. On torch >= 2.1 that +// is an alias of std::optional (identical mangling); spelling the declarations the same way +// as the wheel guarantees identical mangling even on older torch stacks. When the c10 header +// no longer exists, torch has fully migrated to std::optional -- and so have wheels built +// against it. +#if __has_include() +#include +#define INFINICORE_FA_OPTIONAL c10::optional +#else +#define INFINICORE_FA_OPTIONAL std::optional +#endif + +// MetaX flash-attn ships two incompatible forward ABIs: +// INFINICORE_METAX_FA_ABI 253 -- flash_attn 2.5.3 (MACA/HPCC 2.x): +// mha_fwd / mha_varlen_fwd / mha_fwd_kvcache take 13 / 18 / 18 arguments +// INFINICORE_METAX_FA_ABI 263 -- flash_attn 2.6.3+metax (MACA/HPCC 3.x): +// the same functions take 16 / 23 / 21 arguments (softcap, leftpad_k, varlen +// block_table, s_aux and return_max_logit are appended) +// xmake (xmake/metax.lua) injects INFINICORE_METAX_FA_ABI after inspecting the dynamic +// symbols of the actual `flash_attn_2_cuda` wheel that will be linked. When it is not +// injected (e.g. the wheel is absent at configure time), fall back to the HPCC/MACA +// toolkit major-version probe: >= 3 ships the 2.6.3 ABI, otherwise the legacy 2.5.3 ABI. +#if !defined(INFINICORE_METAX_FA_ABI) +#if defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#define INFINICORE_METAX_FA_ABI 263 +#else +#define INFINICORE_METAX_FA_ABI 253 +#endif +#endif +#define INFINICORE_METAX_FA263 (INFINICORE_METAX_FA_ABI >= 263) + +#else // !ENABLE_METAX_API + +#define INFINICORE_FA_OPTIONAL std::optional +#define INFINICORE_METAX_FA263 1 + +#endif // ENABLE_METAX_API + // NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension // exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds // where the namespace is empty. @@ -14,45 +54,46 @@ namespace flash { #endif std::vector -mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads +mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + INFINICORE_FA_OPTIONAL &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + INFINICORE_FA_OPTIONAL &alibi_slopes_, // num_heads or batch_size x num_heads #if defined(ENABLE_METAX_API) - std::optional &attn_mask_, + INFINICORE_FA_OPTIONAL &attn_mask_, #endif const float p_dropout, const float softmax_scale, bool is_causal, int window_size_left, int window_size_right, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 const float softcap, #endif const bool return_softmax, - std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. + INFINICORE_FA_OPTIONAL gen_ +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 + // MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends these arguments vs the + // 2.5.3 wheel and upstream Dao-AILab flash-attn. , - std::optional &s_aux_, + INFINICORE_FA_OPTIONAL &s_aux_, bool return_max_logit_ #endif ); std::vector -mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - std::optional &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &cu_seqlens_q, // b+1 - const at::Tensor &cu_seqlens_k, // b+1 - std::optional &seqused_k, // b. If given, only this many elements of each batch element's keys are used. -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) - std::optional &leftpad_k_, // batch_size - std::optional &block_table_, // batch_size x max_num_blocks_per_seq +mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + INFINICORE_FA_OPTIONAL &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &cu_seqlens_q, // b+1 + const at::Tensor &cu_seqlens_k, // b+1 + INFINICORE_FA_OPTIONAL &seqused_k, // b. If given, only this many elements of each batch element's keys are used. +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 + INFINICORE_FA_OPTIONAL &leftpad_k_, // batch_size + INFINICORE_FA_OPTIONAL &block_table_, // batch_size x max_num_blocks_per_seq #endif - std::optional &alibi_slopes_, // num_heads or b x num_heads + INFINICORE_FA_OPTIONAL &alibi_slopes_, // num_heads or b x num_heads int max_seqlen_q, const int max_seqlen_k, const float p_dropout, @@ -61,53 +102,54 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x he bool is_causal, int window_size_left, int window_size_right, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 const float softcap, #endif const bool return_softmax, - std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. + INFINICORE_FA_OPTIONAL gen_ +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 + // MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends these arguments vs the + // 2.5.3 wheel and upstream Dao-AILab flash-attn. , - std::optional &s_aux_, + INFINICORE_FA_OPTIONAL &s_aux_, bool return_max_logit_ #endif ); std::vector -mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8) - const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &softmax_lse, // b x h x seqlen_q - std::optional &dq_, // batch_size x seqlen_q x num_heads x head_size - std::optional &dk_, // batch_size x seqlen_k x num_heads_k x head_size - std::optional &dv_, // batch_size x seqlen_k x num_heads_k x head_size - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - const float p_dropout, // probability to drop +mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8) + const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size + const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &softmax_lse, // b x h x seqlen_q + INFINICORE_FA_OPTIONAL &dq_, // batch_size x seqlen_q x num_heads x head_size + INFINICORE_FA_OPTIONAL &dk_, // batch_size x seqlen_k x num_heads_k x head_size + INFINICORE_FA_OPTIONAL &dv_, // batch_size x seqlen_k x num_heads_k x head_size + INFINICORE_FA_OPTIONAL &alibi_slopes_, // num_heads or batch_size x num_heads + const float p_dropout, // probability to drop const float softmax_scale, const bool is_causal, int window_size_left, int window_size_right, const float softcap, const bool deterministic, - std::optional gen_, - std::optional &rng_state); + INFINICORE_FA_OPTIONAL gen_, + INFINICORE_FA_OPTIONAL &rng_state); std::vector -mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size - const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &out, // total_q x num_heads x head_size - const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp - std::optional &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - std::optional &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - std::optional &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &cu_seqlens_q, // b+1 - const at::Tensor &cu_seqlens_k, // b+1 - std::optional &alibi_slopes_, // num_heads or b x num_heads +mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size + const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &out, // total_q x num_heads x head_size + const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp + INFINICORE_FA_OPTIONAL &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + INFINICORE_FA_OPTIONAL &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + INFINICORE_FA_OPTIONAL &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &cu_seqlens_q, // b+1 + const at::Tensor &cu_seqlens_k, // b+1 + INFINICORE_FA_OPTIONAL &alibi_slopes_, // num_heads or b x num_heads const int max_seqlen_q, const int max_seqlen_k, // max sequence length to choose the kernel const float p_dropout, // probability to drop @@ -118,38 +160,39 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, int window_size_right, const float softcap, const bool deterministic, - std::optional gen_, - std::optional &rng_state); + INFINICORE_FA_OPTIONAL gen_, + INFINICORE_FA_OPTIONAL &rng_state); std::vector -mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - std::optional &k_, // batch_size x seqlen_knew x num_heads_k x head_size - std::optional &v_, // batch_size x seqlen_knew x num_heads_k x head_size - std::optional &seqlens_k_, // batch_size - std::optional &rotary_cos_, // seqlen_ro x (rotary_dim / 2) - std::optional &rotary_sin_, // seqlen_ro x (rotary_dim / 2) - std::optional &cache_batch_idx_, // indices to index into the KV cache -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) - std::optional &leftpad_k_, // batch_size +mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + INFINICORE_FA_OPTIONAL &k_, // batch_size x seqlen_knew x num_heads_k x head_size + INFINICORE_FA_OPTIONAL &v_, // batch_size x seqlen_knew x num_heads_k x head_size + INFINICORE_FA_OPTIONAL &seqlens_k_, // batch_size + INFINICORE_FA_OPTIONAL &rotary_cos_, // seqlen_ro x (rotary_dim / 2) + INFINICORE_FA_OPTIONAL &rotary_sin_, // seqlen_ro x (rotary_dim / 2) + INFINICORE_FA_OPTIONAL &cache_batch_idx_, // indices to index into the KV cache +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 + INFINICORE_FA_OPTIONAL &leftpad_k_, // batch_size #endif - std::optional &block_table_, // batch_size x max_num_blocks_per_seq - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - std::optional &out_, // batch_size x seqlen_q x num_heads x head_size + INFINICORE_FA_OPTIONAL &block_table_, // batch_size x max_num_blocks_per_seq + INFINICORE_FA_OPTIONAL &alibi_slopes_, // num_heads or batch_size x num_heads + INFINICORE_FA_OPTIONAL &out_, // batch_size x seqlen_q x num_heads x head_size const float softmax_scale, bool is_causal, int window_size_left, int window_size_right, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 const float softcap, #endif bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 int num_splits -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 + // MetaX `flash_attn_2_cuda` 2.6.3+ (MACA/HPCC 3.x) appends this argument vs the + // 2.5.3 wheel and upstream Dao-AILab flash-attn. , - std::optional &s_aux_ + INFINICORE_FA_OPTIONAL &s_aux_ #endif ); diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 39f773e4b..41376f44a 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -224,7 +224,9 @@ void run(void *planned_meta) { std::optional rotary_cos = std::nullopt; std::optional rotary_sin = std::nullopt; std::optional cache_batch_idx = std::nullopt; +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 std::optional leftpad_k = std::nullopt; +#endif const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 && q.size(1) == 1 && q.size(2) > k_cache.size(2) @@ -233,7 +235,7 @@ void run(void *planned_meta) { auto out = use_dynamic_out ? std::optional(std::nullopt) : std::optional(out_tensor); -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 std::optional s_aux = std::nullopt; #endif @@ -247,7 +249,7 @@ void run(void *planned_meta) { rotary_cos, rotary_sin, cache_batch_idx, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 leftpad_k, #endif block_table, @@ -257,12 +259,12 @@ void run(void *planned_meta) { true, -1, -1, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 0.0f, #endif false, 0 -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 , s_aux #endif diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc index 9805f9208..748799528 100644 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc @@ -73,7 +73,7 @@ void run(void *planned_meta) { #if defined(ENABLE_METAX_API) std::optional attn_mask = std::nullopt; #endif -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 std::optional s_aux = std::nullopt; #endif @@ -92,12 +92,12 @@ void run(void *planned_meta) { is_causal, -1, -1, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 0.0, #endif false, std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 , s_aux, false diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index 858009735..3e57da54a 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -294,6 +294,13 @@ void run(void *planned_meta) { } #ifdef ENABLE_FLASH_ATTN +#if defined(ENABLE_METAX_API) && !INFINICORE_METAX_FA263 + // The MetaX flash-attn 2.5.3 (MACA/HPCC 2.x) varlen ABI has no block_table parameter; + // refuse paged KV explicitly instead of silently computing attention over the wrong keys. + if (p->block_table) { + throw std::runtime_error("paged KV varlen attention requires MetaX flash-attn 2.6.3+ (MACA/HPCC 3.x)"); + } +#endif auto out = std::optional(out_work); std::optional seqused_k = std::nullopt; std::optional leftpad_k = std::nullopt; @@ -303,7 +310,7 @@ void run(void *planned_meta) { auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; auto scale = p->scale; -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 std::optional s_aux = std::nullopt; #endif @@ -316,7 +323,7 @@ void run(void *planned_meta) { cu_seqlens_q, cu_seqlens_kv, seqused_k, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 leftpad_k, block_table, #endif @@ -329,12 +336,12 @@ void run(void *planned_meta) { true, -1, -1, -#if !defined(ENABLE_METAX_API) || (defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)) +#if !defined(ENABLE_METAX_API) || INFINICORE_METAX_FA263 0.0, #endif false, std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && INFINICORE_METAX_FA263 , s_aux, false diff --git a/xmake.lua b/xmake.lua index 568410287..07c17c85b 100644 --- a/xmake.lua +++ b/xmake.lua @@ -265,6 +265,14 @@ option("flash-attn") set_description("Path to flash-attention repo. If not set, flash-attention will not used.") option_end() +-- MetaX flash-attn ABI selection +option("metax-fa-abi") + set_default("auto") + set_showmenu(true) + set_values("auto", "253", "263") + set_description("MetaX flash-attn ABI: auto (inspect the wheel's symbols), 253 (flash_attn 2.5.3 / MACA 2.x) or 263 (flash_attn 2.6.3+metax / MACA 3.x)") +option_end() + if has_config("aten") then add_defines("ENABLE_ATEN") if has_config("iluvatar-gpu") then @@ -863,29 +871,9 @@ target("infinicore_cpp_api") stage_infiniops_runtime(os, path.join(INFINI_ROOT, "lib")) end - -- MetaX + flash-attn: `flash_attn_2_cuda` may use a different `mha_fwd_kvcache` ABI - -- depending on the underlying stack version. When building with MACA (`--use-mc=y`), - -- the version file is typically `/opt/maca/Version.txt` (HPCC uses `/opt/hpcc/Version.txt`). - if has_config("metax-gpu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then - local version_txt = "/opt/hpcc/Version.txt" - if not os.isfile(version_txt) and has_config("use-mc") then - version_txt = "/opt/maca/Version.txt" - end - if os.isfile(version_txt) then - local content = os.iorunv("cat", {version_txt}) or "" - content = content:trim() - local major_str = content:match("Version:(%d+)") or content:match("^(%d+)") - if major_str and major_str ~= "" then - local major = tonumber(major_str) - if major then - local define = "INFINICORE_HPCC_VERSION_MAJOR=" .. tostring(major) - target:add("defines", define) - target:add("cxflags", "-D" .. define) - target:add("cxxflags", "-D" .. define) - end - end - end - end + -- MetaX + flash-attn ABI selection (wheel symbol probe with HPCC/MACA Version.txt + -- fallback) lives in `on_load` of `infinicore_cpp_api` in `xmake/metax.lua`, so the + -- define is public and also reaches `infinicore-test`. if has_config("aten") then local outdata = os.iorunv(PYTHON, {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() diff --git a/xmake/metax.lua b/xmake/metax.lua index 85407ed1b..2c4e078ec 100644 --- a/xmake/metax.lua +++ b/xmake/metax.lua @@ -2,12 +2,13 @@ local MACA_ROOT = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT") local FLASH_ATTN_ROOT = get_config("flash-attn") --- MetaX flash-attn (pip `flash_attn_2_cuda`) may append an extra trailing argument --- (`flash_attn_mars_ext_`) depending on the underlying HPCC/MetaX stack version. -do - -- Intentionally empty: HPCC version parsing is deferred to `before_build` - -- on `infinicore_cpp_api` where `os.iorunv` is available in this xmake sandbox. -end +-- MetaX flash-attn (pip `flash_attn_2_cuda`) ships two incompatible forward ABIs +-- (see include/infinicore/adaptor/flash_attention_adaptor.hpp): +-- 253: flash_attn 2.5.3 wheels (MACA/HPCC 2.x) -- mha_fwd/mha_varlen_fwd/mha_fwd_kvcache take 13/18/18 args +-- 263: flash_attn 2.6.3+metax wheels (MACA/HPCC 3.x) -- the same functions take 16/23/21 args +-- The wheel `.so` that will actually be linked is the ground truth, so its dynamic symbols +-- are inspected at load time (same approach as Cambricon in `xmake/bang.lua`); the HPCC/MACA +-- toolkit Version.txt is only a fallback when the wheel cannot be inspected. -- Resolve MetaX flash-attn .so path (used only from this file: `before_link` sandbox cannot see globals from `xmake.lua`). local FLASH_ATTN_METAX_CUDA_SO_CONTAINER_DEFAULT = @@ -41,10 +42,102 @@ local function metax_flash_attn_cuda_so_path() return container_path end --- MetaX flash-attn link flags for pip `flash_attn_2_cuda`. --- Version/ABI macros are set in `xmake.lua` for `infinicore_cpp_api` so they apply to all sources. +-- Classify a MetaX flash-attn wheel as "253"/"263" from its demangled dynamic symbols. +-- Returns nil + reason when the wheel cannot be inspected or recognized. +-- `run(program, argv)` is injected by the caller: xmake (>= 3.x) restricts the script-body +-- sandbox (no os.iorunv/pcall/try there), while hook functions get the full sandbox, so the +-- on_load hook below passes a runner built from its own environment into these body-level +-- helpers. The runner returns nil instead of raising when the program is unavailable. +local function metax_detect_flash_attn_abi(so_path, run) + if not so_path or not os.isfile(so_path) then + return nil, "wheel .so not found" + end + local symbols = run("nm", {"-D", "-C", "--defined-only", so_path}) + if not symbols or symbols == "" then + return nil, "could not read the dynamic symbols of " .. so_path + end + if not symbols:find("mha_fwd(at::Tensor&", 1, true) then + return nil, "no demangled mha_fwd symbol in " .. so_path + end + -- Both markers exist only in flash_attn 2.6.3+metax wheels: + -- * leftpad_k (`optional&`) in mha_varlen_fwd -- primary evidence, + -- it is one of the params the 2.6.3 ABI appends (substring matches the + -- `std::optional` and `c10::optional` demangled spellings alike); + -- * mha_fwd_kvcache_dequant -- corroboration. + local varlen_sig = symbols:match("[^\n]*mha_varlen_fwd%(([^\n]*)") or "" + local has_leftpad = varlen_sig:find("optional&", 1, true) ~= nil + local has_dequant = symbols:find("mha_fwd_kvcache_dequant(", 1, true) ~= nil + if has_leftpad ~= has_dequant then + print(string.format( + "warning: metax+flash-attn: inconsistent ABI markers in %s (varlen leftpad_k=%s, mha_fwd_kvcache_dequant=%s); trusting leftpad_k", + so_path, tostring(has_leftpad), tostring(has_dequant))) + end + if has_leftpad then + return "263" + end + return "253" +end + +-- Legacy fallback: HPCC (`/opt/hpcc/Version.txt`) or MACA (`/opt/maca/Version.txt`, with +-- `--use-mc=y`) toolkit major version. MACA/HPCC 3.x stacks ship flash_attn 2.6.3+metax. +local function metax_stack_version_major(run) + local version_txt = "/opt/hpcc/Version.txt" + if not os.isfile(version_txt) and has_config("use-mc") then + version_txt = "/opt/maca/Version.txt" + end + if not os.isfile(version_txt) then + return nil + end + local content = run("cat", {version_txt}) or "" + content = content:trim() + local major_str = content:match("Version:(%d+)") or content:match("^(%d+)") + if major_str and major_str ~= "" then + return tonumber(major_str) + end + return nil +end + +-- MetaX flash-attn ABI selection + link flags for pip `flash_attn_2_cuda`. +-- `INFINICORE_METAX_FA_ABI` is added {public = true} so it also reaches `infinicore-test`, +-- which depends on this target and compiles the same `mha_*_flashattn.cc` sources. target("infinicore_cpp_api") if get_config("flash-attn") and get_config("flash-attn") ~= "" then + on_load(function (target) + -- This hook body runs in xmake's full sandbox (unlike the restricted + -- script-body scope), so os.iorunv/import are available here. + local find_program = import("lib.detect.find_program") + local function run(program, argv) + if not find_program(program) then + return nil + end + return os.iorunv(program, argv) + end + local abi = get_config("metax-fa-abi") + if not abi or abi == "" or abi == "auto" then + local so_path = metax_flash_attn_cuda_so_path() + local detected, why = metax_detect_flash_attn_abi(so_path, run) + if detected then + abi = detected + print(string.format("metax+flash-attn: %s ABI detected from wheel symbols: %s", abi, so_path)) + else + local major = metax_stack_version_major(run) + if major then + -- Header derives the ABI from the toolkit major version (>= 3 -> 263). + target:add("defines", "INFINICORE_HPCC_VERSION_MAJOR=" .. tostring(major), {public = true}) + print(string.format( + "metax+flash-attn: could not inspect the wheel (%s); falling back to HPCC/MACA major version %d", + why or "unknown reason", major)) + else + print(string.format( + "warning: metax+flash-attn: could not inspect the wheel (%s) and found no HPCC/MACA Version.txt; defaulting to the flash_attn 2.5.3 ABI", + why or "unknown reason")) + end + end + end + if abi == "253" or abi == "263" then + target:add("defines", "INFINICORE_METAX_FA_ABI=" .. abi, {public = true}) + end + end) before_link(function (target) local flash_so_metax = metax_flash_attn_cuda_so_path() local flash_dir_metax = path.directory(flash_so_metax)