Skip to content

[Optim] SWA related performance optimization (part 1/3)#162

Merged
GuoxiaWang merged 2 commits into
PaddlePaddle:mainfrom
Enigmatisms:swa_optim_p0
Jul 8, 2026
Merged

[Optim] SWA related performance optimization (part 1/3)#162
GuoxiaWang merged 2 commits into
PaddlePaddle:mainfrom
Enigmatisms:swa_optim_p0

Conversation

@Enigmatisms

@Enigmatisms Enigmatisms commented Jul 7, 2026

Copy link
Copy Markdown

单看这个 PR 的话,全是死代码(暂时没有用到)。

本 PR 在 flashmask 下导出了一些轻量级的 CPP/CUDA kernel (几乎没有编译开销的那种),用于快速在 paddle 端提供小算子的替代方案(比如带宽利用率不足的 slicing / assign 等操作)。

在 FA4 被 install 时,本 PR 的 utils 模块会默认被编译 install,以 flashmask.utils.xxx 导出。
目前只提供两个 API:

  • bshd_slice_contiguous_kv 用于 dst = src[:, start:end, :, :].contiguous() 操作加速(5~6x)
  • accum_zero_axis1_kv 快速对 dKV accum 进行有限范围的 strided zero 操作。

不走 python 端的原因:上述 kernel 显然确实可以 cute DSL 化,但我个人认为完全没有必要:上述 kernel 用 CUDA copy engine 可以单cuda 函数调用解决,并且不使用 SM 资源,无必要为了完全 python 化而写一个 DSL kernel。

utils 模块的初衷也不是为了将一大堆 CPP/CUDA 函数塞入,放入的内容只是一些轻量级的 utils 函数(无法被 cute DSL 替代或者替代纯属自找麻烦的一些函数),尽可能控制本部分的编译时间。

这是三个拆分优化的第一个:

@GuoxiaWang
GuoxiaWang merged commit a5856db into PaddlePaddle:main Jul 8, 2026
1 check passed
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.

3 participants