已合并
feat: kvallgather cp thd load balance #3820
m0_50947149创建于 7月26日
feat: kvallgather cp thd load balance #3820
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 m0_50947149 的贡献）
7月26日 创建了 pull request，commit 6531cceb
ascend-robot
7月26日 评论：
7月26日 评论：
7月26日 添加了label：ascend-cla/no
ascend-robot
7月26日 评论：
7月26日 评论：
Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-Ascend/MindSpeed | ✅ Liz_, 郭鹏, wuweiqiang24 (3/2) | ✅ 郭鹏 (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
m0_50947149, thanks for your pull request. All authors of the commits have signed the CLA. 👍
ascend-robot
7月26日 评论：
7月26日 评论：
ascend docs pipeline is running...
此处折叠了101条消息 查看更多
11 天前 添加了label：approved
11 天前 合入了pull request
ascend-robot
11 天前 评论：
11 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
What this PR does / why we need it?
This PR adds load-balanced KVAllGather context parallelism for the THD format.
It uses EOD-aware symmetric sequence partitioning and implements a THD attention path with direct rank-major K/V indexing and a single TND fused-attention call. Metadata caching is added to reduce repeated calculations.
MHA, GQA, and expanded MLA are supported.
https://gitcode.com/Ascend/MindSpeed/issues/154
https://gitcode.com/Ascend/MindSpeed/issues/232
Does this PR introduce any user-facing change?
Yes. Users can now use THD packed sequences with
kvallgather_cp_algo, causal attention, and EOD attention-mask reset without introducing new command-line options.Each packed subsequence length must be divisible by
2 * context_parallel_size.How was this patch tested?
Unit tests cover:
The modified files also pass Ruff, Pylint, and whitespace checks.
Test_results:
