@@ -147,14 +147,10 @@ target("flash-attn-nvidia")
147147
148148 before_build (function (target )
149149 if FLASH_ATTN_ROOT ~= nil then
150- local outdata = os .iorunv (" python" , {" -c" , " import torch, os; print(os.path.dirname(torch.__file__))" }):trim ()
151- local TORCH_DIR = outdata
152-
153- local outdata = os .iorunv (" python" , {" -c" , " import sysconfig; print(sysconfig.get_paths()['include'])" }):trim ()
154- local PYTHON_INCLUDE = outdata
155-
156- local outdata = os .iorunv (" python" , {" -c" , " import sysconfig; print(sysconfig.get_config_var('LIBDIR'))" }):trim ()
157- local PYTHON_LIB_DIR = outdata
150+ local TORCH_DIR = os .iorunv (" python" , {" -c" , " import torch, os; print(os.path.dirname(torch.__file__))" }):trim ()
151+ local PYTHON_INCLUDE = os .iorunv (" python" , {" -c" , " import sysconfig; print(sysconfig.get_paths()['include'])" }):trim ()
152+ local PYTHON_LIB_DIR = os .iorunv (" python" , {" -c" , " import sysconfig; print(sysconfig.get_config_var('LIBDIR'))" }):trim ()
153+ local LIB_PYTHON = os .iorunv (" python" , {" -c" , " import sysconfig, os; print(sysconfig.get_config_var('LDLIBRARY'))" }):trim ()
158154 -- Include dirs
159155 target :add (" includedirs" , FLASH_ATTN_ROOT .. " /csrc/flash_attn/src" , {public = false })
160156 target :add (" includedirs" , TORCH_DIR .. " /include/torch/csrc/api/include" , {public = false })
@@ -165,7 +161,7 @@ target("flash-attn-nvidia")
165161
166162 -- Link libraries
167163 target :add (" linkdirs" , TORCH_DIR .. " /lib" , PYTHON_LIB_DIR )
168- target :add (" links" , " torch" , " torch_cuda" , " torch_cpu" , " c10" , " c10_cuda" , " torch_python" , " python3 " )
164+ target :add (" links" , " torch" , " torch_cuda" , " torch_cpu" , " c10" , " c10_cuda" , " torch_python" , LIB_PYTHON )
169165
170166 end
171167
0 commit comments