案例Prefill Micro-Batch 双流流水【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills概述这个案例解决的是 prefill 阶段“计算和通信串行通信空洞明显”的问题。做法是把输入拆成两个 micro-batch让两条流分别承载不同 micro-batch并通过record_event/wait_event把 dispatch、expert、combine 等阶段流水化最适合计算和通信都比较重的 prefill 场景。背景与问题prefill 常常是计算 bound但在大模型的 EP 场景里MoE 的 dispatch 和 combine 也会带来显著通信开销。如果这些阶段完全串行芯片在等待通信或等待计算时都会出现空洞造成端到端吞吐下降。把 prefill 切成两个 micro-batch 后虽然单次 shape 会变化但如果计算和通信的线性度足够好就能用一条流做当前 micro-batch 的计算用另一条流穿插上一个或下一个 micro-batch 的通信从而实现通信隐藏。核心思路创建主流cur_stream和副流stream1。micro-batch 0 和 micro-batch 1 在不同流上生成输入。Attention、gate、dispatch、expert、combine、finalize 不再线性执行而是借助 event 形成流水。这种案例的重点不是“简单双流”而是“多阶段事件编排”。执行编排图关键代码第一段代码是双流入口最核心的是先创建副流self.micro_batch_mode MicroBatchMode(...) if self.micro_batch_mode ! MicroBatchMode.DISABLE: self.stream1 torch.npu.Stream()第二段代码展示最基础的两条流分别处理两个 micro-batchcur_stream torch.npu.current_stream() with torch.npu.stream(cur_stream): hidden_states_mb0, residual_mb0, _, cos_sin_mb0, slot_mapping_mb0, actual_seq_lengths_kv_mb0 \ self.prepare_inputs_for_layer(input_ids_mb0, kv_len_mb0, position_ids, actual_seq_lengths_kv_mb0, is_prefill) with torch.npu.stream(self.stream1): hidden_states_mb1, residual_mb1, _, cos_sin_mb1, slot_mapping_mb1, actual_seq_lengths_kv_mb1 \ self.prepare_inputs_for_layer(input_ids_mb1, kv_len_mb1, position_ids, actual_seq_lengths_kv_mb1, is_prefill)第三段代码展示事件驱动的 dispatch/expert/combine 流水event_routing_dispatch_mb0 cur_stream.record_event() with torch.npu.stream(self.stream1): self.stream1.wait_event(event_routing_dispatch_mb0) tokens_per_expert_group_mb0, gathered_tokens_mb0, gathered_pertoken_scale_mb0, input_splits_mb0, \ output_splits_mb0 decode_layer.forward_dispatch_double_routing( tokens_per_expert_mb0, expanded_x_mb0, pertoken_scale_mb0 ) event_dispatch_expert_mb0 self.stream1.record_event() cur_stream.wait_event(event_dispatch_expert_mb0) new_x_mb0 decode_layer.forward_expert( gathered_tokens_mb0, tokens_per_expert_group_mb0, gathered_pertoken_scale_mb0 ) event_expert_combine_mb0 cur_stream.record_event()复用参考代表实现DeepSeek-R1 prefill。相似实现其他 EP micro-batch 场景可参考其事件编排方式。特化实现SP-TP-EP 与 DP-EP 可能采用不同流水版本。注意事项shape 切半后如果算子性能显著劣化双流可能收益变差。event 顺序稍有错误就可能造成死等、早读或精度异常。双流只是手段真正关键是 dispatch、expert、combine 的排列顺序。关键词torch.npu.Streamrecord_eventwait_eventmicro-batchdispatchcombine【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考