resolve qy compactibility

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/adaptor/aten_adaptor.hpp

@@ -10,7 +10,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #endif
 
-#ifdef ENABLE_NVIDIA_API
+#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API)
 #include <ATen/cuda/CUDAContext.h>
 #endif
 

xmake.lua

@@ -471,25 +471,14 @@ target("infinicore_cpp_api")
         if has_config("metax-gpu") then
             add_deps("flash-attn-metax")
         end
+        if has_config("qy-gpu") then
+            add_deps("flash-attn-qy")
+        end
     end
 
-    -- MetaX: link pip-built flash_attn_2_cuda*.so.
-    -- The `.so` path resolver lives in `xmake/metax.lua` and reads:
-    -- - `FLASH_ATTN_2_CUDA_SO` (exact `.so` override)
-    -- - `FLASH_ATTN_METAX_CUDA_SO_CONTAINER` (override expected container path)
-    -- Path is fixed at target definition time (before_link sandbox has no os.iorunv in xmake 3.x).
-    if get_config("flash-attn") and get_config("flash-attn") ~= "" and has_config("metax-gpu") then
-        local flash_so_metax = _metax_flash_attn_cuda_so_path()
-        local flash_dir_metax = path.directory(flash_so_metax)
-        local flash_name_metax = path.filename(flash_so_metax)
-        before_link(function (target)
-            target:add(
-                "shflags",
-                "-Wl,--no-as-needed -L" .. flash_dir_metax .. " -l:" .. flash_name_metax .. " -Wl,-rpath," .. flash_dir_metax,
-                {force = true}
-            )
-        end)
-    end
+    -- Flash pip `.so` link flags: `before_link` runs in an xmake sandbox that cannot see helpers
+    -- from other included scripts; MetaX and QY each register their own hook in `xmake/metax.lua`
+    -- and `xmake/qy.lua`.
 
     before_build(function (target)
         if has_config("aten") then

xmake/metax.lua

@@ -9,6 +9,38 @@ do
     -- 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-aarch64-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
+
 -- Set numeric HPCC version macro for flash-attn signature/call compatibility.
 -- Must be done before compiling `infinicore_cpp_api` sources.
 target("infinicore_cpp_api")
@@ -39,41 +71,19 @@ target("infinicore_cpp_api")
             end
         end
     end)
-target_end()
-
--- Resolve MetaX flash-attn .so path.
--- `xmake.lua` calls `_metax_flash_attn_cuda_so_path()` during `infinicore_cpp_api`
--- linking, so this helper must live here (and must be global, not `local`).
-local FLASH_ATTN_METAX_CUDA_SO_CONTAINER_DEFAULT =
-    "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so"
-
-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
+    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 flash_dir_metax = path.directory(flash_so_metax)
+            local flash_name_metax = path.filename(flash_so_metax)
+            target:add(
+                "shflags",
+                "-Wl,--no-as-needed -L" .. flash_dir_metax .. " -l:" .. flash_name_metax .. " -Wl,-rpath," .. flash_dir_metax,
+                {force = true}
             )
-        )
+        end)
     end
-    return container_path
-end
+target_end()
 
 add_includedirs(MACA_ROOT .. "/include")
 add_linkdirs(MACA_ROOT .. "/lib")

xmake/qy.lua

@@ -218,3 +218,18 @@ target("flash-attn-qy")
         end)
     end
 target_end()
+
+if FLASH_ATTN_ROOT and FLASH_ATTN_ROOT ~= "" then
+    target("infinicore_cpp_api")
+        before_link(function (target)
+            local flash_so_qy = _qy_flash_attn_cuda_so_path()
+            local flash_dir_qy = path.directory(flash_so_qy)
+            local flash_name_qy = path.filename(flash_so_qy)
+            target:add(
+                "shflags",
+                "-Wl,--no-as-needed -L" .. flash_dir_qy .. " -l:" .. flash_name_qy .. " -Wl,-rpath," .. flash_dir_qy,
+                {force = true}
+            )
+        end)
+    target_end()
+end