diff --git a/docs/metax-flash-attention.md b/docs/metax-flash-attention.md new file mode 100644 index 000000000..0e11807cd --- /dev/null +++ b/docs/metax-flash-attention.md @@ -0,0 +1,35 @@ +# MetaX Flash Attention ABI selection + +MetaX wheels can export different C++ signatures for `mha_varlen_fwd` even +within MACA 3.x. A missing trailing `bool` causes an undefined-symbol error +when loading InfiniCore, before any attention kernel can execute. + +When building MetaX with Flash Attention, the build now inspects the actual +linked extension using `nm -D -C --defined-only`. It selects the declaration +and call for three known signatures: the original arguments, an additional +optional tensor, or that tensor plus `return_max_logit`. The last argument is +passed as `false` when present. NVIDIA and other device adapters are unchanged. + +Use `FLASH_ATTN_2_CUDA_SO=/absolute/path/to/flash_attn_2_cuda.so` to select an +exact wheel. Otherwise the resolver checks `--flash-attn` for a unique +extension, then `FLASH_ATTN_METAX_CUDA_SO_CONTAINER`, then the Python environment. +Invalid explicit paths and unknown/missing/ambiguous signatures fail during +build with a diagnostic. The same resolver is used for compilation and linking. +MACA-version handling for the other Flash Attention entry points remains as +before; this change only detects the varlen ABI. + +Run the detection tests without a GPU: + +```sh +xmake lua tests/xmake/test_metax_flash_abi.lua +``` + +With `FLASH_ATTN_2_CUDA_SO` set, the test also checks the installed extension. +The three recognized signatures and missing/unknown/ambiguous negative cases +are covered by fixtures. Older-wheel execution was not available for testing. + +The current-bool variant passed C500 FP16/BF16 Prefill and Decode reference +checks. Older-wheel execution was not available. Hardware conditions, model +results and archived backend comparisons are in +[PR #1558](https://github.com/InfiniTensor/InfiniCore/pull/1558). The fix restores +loading of affected wheels; it does not introduce a faster attention kernel. diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index b1e778e55..0db925272 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -59,11 +59,15 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea const float softcap, 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. +#if defined(ENABLE_METAX_API) && defined(INFINICORE_METAX_VARLEN_EXT) + // Selected from the linked extension signature at build time. , std::optional &flash_attn_mars_ext_ #endif +#if defined(ENABLE_METAX_API) && defined(INFINICORE_METAX_VARLEN_RETURN_MAX_LOGIT) + , + bool return_max_logit +#endif ); std::vector 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..ffe97ec67 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( @@ -303,7 +303,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) && defined(INFINICORE_METAX_VARLEN_EXT) std::optional flash_attn_mars_ext = std::nullopt; #endif @@ -330,9 +330,13 @@ void run(void *planned_meta) { 0.0, false, std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#if defined(ENABLE_METAX_API) && defined(INFINICORE_METAX_VARLEN_EXT) , flash_attn_mars_ext +#endif +#if defined(ENABLE_METAX_API) && defined(INFINICORE_METAX_VARLEN_RETURN_MAX_LOGIT) + , + false #endif ); diff --git a/tests/xmake/test_metax_flash_abi.lua b/tests/xmake/test_metax_flash_abi.lua new file mode 100644 index 000000000..1d9c3a1ee --- /dev/null +++ b/tests/xmake/test_metax_flash_abi.lua @@ -0,0 +1,24 @@ +function main() + local abi = import("metax_flash_abi", {rootdir = path.join(os.projectdir(), "xmake")}) + local common = "mha_varlen_fwd(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional&, at::Tensor const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, float, bool, std::optional" + local function symbol(suffix) return "0000000000123456 T " .. common .. suffix .. "\n" end + local old = abi.detect_symbols(symbol(")")) + assert(not old.extension and not old.return_max_logit) + local ext = abi.detect_symbols(symbol(", std::optional&)")) + assert(ext.extension and not ext.return_max_logit) + local latest = abi.detect_symbols(symbol(", std::optional&, bool)")) + assert(latest.extension and latest.return_max_logit) + for _, bad in ipairs({"", symbol(", int)"), symbol(")") .. symbol(", std::optional&)")}) do + local failed = false + try {function() abi.detect_symbols(bad) end, + catch {function() failed = true end}} + assert(failed) + end + if os.getenv("FLASH_ATTN_2_CUDA_SO") then + local file = abi.resolve() + local actual = abi.detect_symbols(os.iorunv("nm", {"-D", "-C", "--defined-only", file})) + print("installed extension: " .. file) + print(actual) + end + print("MetaX ABI: three signatures accepted; missing/unknown/ambiguous rejected") +end diff --git a/xmake.lua b/xmake.lua index 568410287..80909cdd2 100644 --- a/xmake.lua +++ b/xmake.lua @@ -867,6 +867,8 @@ target("infinicore_cpp_api") -- 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 abi = import("metax_flash_abi", {rootdir = path.join(os.projectdir(), "xmake")}) + abi.configure(target, get_config("flash-attn")) 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" diff --git a/xmake/metax.lua b/xmake/metax.lua index 85407ed1b..0bfdbaeca 100644 --- a/xmake/metax.lua +++ b/xmake/metax.lua @@ -2,51 +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 - --- 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 = - "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so" - -local function metax_flash_attn_cuda_so_path() - -- Highest priority: override the exact `.so` file to link. - local env_path = os.getenv("FLASH_ATTN_2_CUDA_SO") - if env_path and env_path ~= "" then - env_path = env_path:trim() - if os.isfile(env_path) then - return env_path - end - print(string.format("warning: metax+flash-attn: FLASH_ATTN_2_CUDA_SO is not a file: %s, fallback to container/default path", env_path)) - end - - -- Second priority: allow overriding the "expected" container path via env. - local container_path = os.getenv("FLASH_ATTN_METAX_CUDA_SO_CONTAINER") - if not container_path or container_path == "" then - container_path = FLASH_ATTN_METAX_CUDA_SO_CONTAINER_DEFAULT - end - - if not os.isfile(container_path) then - print( - string.format( - "warning: metax+flash-attn: expected %s; install flash-attn in conda env, or export FLASH_ATTN_2_CUDA_SO.", - container_path - ) - ) - end - 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. target("infinicore_cpp_api") if get_config("flash-attn") and get_config("flash-attn") ~= "" then before_link(function (target) - local flash_so_metax = metax_flash_attn_cuda_so_path() + local abi = import("metax_flash_abi", {rootdir = path.join(os.projectdir(), "xmake")}) + local flash_so_metax = abi.resolve(get_config("flash-attn")) local flash_dir_metax = path.directory(flash_so_metax) local flash_name_metax = path.filename(flash_so_metax) target:add( diff --git a/xmake/metax_flash_abi.lua b/xmake/metax_flash_abi.lua new file mode 100644 index 000000000..bb4c2f73d --- /dev/null +++ b/xmake/metax_flash_abi.lua @@ -0,0 +1,67 @@ +-- Resolve and inspect the exact extension linked into infinicore_cpp_api. +-- MACA's release number alone does not determine the varlen C++ ABI. +function resolve(root) + local override = os.getenv("FLASH_ATTN_2_CUDA_SO") + if override and override ~= "" then + assert(os.isfile(override), "FLASH_ATTN_2_CUDA_SO is not a file: " .. override) + return override + end + if root and root ~= "" then + local files = os.files(path.join(root, "flash_attn_2_cuda*.so")) + assert(#files <= 1, "Multiple MetaX extensions found; set FLASH_ATTN_2_CUDA_SO") + if #files == 1 then return files[1] end + end + local container = os.getenv("FLASH_ATTN_METAX_CUDA_SO_CONTAINER") + if container and container ~= "" then + assert(os.isfile(container), "MetaX container extension is not a file: " .. container) + return container + end + local python = os.getenv("PYTHON") or "python" + local file = os.iorunv(python, {"-c", + "import importlib.util; s=importlib.util.find_spec('flash_attn_2_cuda'); print(s.origin if s else '')"}):trim() + assert(os.isfile(file), "MetaX flash_attn_2_cuda not found; set FLASH_ATTN_2_CUDA_SO") + return file +end + +function detect_symbols(symbols) + local parameters = { + "at::Tensor&", "at::Tensor const&", "at::Tensor const&", + "std::optional&", "at::Tensor const&", "at::Tensor const&", + "std::optional&", "std::optional&", + "std::optional&", "std::optional&", + "int", "int", "float", "float", "bool", "bool", "int", "int", + "float", "bool", "std::optional" + } + local base = "mha_varlen_fwd(" .. table.concat(parameters, ", ") + local signatures = { + [base .. ")"] = {extension = false, return_max_logit = false}, + [base .. ", std::optional&)"] = {extension = true, return_max_logit = false}, + [base .. ", std::optional&, bool)"] = {extension = true, return_max_logit = true} + } + local found = nil + for line in symbols:gmatch("[^\r\n]+") do + local signature = line:match("^%s*%x+%s+[TW]%s+(.+)$") + if signature and signature:find("mha_varlen_fwd(", 1, true) == 1 then + assert(signatures[signature], "Unsupported MetaX Flash Attention ABI: " .. signature) + assert(not found, "Ambiguous MetaX mha_varlen_fwd overloads") + found = signatures[signature] + end + end + assert(found, "MetaX extension does not export a supported mha_varlen_fwd; check FLASH_ATTN_2_CUDA_SO") + return found +end + +function configure(target, root) + local file = resolve(root) + local abi = detect_symbols(os.iorunv("nm", {"-D", "-C", "--defined-only", file})) + local defines = {} + if abi.extension then table.insert(defines, "INFINICORE_METAX_VARLEN_EXT") end + if abi.return_max_logit then table.insert(defines, "INFINICORE_METAX_VARLEN_RETURN_MAX_LOGIT") end + for _, define in ipairs(defines) do + target:add("defines", define) + target:add("cxflags", "-D" .. define) + target:add("cxxflags", "-D" .. define) + end + print(string.format("MetaX varlen ABI: extension=%s, return_max_logit=%s (%s)", + tostring(abi.extension), tostring(abi.return_max_logit), file)) +end