Repository navigation
Conversation
weich97
requested review from
liuchongming74,
suteng,
yangzhenzhang02 and
yaoyifanx
as code owners
October 11, 2026 03:53
Author
|
/check-pr |
Author
|
/retest |
|
|
@weich97, your gate pipeline failed at the code check stage. You can configure some code check tools locally following the (pre-push quick guide), to enable automatic checking when you push your code every time. Tools included are: Check_CMakelint, Check_Codespell, Check_Cpplint, Check_Lizard, Check_Pylint, Check_ClangFormat, Check_ShellCheck. |
…ndex tensors The candidate chain and the Indexer KL loss gather compressed keys at selected positions by indexing with two index tensors, a batch id and a position. That form of advanced indexing has no AI Core implementation, so it runs on AI CPU and then accounts for most of the time in both paths. A profiling breakdown of `compressed_candidate_topk` on the torch path (Ascend 910, 16 index heads, head_dim 128, bf16, sequence 4096, compressed length 2048) put 97.8% of device time on AI CPU, and named `Index` as 85.5% of it: 383 ms of a 449 ms step, with the position sort second at 50 ms. Folding the batch offset into the position index turns the gather into one `index_select` over the flattened key, which runs on device. Measured on the same card, with the gathered keys identical bit for bit: | gather shape | two tensors | index_select | ratio | |----------------------------------|-------------|--------------|--------| | 1024 rows x 1024 positions | 101.04 ms | 0.719 ms | 140.6x | | 1024 rows x 512 positions | 51.52 ms | 0.351 ms | 146.9x | | 1024 rows x 512, comp_len 16384 | 59.28 ms | 0.350 ms | 169.2x | End to end on the chain, with the selected positions identical bit for bit: `compressed_candidate_topk` at sequence 4096 and compressed length 2048 goes from 406.22 ms to 55.91 ms, a 350.30 ms saving per call. The first shape is what the candidate chain gathers per query chunk (`topk_blocks * block_size` positions); the second is what the KL loss gathers twice per chunk (`sparse_count` positions), once for the index key and once for the attention key. Both paths run per Indexer source layer per step, and the KL loss runs whenever the post-training Indexer stage is exercised. The gather is extracted into one helper so the three call sites stay consistent and the reason for the form is stated once. Unit tests cover equality with the advanced indexing for one batch and for several, that batches do not mix, the result shape and dtype, a non-contiguous key, repeated positions, and the unchanged output of the candidate chain at batch 1 and batch 3 (CPU only).
weich97
force-pushed
the
perf/gather-index-select
branch
from
October 11, 2026 04:05
c866d7c to
ce675d3
Compare
Author
|
/check-pr |
Author
|
/retest |
Author
|
/retest |
|
|
|
@weich97, your gate pipeline failed at the code check stage. You can configure some code check tools locally following the (pre-push quick guide), to enable automatic checking when you push your code every time. Tools included are: Check_CMakelint, Check_Codespell, Check_Cpplint, Check_Lizard, Check_Pylint, Check_ClangFormat, Check_ShellCheck. |
mindspore-ai-bot
added
ci-pipeline-passed-but-overtime
and removed
ci-pipeline-running
labels
Author
|
/retest |
mindspore-ai-bot
added
ci-pipeline-running
and removed
ci-pipeline-passed-but-overtime
labels
|
|
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.
问题
candidate 链与 Indexer KL loss 都要按选中的位置取压缩 key，取法是用两个索引张量
做高级索引（一个 batch id、一个 position）。这种形式没有 AI Core 实现，直接回退
AI CPU，并随即吃掉这两条路径的绝大部分时间。
三处（均在
shared_compressed_dsa_attention.py，master 上存在）：compressed_candidate_topk→key[batch_indices, safe_positions]_SharedCompressedIndexerKLLoss.forward→index_key[batch_indices, safe_indices]compressed_key[batch_indices, safe_indices]定位
Ascend 910 单卡，16 个 index 头、head_dim 128、bf16、seq 4096、comp_len 2048，
对
compressed_candidate_topk的 torch 路径做 profiling 拆解：IndexSortMaskedFill该链 Stage 448.6 ms/步，与直接计时的 454.5 ms 吻合。
Index是大头，不是矩阵乘。（
einsum此前被怀疑过，经单独画像洗清：孤立仅 2.1 ms。）修法
把 batch 偏移折进位置索引，使之成为对展平后 key 的单次
index_select，在设备上执行：实测（同卡，逐位相同）
孤立 gather：
index_select整条链（选中位置逐位相同）：
index_selectcompressed_candidate_topkseq 4096 / comp 2048第一种形状是 candidate 链每个 query chunk 取的量（
topk_blocks × block_size个位置）；第二种是 KL loss 每个 chunk 取两次的量（
sparse_count个位置，index key 与attention key 各一次）。两条路径都是每个 Indexer source 层每步各跑一遍。
对照（同进程内切回高级索引做 A/B，不是跨跑比较）。
Indexer KL loss：按默认 recipe 的真实几何端到端实测
发布配置
head_dim=512、index_head_dim=128、index_topk=512、num_attention_heads=64、index_n_heads=32；默认验证 recipe 的text_parameter_divisor=8只缩放头数、不缩放头维，所以实跑形状是4 个 128 维 index 头 + 8 个 512 维 attention 头。该 recipe 下
index_n_heads=4使融合算子门控恒假，这条 torch 路径就是唯一会跑的路径。
index_select即每个 Indexer source 层每步约 1.0 s → 约 32 ms，loss 值逐位相等。
该配置有 8 个 Indexer source 层。
机理核对：每个 chunk 取
index_key134 MB +compressed_key537 MB（后者最后一维是 512），seq 4096 / chunk 1024 共 4 个 chunk ⇒ 每层每步搬约 2.7 GB。
按孤立测得的 AI_CPU 约 2.6 GB/s，2.7 GB 约需 1.0 s，与实测 1011.91 ms 吻合；
index_select侧约 383 GB/s，余下 32 ms 为 einsum 与 softmax。量级自洽。适用范围
本修改在 torch 参考路径上，不依赖任何融合算子、不涉及头数门控，因此对
默认验证 recipe 直接生效（该 recipe 下
text_parameter_divisor=8使index_n_heads=4，融合算子要求头数 ∈ (16,24,32,48,64)，融合路径本就不生效）。与 #1093（位置排序走 float32）互相独立、可分别合入：两者改的是不同算子、不同行；
按本次画像，
Index占该链 85.5%、Sort占 11.2%，合入两者后该链实测454.50 ms → 55.91 ms（8.1×）。
测试
tests/ut/components/modules/test_compressed_key_gather.py，8 条，CPU only：与高级索引在 batch=1 与 batch>1 下相等、批之间不串（最容易写错的就是折叠后的
批偏移）、结果 shape 与 dtype、非连续 key、重复位置、以及 candidate 链在
batch 1 与 batch 3 下输出不变。
pytest tests/ut/components tests/ut/core/context_parallel tests/ut/auto_models289 条 + 165 子测试全过；产品文件 pylint 零违规、lizard 无超限。
关键结果一览
index_selectcompressed_candidate_topkseq 4096两条每层每步都要跑一遍的路径（ms，柱=两索引张量高级索引，线=
index_select）xychart-beta title "每个 Indexer source 层每步的耗时（ms）" x-axis ["KL loss 4096/2048", "KL loss 4096/4096", "candidate 链 4096/2048"] y-axis "毫秒" 0 --> 1100 bar [1011.91, 1065.43, 406.22] line [31.91, 31.92, 55.91]该配置有 8 个 Indexer source 层，两条路径都是每层每步各跑一遍。
KL loss 侧 loss 值逐位相等、candidate 链选中位置逐位相同。
机理核对（为什么是 1.0 s 而不是别的量级）：每个 chunk 搬
index_key134 MB +compressed_key537 MB，seq 4096 / chunk 1024 共 4 chunk ⇒ 每层每步约 2.7 GB；按孤立测得的 AI_CPU 约 2.6 GB/s，2.7 GB 约需 1.0 s，与实测 1011.91 ms 吻合。
index_select侧约 383 GB/s，余下 32 ms 为 einsum 与 softmax。量级自洽。与 #1093 合入后该链实测 454.50 ms → 55.91 ms（8.1×）。