基于Atlas A3集群的LongCat-Flash模型推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍基于Atlas A3系列产品的LongCat-Flash模型的优化方式。性能优化通用优化LongCat-Flash结构中的与Llama类似的部分可参考通用优化点Llama的改动如固定KV Cache大小、cos/sin优化、AddRMSNorm融合等。使能图模式使用静态图可以获得更好的推理性能。LongcatFlashRunner通过覆写executor/model_runner.py中的ModelRunner的graph_compile函数将模型编译为静态图当前暂不支持acl_graph。使能图编译缓存torch.compile是一种即时编译器Just-In-Time compiler成图首次编译时间通常较长在时延敏感的模型推理场景下使能图编译缓存可以缓存编译后的静态图有效缩短服务启动后的首次推理时延从而提高推理性能。可参考LongcatFlashRunner中graph_compile函数中的使用if self.enable_cache_compile: case_name compile_cache/ os.getenv(CASE_NAME) cache_model self.model.decode if self.is_mtp: case_name _spec cache_model self.model.mtp_compile_decode cache_dir os.path.join(os.path.dirname(os.path.abspath(__file__)), case_name) self.model.decode tng.inference.cache_compile(cache_model, cache_dircache_dir, configcompiler_config, dynamicFalse, fullgraphTrue, ge_cacheTrue)主模型缓存默认路径为./compile_cache/CASE_NAMEmtp模型缓存默认路径为./compile_cache/CASE_NAME_spec。多流并行与控核大模型推理场景下对于一些可并行的场景可以划分多个stream做并行计算多个stream上的计算形成overlap从而降低整体计算耗时。多流并行技术的详细介绍请参考官方文档多流场景下会出现所有核Core都被一个流占用的情况导致算子执行并行度降低因此需要把核分给不同的流用从而保证算子并行执行的收益。控核技术的详细介绍请参考官方文档原始的LongCat-Flash模型在论文中提供了四阶段的并行策略其方案图如下所示。我们对并行策略进行了调整调整后的多流并行和控核方案图如下所示。将第二段attention和FFN专家提前执行并通过控制多流上的ai core和vector core核数使得双流的计算时间接近无明显拖尾提升性能。其中stage1不做控核默认占用全部的ai core和vector core核数stage2里的stream0采用c16v32控核方案stream1采用c8v16。实现多流并行和控核可以参考以下伪代码。attn layernorm with npu_stream_switch(True, 1): with limit_core_num(True, 8, 16): router dispatch gmm combine with limit_core_num(True, 16, 32): dense attn layernorm dense权重预取该优化提供网络weight预取功能在算子计算的同时利用空闲的带宽提前将一些访存bound算子的权重从HBM搬运到L2 Cache中提升算子性能。npu_prefetch技术的详细介绍请参考官方文档。npu_prefetch优化功能可通过enable_prefetch开关使能。以下图为例Matmul的权重有20 MB前序的算子RmsNorm对内存带宽的需求较小执行时间为10us此时可以通过预取技术将Matmul的权重提前写入到L2 Cache中理想情况下最大可预取的大小为HBM_brandwidth / time_available实际可预取大小因网络与前序算子对带宽的占用程度而异。假设此时可预取的大小为10 MB与RmsNorm并行执行则在Matmul算子执行时只需要再额外读10 MB权重访存的耗时开销减少明显。判断预取的大小和位置是否合理可以通过profiling分析如果与预取并行的算子性能有明显劣化则需要进一步调整预取的位置可以将预取的位置提前开始或与劣化算子错开如没有可调整的位置则需要减少预取的大小使得并行算子尽量不劣化。如果想要获取最大化的预取收益可以通过实验对比调整。如果前序算子是访存密集型如Matmul等则不适合并行避免带宽抢占导致劣化。预取分析和调试的通用流程如下注各模型的预取情况可能有所不同可按需调整上述流程。下图为LongCat-Flash模型的预取位置。我们针对访存bound的算子如QuantBatchMatmul (QBMM)、MLAProlog、Matmul等算子提前预取了对应的权重。其中MLAProlog算子包含了多个矩阵乘计算搬运bound较大因此我们提前预取了前置的QBMM的权重为MLAProlog提供了更大的预取空间获取较大的性能收益。具体的预取大小及预取位置可在models/modeling_longcat_flash.py中搜索npu_prefetch接口查看。使能SuperKernelSuperKernel优化功能在decode启用ge_graph图模式的场景下根据用户定义的范围对模型的计算图进行优化。SuperKernel技术的详细介绍请参考官方文档。SuperKernel优化功能将通过enable_superkernel开关使能可将部分算子优化在一个SuperKernel scope内从而实现对任务调度的等待时间和调度开销的优化提升整体性能。由于我们在不同流上采取了不同的分核策略按照分核、分流的范围标定各SuperKernel scope的范围即可。下图为针对Longcat-Flash模型标定的SuperKernel范围。MLA (Multi-Head Latent Attention)低秩压缩优化Decode阶段参考Deepseek论文中提及的低秩压缩方法可以减少KV Cache占用的内存提升推理效率相关实现可以参考LongcatFlashAttention类中的forward_page_attention_absorb函数。MLA整体采用Data Parallelism (DP) 数据并行并针对o_proj matmul单独采用Tensor Parallelism (TP)切分提升性能。融合算子优化MLA前置计算性能优化使能[npu_mla_prolog_v3]融合kernel替换attention计算前的计算其中包含Q、K、V的线性层计算、旋转位置编码 (ROPE)、RmsNorm计算及KV Cache更新等计算处理Attention性能优化使能npu_fused_infer_attention_score融合kernel实现对MLA计算的加速。o_proj Tensor Parallelism K轴切分对线性层o_proj进行行切分处理o_proj的输入维度为num_heads * v_head_dim输出维度为hidden_size。在进行o_proj TP时将这一输入维度按照o_proj_tp_size进行切分要求num_heads * v_head_dim能被o_proj_tp_size整除每个rank上只保存一部分输入通道对应的权重。在推理过程中attn_output先经过all_to_all使得每个rank只拿到自己负责的那一段输入特征然后每个rank对局部张量执行线性层变换得到局部输出最后通过reduce_scatter对各rank的局部输出按元素求和并在batch维度切分得到各rank的输出。MLA (Multi-Head Latent Attention) KVP切分优化KVP特性旨在面向长序列推理场景通过将KV Cache沿S维度切分到多个rank上缓解单rank的KV Cache访存压力从而降低时延。KVP特性暂未支持Multi-Token Prediction投机推理。Prefill阶段使能KVP切分优化将根据kvp_size和kvp_rank对slot_mapping进行切分并将slot_mapping作为npu_kv_rmsnorm_rope_cache算子的index参数将完整的KV Cache分片轮转写入到各个kvp_rank上每个kvp_rank上存储1/kvp_size长度的KV Cache此外Prefill阶段在计算Attention时复用了kvp_size同时约束attn_tp_size 1对npu_fused_infer_attention_score算子的输入Q、K、V按头维度进行TP切分。Decode阶段使能KVP切分优化启用npu_fused_infer_attention_score算子的softmax_lse_flag功能输入Q、K、V头维度保持完整与各个kvp_rank上存储的部分KV Cache进行Attention计算并输出相应的softmax lse每个kvp_rank拿到完整num_heads_per_rank个Q head对部分KV Cache的计算结果attn_partial和lse_partial通过_all_to_all_along_headdim从各个kvp_rank拿到kvp_size份num_heads_per_rank // kvp_size个Q head对各个部分KV Cache的计算结果attn_scatter和lse_scatter并通过npu_attention_update算子对attn_scatter和lse_scatter执行归约操作。Decode阶段actual_seq_lengths_kv表示当前kvp_rank实际存储的KV Cache长度因此需要根据kvp_size和kvp_rank进行调整。使能KVP特性的场景下kvp_size 1新增了o_proj_tp_size kvp_size约束在o_proj_forward中配合执行o_proj TP切分并后续执行all_reduce通信。MoE (Mixture of Experts)模块实现Expert Parallel (EP)及使能融合算子MoE计算阶段采用EP (Expert Parallelism)切分策略将路由专家和零计算专家均匀分布到每张卡上。Router使能融合算子使用torch_npu.npu_moe_gating_top_k算子对router计算的结果排序并选取前top-k个专家。Prefill阶段优化Prefill阶段路由专家采用Double-Routing的计算策略完成计算,具体计算步骤可参见基于Atlas A3集群的DeepSeek-R1模型prefill阶段推理性能优化实践的MoE部署策略优化章节。Decode阶段优化高性能专家计算使用torch_npu.npu_grouped_matmul算子可以同时处理多个专家的计算提高计算和搬运效率多卡间高性能通信路由使能torch_npu.npu_moe_distribute_dispatch_v2 和torch_npu.npu_moe_distribute_combine_v2算子实现EP并行下多卡间的通信。通过传入参数copy_expert_num使能dispatch_v2和combine_v2算子支持零专家处理计算公式为MoE(ori_x) ori_x。在使用前可参考上述的算子文档检查HCCL_BUFFSIZE等环境变量的配置是否合理了解该算子的使用场景和约束。MLP线性层优化MLP线性层计算合并原始LongCatFlashMLP实现中存在gate_proj、up_proj与down_proj三个matmul运算可通过将gate_proj与up_proj进行合并计算得到gate_up_proj提升整体计算效率。MLP线性层TP切分在MLP的计算中gate_up_proj和down_proj需要全量存储到每一个device上造成device内存压力本优化将gate_up_proj沿N轴切分、down_proj沿K轴切分到dense_tp域内的不同device上以完成MLP线性层的TP切分降低单个device的内存使用并减少matmul矩阵运算时的权重搬运开销。需要注意的是MLP阶段采用TP切分但前后的Attention模块采用的是DP切分在MLP计算之前和之后需要分别进行AllGather和ReduceScatter完成DP - TP - DP的并行方式转化。虽然有额外的通信开销但整体仍有较好的性能收益。在当前的优化实践中Dense FFN专家采用TP8切分。支持Multi-Token Prediction (MTP)实现了MTP投机推理在未达到计算bound的场景下MTP计算可以实现较好的推理加速效果。可通过next_n参数使能MTP。当前支持了MTP1MTP2。Attention-FFN Disaggregation(AFD)优化优化出发点在非分离场景中Attention 模块和 MoE 模块部署在同一个卡上由于 ScMoE 结构的固有特性为了获取最佳的性能在 多流并行与控核 小节中对 Stage2 设计了多流并行策略通过分配不同的核数将 Attention 计算流和 MoE 计算流进行了并行流水达到最佳性能。然而对 Stage2 通过控核后由于两条计算流在计算时只能使用部分 Cube 和 Vector 核算子执行时会受到算力的约束。在 Atlas A3 环境上实测对比发现在不控核情况下MLAProlog/FA/MM/TBMM/QBMM算子耗时相比控核时均有所下降耗时对比如下表格。Cube_Vector核数MLAPrologusFAusMMusTBMMusQBMM1usQBMM2us12C_24V68.11181819.618.720.124C_48V45.291.814.212.514.616.3以上数据除了核数不同其他都在相同的配置情况下从整网中采集拆解得到的。根据以上分析针对 LongCat-Flash-560B 模型为了在 Decode 阶段进一步降低 TPOT 耗时可采用 Attention-FFN Disaggretation(AFD) 技术方案它将 MoE 模块从整网中剥离出来进行独立部署也即 Attention 模块 和 MoE 模块单独部署在不同的节点上中间通过 Send/Recv 算子进行节点间的数据交互使能 AFD 技术前后的网络结构示意图如下。Attention 和 MoE 独立部署后的示意图如下计算流图使能 AFD 之后 Send/Recv 算子跨节点通信走 SDMA 不占用 Vector 核Attention 模块中计算算子和 Send/Recv 并行时不再受控核的约束。计算流上如下图所示 Attention 节点上第一个 LayerNorm(LN) 算子的结果一方面通过 Send 算子发送到 FFN 节点上另外一方面同时送给下一个算子完成后续的计算操作。在 Send 算子完成发送动作后会同时通过 Recv 算子等待并接收从 FFN 节点发送回来的数据。Attention 节点上的 Send/Recv 算子在 Stream1 上和主流 Stream0 上的计算算子进行 overlap达到通信计算隐藏的目的。同理FFN 节点作为服务端任务执行时图上第一个是 Recv 算子它一直等待直到 Attention 节点通过 Send 算子发数据过来。当前 FFN 节点接收到数据后会进行 Router/Dispatch/MoE-GMM/Combine 等操作最后再通过 Send 算子将数据发回 Attention 节点。FFN 权重预取非分离场景中对 QuantBatchMatmul (QBMM)、MLAProlog、Matmul 进行权重预取后一方面由于 Stage2 里 Stream0 上算子的计算耗时与Stream1 上 Router/Dispatch/MoE GMM/Combine 算子耗时相当另一方面也没有较好的时间窗口对 MoE GMM 进行权重预取故对 MoE GMM 算子没有采取权重预取。但使能 AFD 之后Attention 侧再不控核后MLAProlog/FA/MM/TBMM/QBMM 算子耗时降低可能出现 FFN 侧计算瓶颈导致 Attention 侧的 Recv 算子长时间没能收到 FFN 侧发来的数据在计算流上出现 EVENT_WAIT 间隙的拖尾现象如下图所示。针对 FFN 侧计算瓶颈问题可对 GMM/Matmul 访存 Bound 类算子在通信间隙时提前预取对应的权重从而降低 FFN 侧的整体耗时示意图如下。AFD场景下Attention 侧的权重预取保持和非分离场景时一样可参看权重预取小节。Benchmark详细的 benchmark 性能数据请参考longcat-flash模型README。附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考