Skip to content

[Feat] CP-balance formal incorporation as flash_mask sub-module via build_ext#128

Open
Enigmatisms wants to merge 3 commits intoPaddlePaddle:mainfrom
Enigmatisms:cp_balance_unified
Open

[Feat] CP-balance formal incorporation as flash_mask sub-module via build_ext#128
Enigmatisms wants to merge 3 commits intoPaddlePaddle:mainfrom
Enigmatisms:cp_balance_unified

Conversation

@Enigmatisms
Copy link
Copy Markdown

#127 的优点和缺点都十分明显:

  • 优点在于(1)安装管理很简单:子模块自己内部管理 setup.py 就行。(2)可以进行编译环境隔离:Paddle 目前确认自定义算子无法简单地使用单独一个顶层模块管理若干个子 Extension 并且使他们的编译选项不同。这导致:
    • 如果要统一管理,一般来说必须将 cp_balance 的 source code 让 FA3 Extension 可见,则 cp_balance 将与 FA3 module 共用编译 flags。这将导致 cp_balance 被限制在 SM90 下编译(FA3 flags 带 sm_90 arch flag,没有别的 arch flag)。也即:一个 Extension 无法允许多种 nvcc_flags。 编译灵活性受到严重的限制。
  • 缺点在于:[Feat] CP-balance formal incorporation as flash_mask sub-module #127 install 之后,在环境中将同时出现 flash_mask 库以及 flashmask_cpbalance_ops 库,两个独立的库。如果要进行包管理则比较不优雅:卸载 flash_mask,cp_balance 功能还在。

本 PR 探索了一种统一管理的方法,理论上这个方法能扩展到任意多的子模块。基本方法是:用 build_ext,对 cp_balance 只编译,不安装。生成 .so 文件后,将本地的 python wrapper + .so 文件 copy 到 flash_mask 库安装目录下,由 flash_mask 统一安装。build_ext 对 .so 的裸符号会自动生成适用于paddle动静态图的 wrapper code。不过总的来说,本 PR 的 code 改动相对 #127 还是抽象了一些,写法更 hack,但也能达到正确的 pip wheel 编 whl 以及 pip install -e .,这两种安装方式均已经测试通过。

PR 中的:_build_cuda_submodule 就是用于扩展不同子模块的代码。新的子模块只需要在 flash_mask 中 setup.py 调用此函数。

@baoqiwen
Copy link
Copy Markdown

LGTM

@Enigmatisms Enigmatisms changed the title Cp balance unified [Feat] CP-balance formal incorporation as flash_mask sub-module via build_ext Apr 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants