基于Atlas A2的GPT-OSS模型推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍GPT-OSS系列模型基于NPU的推理的优化方式。基于Atlas A2 系列产品GPT-OSS-120B模型可以采用8卡部署GPT-OSS-20B模型可以在单device上进行部署。性能优化Attention Tensor Parallel (TP)优化对Attention的张量切分策略可以分为对QKV头的切分和对线性层的切分。 在对QKV头切分时attention的多头计算机制可以方便进行张量切分每个头先独立计算再将结果concat起来。假设模型的attention层需要对num_heads个query按照切分数量tp_size进行切分要求num_heads必须能被tp_size整除每张卡放置query头个数为num_heads_per_rank num_heads // tp_sizekey和value头数相等且可能小于等于query头个数在MQA和GQA的场景下会小于。为了确保每张卡至少放置一个key和value头每张卡放置的key或value头数计算方法为num_key_value_heads_per_rank max(num_key_value_heads // tp_size, 1)。QKV头在卡上排布情况如下图所示。针对q_proj、k_proj和v_proj三个线性层可以进一步合并为一个线性层qkv_proj即在推理中将三个矩阵乘替换为一个矩阵乘可以最大化使用NPU的计算能力提升性能。随后将结果切分得到Q、K、V并进行attention计算最后通过o_proj层输出。MoE Tensor Parallel优化假设模型的MoE层的切分数量为tp_size专家个数为expert_num。对MoE层进行张量切分时对gate_proj与up_proj进行列切分对down_proj进行行切分。对gate_proj和up_proj可以采取合并matmul的优化方式得到gate_up_proj实现计算性能优化。固定KV Cache大小首先申请一块固定大小的KV Cache tensor以优化源码动态申请内存的性能损失。KV Cache的大小为(batch * seq_len * num_key_value_heads_per_rank * head_dim)其中seq_len至少为输入和输出的最大总和以保证在推理过程中KV Cache的更新不会越界。具体实现体现在GptOssForCausalLM类的init_cache函数中。在计算完Q、K、V后可以通过scatter_update算子对指定位置上的KV Cache进行更新可参考torch_npu.scatter_update接口。使能融合算子RMSNorm融合与AddRMSNorm融合替换GptOssRMSNorm的forward函数使用融合算子接口torch_npu.npu_rms_norm和torch_npu.npu_add_rms_norm将多个算子的计算替换成融合kernel提升性能。GMM使能与Routing优化在GPT-OSS的MoE专家模块中通过使能融合算子进一步提升性能可以参考GptOssMLP类中的实现。Flash Attention融合算子优化通过使能torch_npu.npu_fused_infer_attention_score_v2推理场景下支持图模式的FlashAttention算子既可以支持全量计算场景,也可支持增量计算场景。快速选择专家在计算专家和token之间的路由分数时可以使用torch_npu.npu_moe_gating_top_k_softmax融合算子代替原来先topk再softmax的多算子操作可以更快速地计算出token和专家的分数。高效排序和token路由使能torch_npu.npu_moe_init_routing融合算子实现MoE routing计算获取专家的排序使能torch_npu.npu_moe_compute_expert_tokens融合算子获取每个专家需要计算的token数使能torch_npu.npu_moe_finalize_routing融合算子将专家计算完成后的token重新排布并加权求和获得最终输出。高性能专家计算使能torch_npu.npu_grouped_matmul融合算子实现多个专家的矩阵乘计算提高计算和搬运效率。集合通信使能AIV展开利用Device的Vector Core计算单元来加速AllReduce操作可通过环境变量HCCL_OP_EXPANSION_MODE使能AIV配置如下export HCCL_OP_EXPANSION_MODEAIV附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考