提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。

性能优化具体描述

在 NPU IVFPQ 大规模建库场景中，add 阶段耗时较高。典型配置为 dim=128、nlist=262144、M=16、nbits=8，使用 8 张 NPU 进行多卡 add，目标数据规模为 2 亿向量。

分析发现 add 阶段主要瓶颈包括：

1. L2 PQ 编码阶段逐向量执行编码，批量数据场景下开销较高。
2. 大 nlist 场景下 L1 coarse assign 的 DistanceFlatL2MinsAtFp32 算子 tiling 参数不够适合当前输入规模。
3. 多卡 sharded add 时，同一个 IVF list 可能被拆分到多个 device，导致后续产生较多小粒度 list/device 写入和元数据更新。

本优化针对上述问题进行改进：

1. 使用 ProductQuantizer::compute_codes 批量完成 PQ 编码，替代逐向量编码。
2. 将 DistanceFlatL2MinsAtFp32 的 BASE_CODE 从 8192 调整为 4096，改善 dim=128、nlist=262144 下的 L1 assign 性能。
3. sharded add 中同一个 IVF list 在一次 add 内绑定到同一个 device，减少碎片化写入。

性能劣化说明

优化前，在 dim=128、nlist=262144、M=16、nbits=8、8 卡 NPU、单批 200 万向量 add 的测试中，add 阶段耗时较高。

优化前关键耗时：

• L1 coarse assign: 约 30.6s
• L2 PQ encode: 约 86.2s
• add 总耗时: 约 122s
• add 吞吐: 约 1.6 万 vectors/s

第一阶段优化 L2 批量编码后：

• L1 coarse assign: 约 30.6s
• L2 PQ encode: 约 0.5s
• add 总耗时: 约 36s
• add 吞吐: 约 5.5 万 vectors/s

继续优化 L1 distance tiling 后：

• L1 coarse assign: 约 16.4s
• L2 PQ encode: 约 0.4~0.5s
• add 总耗时: 约 21~22s
• add 吞吐: 约 9.0 万 vectors/s

整体看，单批 200 万向量 add 从约 122s 降低到约 22s，性能提升约 5 倍。按 2 亿向量建库估算，add 时间从约 2 小时降低到约 35~40 分钟。

其他相关讨论

环境信息

- 操作系统：Linux
- 昇腾硬件信息：8 卡 NPU，芯片类型 ascend910b
- CANN 软件版本：CANN 9.0.0
- 测试配置：
  - dim=128
  - nlist=262144
  - M=16
  - nbits=8
  - metric=IP
  - add batch=2000000
  - deviceList=[0,1,2,3,4,5,6,7]
  - useNpuTrain=True
  - useDistributedCoarse=True

欢迎加入社区，感谢您对社区的贡献 🎉!