fix(metax): detect Flash Attention varlen ABI from linked library

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.

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<at::Generator> 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<at::Tensor> &flash_attn_mars_ext_
 #endif
+#if defined(ENABLE_METAX_API) && defined(INFINICORE_METAX_VARLEN_RETURN_MAX_LOGIT)
+               ,
+               bool return_max_logit
+#endif
 );
 
 std::vector<at::Tensor>

src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc

@@ -198,11 +198,11 @@ void run(void *planned_meta) {
         const std::optional<infini::ops::Tensor> no_tensor;
         const std::optional<infini::ops::Tensor> block_table = p->block_table
                                                                  ? std::optional<infini::ops::Tensor>{
-                                                                     p->infiniops_block_table->tensor(*p->block_table)}
+                                                                       p->infiniops_block_table->tensor(*p->block_table)}
                                                                  : std::nullopt;
         const std::optional<infini::ops::Tensor> alibi_slopes = p->alibi_slopes
                                                                   ? std::optional<infini::ops::Tensor>{
-                                                                      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<at::Tensor>(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<at::Tensor> 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
     );
 

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>&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor>&, std::optional<at::Tensor const>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, int, int, float, float, bool, bool, int, int, float, bool, std::optional<at::Generator>"
+    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<at::Tensor>&)"))
+    assert(ext.extension and not ext.return_max_logit)
+    local latest = abi.detect_symbols(symbol(", std::optional<at::Tensor>&, bool)"))
+    assert(latest.extension and latest.return_max_logit)
+    for _, bad in ipairs({"", symbol(", int)"), symbol(")") .. symbol(", std::optional<at::Tensor>&)")}) 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

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"

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(

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>&", "at::Tensor const&", "at::Tensor const&",
+        "std::optional<at::Tensor>&", "std::optional<at::Tensor const>&",
+        "std::optional<at::Tensor>&", "std::optional<at::Tensor>&",
+        "int", "int", "float", "float", "bool", "bool", "int", "int",
+        "float", "bool", "std::optional<at::Generator>"
+    }
+    local base = "mha_varlen_fwd(" .. table.concat(parameters, ", ")
+    local signatures = {
+        [base .. ")"] = {extension = false, return_max_logit = false},
+        [base .. ", std::optional<at::Tensor>&)"] = {extension = true, return_max_logit = false},
+        [base .. ", std::optional<at::Tensor>&, 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