CANN/cannbot-skills Indexer Prolog多流并行案例

案例Indexer Prolog 多流并行【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills概述这个案例解决的是 Lightning Indexer 前处理链路串行过长的问题。做法是把 Q 路径和权重投影路径拆到不同 tagged stream 中执行让前处理阶段出现 overlap最适合 Attention 前处理或 Prolog 类算子链的时延优化。背景与问题Indexer Prolog 往往由多个前处理算子组成例如线性投影、RoPE、量化和权重路径准备。如果全部堆在主流上前处理会形成一个很长的串行段后续 attention 计算即使再快也要等前面的准备工作全部完成。这类场景适合多流的原因在于Q 路径和部分权重路径之间可形成局部并行窗口。前处理中既有 Cube 类算子也有 Vector / quant 路径容易形成硬件互补。只要同步点设计得当就不会改变后续 attention 的输入语义。核心思路使用一条副流先跑q_b和 Q 相关预处理。再用另一条流提前跑weights_proj。在进入后续 quant 或 attention 前通过wait_tensor或 tagged event 保证依赖满足。这种写法本质上是“前处理切流”而不是完整的双网络并行。执行编排图关键代码第一段代码展示 Q 路径被放到22号流里enable_multi_streams self.enable_multi_streams and not is_prefill with npu_stream_switch(enable_multi_streams, 22): if enable_multi_streams: tng.scope.npu_wait_tensor(qr, query_states[0]) q_b self.wq_b(qr, c8_input_dict.get(pertoken_scale, None)) q q_b.view(bsz, seqlen, self.n_heads, self.head_dim) q_pe, q_nope torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim-1) q_pe torch_npu.npu_rotary_mul(q_pe.view(-1, self.n_heads, 1, self.rope_head_dim), cos, sin) q torch.cat([q_pe.view(bsz, -1, self.n_heads, self.rope_head_dim), q_nope], dim-1)第二段代码展示weights_proj提前在33号流里执行with npu_stream_switch(enable_multi_streams, 33): if enable_multi_streams: tng.scope.npu_wait_tensor(x, q_b) weights self.weights_proj(x.view(-1, self.dim))如果图模式开启常会配合 tagged stream eventif enable_multi_streams and self.enable_aclgraph: tng.ops.npu_record_tagged_stream(qr, 22) tng.ops.npu_tagged_event_record(indexer_npu_events[0])复用参考代表实现DeepSeek-V3.2-Exp。相似实现GLM-5。特化实现和 MoE 双流不同这类案例通常不直接并行完整模块而是拆分前处理子链。注意事项前处理切流后同步点设计不清楚时最容易出现输入未准备好的问题。如果q_b、weights_proj之间实际共享更多隐藏依赖盲目切流会导致图编译失败或精度问题。图模式下要统一 stream tag 的编号和事件生命周期。关键词npu_stream_switchIndexer Prologweights_projq_bwait_tensortagged stream【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考