CANN/cannbot-skills:LongCat-Flash多流与控核联动案例
案例LongCat-Flash 多流与控核联动【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills概述这个案例解决的是 LongCat-Flash Stage2 中双路径并行时“有 overlap 但仍有明显拖尾”的问题。做法是把 Attention 路径和 MoE/FFN 路径放到不同流中同时用limit_core_num给不同流分配不同核数让两条流时长更接近最适合 decode 阶段的低时延优化。背景与问题单纯多流并行并不一定稳定获益因为一个流可能吃掉过多 Core另一个流即使被切出去也跑不起来。LongCat-Flash 的特点在于双路径都很重如果不控核就容易出现一个流拖尾、另一个流提前结束整体关键路径仍然长。因此这个案例的关键不是“多流”本身而是“多流之后还要分核”。核心思路主流跑 dense / attention 主路径。副流提前执行 shortcut MoE 路径。通过limit_core_num(True, aic, aiv)给两个流分配不同的 AI Core / Vector Core 数量。必要时叠加npu_prefetch和 superkernel让双流窗口更完整。执行编排图关键代码第一段代码展示 decode 阶段进入多流分支if (self.enable_multi_stream 0) and not is_prefill: return self.multi_stream_forward( hidden_states, kv_len, actual_seq_lengths_kv, position_embeddingsposition_embeddings, attention_maskattention_mask, position_idsposition_ids, past_key_valuepast_key_value, is_prefillis_prefill, slot_mappingslot_mapping, past_residualpast_residual, cur_topk_listcur_topk_list, next_layernext_layer, )第二段代码是核心并行段把 shortcut MoE 放到副流并单独控核if not self.enable_afd: npu_prefetch(self.enable_prefetch, self.mlp.router.classifier.weight.data, o_proj, 18 * 1024 * 1024, 0) with npu_stream_switch(True, 1): with limit_core_num(True, self.aic_num1, self.aiv_num1): shortcut_mlp_output self.mlp(hidden_states_norm, is_prefill, cur_topk_listcur_topk_list)主路径继续在另一套核数配置下运行with limit_core_num(not self.enable_afd, self.aic_num2, self.aiv_num2): hidden_states, _, dsq self.mlps[0](https://link.gitcode.com/i/051379f4dc2d15cffd2bacf3b29b099c) hidden_states, residual self.input_layernorm[1](https://link.gitcode.com/i/01e8779bd8cf45ac2de04cc683af5413) hidden_states, o_proj self.self_attn[1].fused_infer_attention_score( query_statesquery_states, k_nopek_nope, k_ropek_rope, attention_maskattention_mask, actual_seq_lengths_kvactual_seq_lengths_kv, )复用参考代表实现LongCat-Flash。相似实现别的模型可以复用“多流后再分核”的思想但核数配置要重调。特化实现AFD 场景会把通信 overlap 也纳入设计。注意事项多流但不控核收益可能会被拖尾抵消。控核参数不是通用值要结合 profile 调整。prefetch、superkernel、控核、多流四者在这个案例里是耦合设计拆开后收益可能不成立。关键词npu_stream_switchlimit_core_numnpu_prefetchdecodeshortcut_mlp_output【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考