CANN Qwen Dense推理优化

基于Atlas A2/A3的Qwen Dense模型推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍Qwen2/Qwen3 Dense非MoE模型基于NPU的推理优化策略。基于Atlas A2/A3 训练/推理系列产品实现BF16场景下的推理。Tensor Parallel (TP)优化Attention TP优化切分策略对Attention的张量切分策略可以分为对QKV头的切分和对线性层的切分。在对QKV头切分时attention的多头计算机制可以方便进行张量切分每个头先独立计算再将结果concat起来。假设模型的attention层需要对num_heads个query按照切分数量attn_tp_size进行切分要求num_heads必须能被attn_tp_size整除每张卡放置query头个数为num_heads_per_rank num_heads // attn_tp_sizekey和value头数相等且可能小于等于query头个数在GQA场景下会小于。为了确保每张卡至少放置一个key和value头每张卡放置的key或value头数计算方法为num_key_value_heads_per_rank max(num_key_value_heads // attn_tp_size, 1)。在对线性层o_proj进行切分时按照行切分即可。计算分解该优化策略先将Q、K、V的线性层计算合并为一次Matmul计算merged_qkv_proj提升计算性能。将merged_qkv_proj的输出结果按Q、K、V拆分后对Q和K进行归一化操作Qwen3启用QK-NormQwen2无此步骤并使用旋转位置编码再计算attention最后通过o_proj层输出。MLP TP优化切分策略对MLP层的gate_proj与up_proj进行列切分对down_proj进行行切分。同时对gate_proj与up_proj线性层采用合并计算的优化方式得到gate_up_proj。计算分解MLP层存在gate_proj、up_proj与down_proj三个matmul运算具体运算为 x down( SiLU(gate(x)) * up(x) )。本优化将张量切分后的gate_proj和up_proj合并为gate_up_proj一次Matmul计算将输出按中间维度切分为两块后分别执行SiLU激活和element-wise乘法最终通过down_proj输出。通过将gate_proj与up_proj合并计算减少kernel launch次数提升整体计算效率。残差融合优化在原始transformers实现中每个DecoderLayer的残差连接和RMSNorm是分开执行的residual hidden_states hidden_states RMSNorm(hidden_states) hidden_states Attention(hidden_states) hidden_states residual hidden_states ← 独立residual add本优化通过使能torch_npu.npu_add_rms_norm融合算子将残差加法与RMSNorm合并为一次kernel调用hidden_states, residual npu_add_rms_norm(residual, hidden_states, weight, eps)每层DecoderLayer包含两处残差连接Attention前和MLP前融合后每层减少两次独立的tensor加法算子及对应的显存读写开销。残差residual作为参数在层间传递不再在每层内部做独立的add操作。使能融合算子RmsNorm算子优化通过使能torch_npu.npu_rms_norm算子替代原始手动计算variance rsqrt的实现提升模型推理性能。RoPE算子优化通过使能torch_npu.npu_apply_rotary_pos_emb融合算子替代原始rotate_half 手动乘加的实现。同时对cos/sin进行预计算并缓存避免每次forward动态计算。FlashAttention融合算子优化通过使能torch.ops.npu.npu_fused_infer_attention_score_v2推理场景下的FlashAttention融合算子既可以支持全量计算场景也可支持增量计算场景。使能图模式使用静态图可以获得更好的推理性能。通过覆写executor/model_runner.py中的ModelRunner的graph_compile函数将模型编译为静态图。附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考