{{ message }}
Repository navigation
[Optim] FlashMask-4 Overlap support part 1 (NVSHMEM backend) - #166
Merged
Merged
Conversation
- RS & AG multi-stage bwd overlap - AG fwd overlap - FM-3 overlap bridge
baoqiwen
approved these changes
Jul 23, 2026
umiswing
reviewed
Jul 24, 2026
| block_bos is None, | ||
| ) + ( | ||
| # SRBuffer K/V require a distinct artifact only when overlap is active. | ||
| (overlap_bhsd_layout, overlap_kv_chunk_size) if enable_overlap else () |
Member
There was a problem hiding this comment.
overlap_view.shape 不需要加到compile key里吗
umiswing
approved these changes
Jul 24, 2026
GuoxiaWang
approved these changes
Jul 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
FM-4 Overlap 支持第一阶段，将 FM-3 已经大量测试的 overlap 组件通过轻量的 bridge 方式引入 FM-4。
FM-4 端的主要改动在 interface.py，尤其是 bwd，需要适配 multi-stage AG & RS overlap 的 orchestration。
part 1 只支持 NVSHMEM 作为 backend，故在后续封装 Fleet 端 PyLayer 时会限制 EP group 不跨机下才能使用 Overlap。EP 跨机会导致 NVSHMEM 通信单例冲突，出现 hang 或者 core dump。此前已经论证了 NVSHMEM 共用没有简单而不 hacky 的解法。
本 PR 支持并实现了：
只有在 FM-4 安装时带 ovl（或者 all）情况下，overlap 组件才会被编译。由于 overlap 组件强依赖手写的 CUDA kernel，编译肯定是没办法避免的，overlap 组件编译时间并不短（通常2~3min）。但相比 FM-3 还是好很多。
另外，interface.py 做了这样的一个接口小改动：将所有的 non keyword arguments 转为了 keyword args 的写法。考虑到 FM-4 的 API 形参列表可能会发生变动，一旦在列表非尾部位置增加了一个参数，原始写法如果不注意将很容易导致大量参数错位。