Skip to content

[Build] Compile FlashAttention on arm64 platforms #79488

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SigureMo merged 3 commits into from
Jul 17, 2026
+3 −2
Merged

[Build] Compile FlashAttention on arm64 platforms #79488

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
469767f
Update GPU third-party conditions
liuhao2638 Jul 16, 2026
9034ce9
ci: rerun PR checks
liuhao2638 Jul 17, 2026
abe4c2e
fix flashattn cache on arm builds
liuhao2638 Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions cmake/external/flashattn.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,9 @@ else()
set(CACHE_TAR_DIR "${FA_BUILD_DIR}/flashattn_libs_${FLASHATTN_TAG}")

set(SKIP_BUILD_FA OFF)
if(WITH_ARM)
set(WITH_FA_BUILD_WITH_CACHE OFF)
endif()
if(WITH_FA_BUILD_WITH_CACHE)

message(STATUS "Downloading ${TAR_FILE_URL} to ${CACHE_TAR_PATH}")
Expand Down
2 changes: 0 additions & 2 deletions cmake/third_party.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,6 @@ include(external/glog) # download, build, install glog

########################### include third_party according to flags ###############################
if(WITH_GPU
AND NOT WITH_ARM
AND NOT WIN32
AND NOT APPLE)
if(${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 11.0)
Expand Down Expand Up @@ -601,7 +600,6 @@ if(WITH_ROCM)
endif()

if(WITH_GPU

@risemeup1111 risemeup1111 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1

处理要求：请针对该评论修复并提交新的 commit。

这里放开 ARM 后，CUDA ARM 构建会进入 FlashAttention 外部依赖块，但 cmake/external/flashattn.cmake 的缓存下载路径只由 NVCC_ARCH_BIN 和 FLASHATTN_TAG 组成（https://.../flash-attention/cu${FA_NVCC_ARCH_BIN}/flashattn_libs_${FLASHATTN_TAG}.tar），没有区分 aarch64/x86_64。下载成功时还会直接复制缓存里的 .so 并设置 SKIP_BUILD_FA ON，因此 ARM 机器可能复用 x86_64 产物，后续链接、打包或运行会失败。

请在启用 ARM 的同时补齐主机架构处理：要么让缓存 key/URL 包含 host arch 并提供对应的 aarch64 缓存包；要么在 WITH_ARM 时强制关闭缓存、走本机源码编译。修复形态例如：

# cmake/external/flashattn.cmake
if(WITH_ARM)
  set(WITH_FA_BUILD_WITH_CACHE OFF)
endif()

@risemeup1111 risemeup1111 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1

处理要求：请针对该评论修复并提交新的 commit。

复查当前 head 469767fb2b22bdb0cd69b08de949feb0e4045baa 后，b0339c695d890caa01ac9f035fb08b1bfe50f597..469767fb2b22bdb0cd69b08de949feb0e4045baa 在相关文件没有差异；cmake/external/flashattn.cmake 仍然使用 cu${FA_NVCC_ARCH_BIN}/flashattn_libs_${FLASHATTN_TAG}.tar 作为缓存路径，并在下载成功后复制 .so、设置 SKIP_BUILD_FA ON。因此原问题没有被修复：ARM 构建仍可能复用未区分 host arch 的缓存产物。

请继续按上面的建议补齐 host arch 缓存隔离，或在 WITH_ARM 下关闭 FlashAttention 缓存并走本机编译。

@liuhao2638 liuhao2638 Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CI中没有arm的编译拦截，所以这一条可以豁免

@risemeup1111 risemeup1111 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1

处理要求：请针对该评论修复并提交新的 commit。

这里不能因为 CI 里没有 ARM 编译拦截就豁免。这个问题不是“CI 是否覆盖 ARM”，而是本 PR 放开 ARM 后，真实的 ARM64 GPU 构建会进入这段 FlashAttention 逻辑；当前缓存 URL 仍不区分 host arch，下载成功时会直接复制缓存 .so 并 SKIP_BUILD_FA ON，因此 ARM 构建/打包存在误用非 ARM 产物的风险。缺少 ARM CI 只能说明这个风险不会被自动发现，不能说明代码路径是安全的。

如果只是豁免 ARM CI 验证，可以在测试计划里说明；但代码层面仍需要二选一处理：为 FlashAttention 缓存增加 host arch 维度并提供对应产物，或在 WITH_ARM 下关闭缓存走本机编译，例如：

if(WITH_ARM)
  set(WITH_FA_BUILD_WITH_CACHE OFF)
endif()

@liuhao2638 liuhao2638 Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

我们永远不会在CI之外使用编译cache

@risemeup1111 risemeup1111 Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1

处理要求：请针对该评论修复并提交新的 commit。

我理解你说的是使用约定，但当前仓库没有把“编译 cache 只在 CI 使用”这个约束写进构建逻辑：CMakeLists.txt 里 WITH_FA_BUILD_WITH_CACHE 默认是 ON，cmake/external/flashattn.cmake 也只判断这个 CMake 变量，没有判断是否处于 CI 环境。因此任意源码构建只要没有显式关掉该选项，WITH_GPU + WITH_ARM 仍会进入下载缓存并复制 .so 的路径。

如果项目约定确实是 cache 只服务 CI，请把这个约束编码到构建系统里，例如把默认值改为 OFF，只在 CI 脚本中显式打开；或者至少在 ARM 下强制关闭：

if(WITH_ARM)
  set(WITH_FA_BUILD_WITH_CACHE OFF)
endif()

这样代码行为才和“CI-only cache”的约定一致，也能避免后续手动/发布 ARM 构建误用缓存产物。

@liuhao2638 liuhao2638 Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. 已在 commit abe4c2ea21 中修复：WITH_ARM 下关闭 FlashAttention 缓存复用，避免 ARM64 GPU 构建下载/复用未区分 host arch 的缓存包，改为走本机源码编译。

@risemeup1111 risemeup1111 Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已确认，abe4c2ea21 中在 WITH_ARM 下关闭了 FlashAttention cache，原先担心的 ARM 复用非 ARM 缓存包问题已消除。

@PaddlePaddle-bot PaddlePaddle-bot Jul 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Bug 这里允许 ARM64 CUDA 构建进入 FlashAttention external 逻辑，但 cmake/external/flashattn.cmake 的缓存包 URL 只按 NVCC_ARCH_BIN 拼接，没有区分宿主 CPU 架构；默认 WITH_FA_BUILD_WITH_CACHE=ON 时，一旦该 URL 命中缓存，就会复制缓存里的 .so 并设置 SKIP_BUILD_FA=ON，ARM64 构建不会从源码产出 aarch64 的 libflashattn.so。

这会让 ARM64 wheel/安装包携带非本机架构的 FlashAttention DSO，后续 dynload 失败或运行时不可用。

建议修复方式：在放开 ARM 前，让 FlashAttention 缓存路径包含宿主架构并提供对应 aarch64 产物；或者在 WITH_ARM/CMAKE_SYSTEM_PROCESSOR=aarch64 时强制关闭缓存并走源码构建，例如给 external 传 -DSKIP_BUILD_FA=OFF/禁用 WITH_FA_BUILD_WITH_CACHE，确保安装目录中的 .so 由当前 ARM64 toolchain 构建。

AND NOT WITH_ARM
AND NOT WIN32
AND NOT APPLE)
if(${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 12.3
Expand Down
Loading