[Optim] SWA related performance optimization (part 2/3)#161
Merged
Conversation
Enigmatisms
changed the title
Enigmatisms
changed the title
This was referenced Jul 7, 2026
GuoxiaWang
approved these changes
Jul 8, 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
2 participants
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.
提升 SWA 层性能。本 PR 的主要操作是：
考虑到 SWA 在 dKV postprocess CP 情况下只需要对一小段进行 post-process（本地段以及本地段前面的 window size 个 token），故新增两个可选参数（start end range），专门用于 contiguous 切分情况下的 dKV postprocess。根据 start end range，进入 post-process 前会对 tensor 进行 seqlen 维度 slice，以降低实际 dKV postprocess 访存量（纯 mem bound 操作）。
本 PR 的功能，在不传入对应 range 时，不影响精度行为。
这是三个拆分优化的第二个：