[Build] Compile FlashAttention on arm64 platforms

[liuhao2638]

PR Category

Environment Adaptation

PR Types

Improvements

Description

1. Remove the NOT WITH_ARM guard from GPU third-party dependency conditions.
2. Allow FlashAttention-related GPU third-party dependency logic to be enabled on ARM64 GPU builds.

Tests

• git diff --check -- cmake/third_party.cmake
• Planned: ARM64 + GPU CMake/build validation with WITH_ARM=ON and WITH_GPU=ON to verify FlashAttention-related third-party dependency logic is enabled on ARM64 GPU builds.

是否引起精度变化

否。本变更仅调整 ARM64 GPU 构建依赖开关，不修改数值计算逻辑或算子实现。

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[CLAassistant]

All committers have signed the CLA.

[risemeup1111]

代码层面有一个阻塞风险已放在行级评论里，请优先处理后再推进。另有一个非行级的 PR/CI hygiene 问题：

• 优先级：P3 非行级：PR 模板和 CI 检查不对应具体 diff 行。
  处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。
  当前 Check PR Template 失败，日志显示 PR Category 不在允许列表内，并且缺少“是否引起精度变化”；PR 描述里的 Tests 也还是空的。请按最新模板更新分类（本变更可考虑 Environment Adaptation，以项目规范为准）、补充精度影响说明，并填写已执行或计划执行的 ARM+GPU CMake/build 验证。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

优先级：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]

优先级：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]

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

[risemeup1111]

优先级：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]

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

[risemeup1111]

优先级：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]

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

[risemeup1111]

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

[risemeup1111]

复查当前提交后，上一轮 P1 仍未解决；我已在原行级线程补充当前 head 的核验证据，仍需修复 commit。

• 优先级：P3 非行级：PR 模板检查仍失败，不对应具体 diff 行。
  处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。
  PR 描述已补充 Category 和测试项，但当前 CI 日志仍显示 必须填写是否引起精度变化。请按模板补上“是否引起精度变化”字段，并说明本次 CMake 条件调整对数值精度无影响或具体影响。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

已复查本次编辑，PR 描述里的“是否引起精度变化”已补上；原来的阻塞点仍在 inline 评论里。另有一条非行级建议：

• 优先级：P3 非行级：标题仍使用 [Build]，与当前 Environment Adaptation 定位不一致。
  处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。
  建议把标题改成类似 [Environment Adaptation] Enable FlashAttention GPU dependencies on ARM64 builds，这样和模板分类、变更内容更一致。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[liuhao2638]

done

[liuhao2638]

Done. 已按 P3 评论更新 PR 模板：分类使用 Environment Adaptation，补充精度影响说明，并在 Tests 中填写 ARM64 + GPU CMake/build 计划验证。

[SigureMo]

@risemeup1111 不许再参考垃圾 @PaddlePaddle-bot 提出没有任何没有意义的建议！

[risemeup1111]

已复查最新提交，之前 inline 线程里的 ARM 下 FlashAttention cache 问题已修复，PR 模板信息也已补齐。详细核对见之前的 inline 线程，本轮没有新增需要阻塞合入的问题。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

✅ Cherry-pick successful! Created PR: #79496