Skip to content

[Build] Keep NVSHMEM available for FlashMask #79240

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 5 commits into from
Jun 5, 2026
Merged

[Build] Keep NVSHMEM available for FlashMask #79240

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8ed0a10
[Build] Keep NVSHMEM available for FlashMask
ShigureNyako Jun 4, 2026
89bbdcf
Merge develop into fix/keep-nvshmem-for-flashmask
ShigureNyako Jun 4, 2026
da334fd
Merge develop into fix/keep-nvshmem-for-flashmask
ShigureNyako Jun 4, 2026
0dc25f4
Merge develop into fix/keep-nvshmem-for-flashmask
ShigureNyako Jun 4, 2026
6d8c566
[Build] Respect explicit NVSHMEM CMake option
ShigureNyako Jun 4, 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
89 changes: 57 additions & 32 deletions cmake/third_party.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -575,33 +575,9 @@ if(WITH_CUSPARSELT)
list(APPEND third_party_deps extern_cusparselt)
endif()

string(FIND "${CUDA_ARCH_BIN}" "90" ARCH_BIN_CONTAINS_90)
if(NOT WITH_GPU
OR NOT WITH_DISTRIBUTE
OR (ARCH_BIN_CONTAINS_90 EQUAL -1))
set(WITH_NVSHMEM OFF)
endif()
# DeepEP GPU retired, force disable NVSHMEM for non-XPU builds
if(NOT WITH_XPU)
set(WITH_NVSHMEM OFF)
endif()
Comment on lines -584 to -587

@SigureMo SigureMo Jun 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

是否可以，在 XPU OR FA 的场景下开启，感觉不用 FA 的场景也不用 nvshmem

@ShigureNyako ShigureNyako Jun 4, 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.

已按这个方向收窄：

  • 先预计算本次构建是否会启用 FlashAttention。
  • 非 XPU 下，WITH_NVSHMEM=ON 只有在 WITH_GPU + WITH_DISTRIBUTE + sm90 + FlashAttention 都满足时才保留；不用 FA 的非 XPU 场景会继续关掉 NVSHMEM。
  • XPU 场景不走这个非 XPU gate。
  • external/nvshmem 仍然在 external/flashattn 前 include，避免 FA overlap 依赖 extern_nvshmem 时顺序出问题。

也确认没有恢复 GPU DeepEP。

if(WITH_SLEEF
AND NOT WITH_ROCM
AND NOT WIN32)
include(cmake/sleef.cmake)
if(TARGET extern_sleef)
list(APPEND third_party_deps extern_sleef)
endif()
endif()
if(WITH_NVSHMEM)
include(external/nvshmem)
list(APPEND third_party_deps extern_nvshmem)
endif()

set(FLASHATTN_ENABLED OFF)
if(WITH_ROCM)
include(external/flashattn)
list(APPEND third_party_deps extern_flashattn)
set(WITH_FLASHATTN ON)
set(FLASHATTN_ENABLED ON)
endif()

if(WITH_GPU
Expand All @@ -618,24 +594,73 @@ if(WITH_GPU
endforeach()
foreach(arch ${NVCC_ARCH_BIN})
if(${arch} GREATER_EQUAL 80)
include(external/flashattn)
list(APPEND third_party_deps extern_flashattn)
set(WITH_FLASHATTN ON)
set(FLASHATTN_ENABLED ON)
break()
endif()
endforeach()
elseif(${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 11.4)
foreach(arch ${NVCC_ARCH_BIN})
if(${arch} GREATER_EQUAL 80)
include(external/flashattn)
list(APPEND third_party_deps extern_flashattn)
set(WITH_FLASHATTN ON)
set(FLASHATTN_ENABLED ON)
break()
endif()
endforeach()
endif()
endif()

set(ARCH_BIN_CONTAINS_90 -1)
foreach(arch ${NVCC_ARCH_BIN})
if(${arch} GREATER_EQUAL 90)
set(ARCH_BIN_CONTAINS_90 0)
break()
endif()
endforeach()

set(WITH_NVSHMEM_USER_DEFINED OFF)
if(DEFINED CACHE{WITH_NVSHMEM} AND WITH_NVSHMEM)
set(WITH_NVSHMEM_USER_DEFINED ON)
endif()

set(NVSHMEM_AVAILABLE ON)
if(NOT WITH_XPU)
if(NOT WITH_GPU
OR NOT WITH_DISTRIBUTE
OR (ARCH_BIN_CONTAINS_90 EQUAL -1)
OR NOT FLASHATTN_ENABLED)
set(NVSHMEM_AVAILABLE OFF)
endif()
endif()

if(WITH_NVSHMEM AND NOT NVSHMEM_AVAILABLE)
if(WITH_NVSHMEM_USER_DEFINED)
message(
FATAL_ERROR
"For non-XPU builds, WITH_NVSHMEM=ON requires WITH_GPU=ON, "
"WITH_DISTRIBUTE=ON, NVCC_ARCH_BIN containing sm90 or higher, "
"and FLASHATTN_ENABLED=ON.")
else()
set(WITH_NVSHMEM OFF)
endif()
endif()
if(WITH_SLEEF
AND NOT WITH_ROCM
AND NOT WIN32)
include(cmake/sleef.cmake)
if(TARGET extern_sleef)
list(APPEND third_party_deps extern_sleef)
endif()
endif()
if(WITH_NVSHMEM)
include(external/nvshmem)
list(APPEND third_party_deps extern_nvshmem)
endif()

if(FLASHATTN_ENABLED)
include(external/flashattn)
list(APPEND third_party_deps extern_flashattn)
set(WITH_FLASHATTN ON)
endif()

if(WITH_CUDNN_FRONTEND)
include(external/cudnn-frontend) # download cudnn-frontend
list(APPEND third_party_deps extern_cudnn_frontend)
Expand Down