从GPU到NPUQwen3-Embedding模型昇腾适配与性能优化实战作者: 昇腾实战派项目背景跨平台迁移的技术挑战在AI大模型蓬勃发展的今天文本嵌入Text Embedding作为自然语言处理的基础能力广泛应用于语义检索、文本聚类、相似度计算等场景。HuggingFace开源的text-embedding-inference框架因其高效的推理性能和灵活的模型支持成为业界广泛采用的Embedding服务解决方案。然而该框架原生仅支持GPU和HPU平台。随着国产AI芯片生态的快速发展将主流推理框架迁移到华为昇腾NPU平台成为推动AI基础设施国产化的重要一环。本项目正是在这一背景下启动——将text-embedding-inference框架完整适配到昇腾NPU使Qwen3-Embedding模型能够在国产硬件上高效运行。本文将详细分享从底层算子适配到上层模型重构的全链路优化工作过程中的技术决策、实现细节与性能成果。核心优化一TND格式适配——打破BNSD的性能瓶颈什么是TND格式为什么要做格式转换在传统的Transformer推理中数据通常以BNSD格式Batch × Sequence_length × Num_heads × Head_dim组织。这种格式需要将不同长度的序列通过padding对齐到相同长度存在两个明显缺陷内存浪费短序列需要大量padding无效计算占比高带宽压力padding位置的数据搬运消耗宝贵的显存带宽TND格式Total_tokens × Num_heads × Head_dim采用变长序列压缩策略将批次内所有token展平为连续存储通过cu_seqlens累积序列长度记录每个序列的边界。这种格式天然适合NPU的融合算子可显著减少无效计算。BNSD vs TND 格式对比图代码实现要点在flash_qwen3.py中TND格式适配贯穿整个模型前向传播流程# embed函数中的TND格式处理defembed(self,batch):ifisinstance(batch,FlashBatch):cu_seqlensbatch.cu_seqlens# 累积序列长度# 模型输出为 [T, H] 格式T总token数H隐藏维度outputself.model.forward(input_idsbatch.input_ids,...)hidden_statesoutput.last_hidden_state# shape: [T, H]# 直接通过cu_seqlens索引取最后token无需重组last_token_indicescu_seqlens[1:]-1embeddinghidden_states[last_token_indices]在flash_attn.py中注意力算子同样需要适配TND格式defnpu_attn(q,k,v,num_heads,out,seqlen_q,...):# npu_fusion_attention 原生支持TND layoutout_torch_npu.npu_fusion_attention(queryq,keyk,valuev,head_numnum_heads,input_layoutTND,# 关键指定TND格式actual_seq_qlenseqlen_q.tolist(),actual_seq_kvlenseqlen_k.tolist(),)[0]技术细节cu_seqlens设备放置在TND格式适配过程中一个容易被忽视但影响性能的细节是cu_seqlens​的设备放置。npu_fusion_attention​算子要求actual_seq_qlen​参数为Python list类型。如果cu_seqlens​存储在NPU设备上调用.tolist()方法时会触发D2HDevice to Host数据传输带来额外的延迟开销。解决方案在types.py中将cu_seqlens强制放在CPU上cu_seqlenstorch.tensor(pb.cu_seq_lengths,dtypetorch.int32,devicecpu)这样.tolist()操作直接在CPU内存中完成避免了跨设备数据传输。性能收益TND格式适配带来的收益体现在两个方面内存效率提升消除padding开销实际有效计算占比从约60%提升至接近100%算子调用优化NPU融合注意力算子对TND格式有原生优化减少数据搬运次数核心优化二NPU融合算子替换——释放硬件极致性能昇腾NPU提供了丰富的融合算子库通过torch_npu模块访问这些算子针对硬件特性深度优化相比通用PyTorch算子可实现数倍性能提升。本项目重点替换了以下核心算子NPU融合算子替换架构图2.1 注意力层npu_fusion_attention这是最核心的优化点。原始代码依赖CUDA的Flash Attention在NPU上无法运行。我们通过torch_npu.npu_fusion_attention实现了等效功能# flash_attn.py 核心实现defnpu_attn(q,k,v,num_heads,out,seqlen_q,seqlen_k,...):ifis_causal:# 构造causal maskattn_mask_nputorch.triu(torch.ones((2048,2048),dtypetorch.bool,deviceq.device),diagonal1)out_torch_npu.npu_fusion_attention(queryq,keyk,valuev,head_numnum_heads,input_layoutTND,scalesoftmax_scale,actual_seq_qlenseqlen_q.tolist(),actual_seq_kvlenseqlen_k.tolist(),sparse_mode3,# causal模式atten_maskattn_mask_npu)[0]2.2 归一化层npu_rms_normRMS Norm是Qwen3模型使用的归一化方式原始实现需要多次数据类型转换# 原始实现已弃用# hidden_states hidden_states.to(torch.float32)# variance hidden_states.pow(2).mean(-1, keepdimTrue)# hidden_states hidden_states * torch.rsqrt(variance eps)# NPU融合实现returntorch_npu.npu_rms_norm(hidden_states.to(input_dtype),self.weight,epsilonself.variance_epsilon)[0]2.3 线性层npu_linear将所有投影层的F.linear​替换为npu_linear减少算子调度开销# flash_qwen3.py Qwen3Attention.forwardqself.q_norm.forward(torch_npu.npu_linear(hidden_states,self.q_proj_weight,biasNone).view(*input_shape,self.num_heads,self.head_dim))kself.k_norm.forward(torch_npu.npu_linear(hidden_states,self.k_proj_weight,biasNone).view(*input_shape,self.num_key_value_heads,self.head_dim))为什么npu_linear​比nn.Linear更快这是一个关键的技术问题。经过深入分析我们发现两者最终都调用同一个底层CANN算子aclnnAddmm但调用路径不同调用链对比性能对比测试我们在不同shape下进行了详细测试Config (Batch, IN, OUT)nn.Linear (µs)npu_linear (µs)性能差异B4, IN128, OUT6437.345.24npu_linear快7.1×B32, IN1024, OUT51224.856.99npu_linear快3.6×B128, IN4096, OUT2048103.21103.36基本持平B1024, IN4096, OUT2048759.41759.40基本持平原因分析NPU异步执行模型结论何时使用npu_linear场景建议原因小batch推理✅ 优先用npu_lineardispatch开销占比大可节省大量时间大矩阵计算密集两者无差异kernel时间远大于dispatch用nn.Linear可读性更好Embedding服务✅ 推荐npu_linear典型小batch场景性能提升显著2.4 旋转位置编码npu_rotary_mulRoPERotary Position Embedding是Transformer中关键的位置编码方式。我们实现了完整的NPU版本defapply_rotary_pos_emb_npu(q,k,cos,sin,unsqueeze_dim1):# 预处理将3D [T, N, D] 转换为 4D [1, T, N, D]def_pre_process(x,cos,sin):origin_shapex.shapeiflen(origin_shape)3:xx.unsqueeze(0)# TND → BTNDreturnx,cos,sin,origin_shape,x.dtype q_,cos,sin,q_shape,q_dtype_pre_process(q,cos,sin)k_,cos,sin,k_shape,k_dtype_pre_process(k,cos,sin)# cos/sin reshape为 [1, T, 1, D]coscos.reshape(1,-1,1,head_dim)sinsin.reshape(1,-1,1,head_dim)# 调用NPU融合算子output_qtorch_npu.npu_rotary_mul(q_,cos,sin)output_ktorch_npu.npu_rotary_mul(k_,cos,sin)# 后处理还原为3D格式returnoutput_q.squeeze(0),output_k.squeeze(0)技术细节TND格式的维度适配​npu_rotary_mul​算子要求输入为4D张量[B, T, N, D]​而TND格式的q/k是3D张量[T, N, D]​。上述代码通过_pre_process​和_post_process函数实现了优雅的维度转换预处理unsqueeze(0)​将[T, N, D]​扩展为[1, T, N, D]模拟batch维度后处理squeeze(0)​将输出还原为[T, N, D]保持TND格式一致性这种设计既满足了算子接口要求又保持了数据流的TND格式连贯性。2.5 MLP层npu_swiglu融合这是性能提升最大的优化点。Qwen3使用SwiGLU激活函数原始实现需要三次矩阵乘法MLP层优化前后对比图# 原始实现已弃用# gated F.linear(hidden_state, self.gate_proj_weight)# uped F.linear(hidden_state, self.up_proj_weight)# output F.linear(silu(gated) * uped, self.down_proj_weight)# 优化实现权重合并 融合算子self.gate_up_proj_weighttorch.cat([self.gate_proj_weight,self.up_proj_weight],dim0)defforward(self,hidden_state):# 一次linear得到gate和up的拼接结果gate_up_statestorch_npu.npu_linear(hidden_state,self.gate_up_proj_weight,)# npu_swiglu融合SiLU激活和门控乘法hidden_statestorch_npu.npu_swiglu(gate_up_states,dim-1)returntorch_npu.npu_linear(hidden_states,self.down_proj_weight,)这一优化将MLP层从3次矩阵乘法减少到2次同时消除了中间结果的显存分配。核心优化三Pooling层重构——去除第三方依赖原始代码依赖sentence_transformers库的Pooling实现在NPU环境存在兼容性问题。我们完全重写了Pooling层实现纯PyTorch原生版本# pooling.py 核心实现classLastTokenPooling(_Pooling):取每个序列最后一个真实token适合decoder-only模型defforward(self,model_output,attention_mask)-Tensor:token_embeddingsmodel_output[0]attention_maskattention_mask.to(dtypetorch.bool)# 计算每个序列的实际长度last_indicesattention_mask.sum(dim1,keepdimTrue)-1last_indiceslast_indices.clamp(min0)# 使用gather提取最后tokenlast_indices_expandedlast_indices.unsqueeze(-1).expand(-1,-1,token_embeddings.shape[-1])returntoken_embeddings.gather(1,last_indices_expanded).squeeze(1)classDefaultPooling(_Pooling):工厂路由类根据pooling_mode派发到对应实现def__init__(self,hidden_size,pooling_modelast):ifpooling_modemean:self.poolingMeanPooling(hidden_size,pooling_mode)elifpooling_modemax:self.poolingMaxPooling(hidden_size,pooling_mode)elifpooling_modecls:self.poolingCLSPooling(hidden_size,pooling_mode)elifpooling_modein(last,last_token):self.poolingLastTokenPooling(hidden_size,pooling_mode)新增的LastTokenPooling​类专为Qwen3等decoder-only模型设计通过attention_mask计算真实序列长度精确提取最后一个有效token的embedding。核心优化四npu_add_rms_norm融合算子——残差与归一化的深度合并优化背景在Transformer架构中每个Decoder层包含两个关键路径注意力计算和前馈网络MLP。传统实现中残差连接与层归一化是分离的两个操作# 传统实现两次独立操作residualhidden_states hidden_statesself.input_layernorm.forward(hidden_states)attn_outputself.attention.forward(...)hidden_statesresidualattn_output# 残差相加hidden_statesself.post_attention_layernorm.forward(hidden_states)# 归一化这种分离实现存在两个问题显存访问开销残差相加需要读取两个张量归一化需要再次读取结果算子调度开销两次独立算子调用增加CPU调度负担融合算子实现昇腾NPU提供了npu_add_rms_norm融合算子将残差相加 RMSNorm合并为单一操作# flash_qwen3.py Qwen3DecoderLayer.forwardhidden_states,_,residualtorch_npu.npu_add_rms_norm(residual,# 残差张量attn_output,# 注意力输出self.post_attention_layernorm.weight,self.post_attention_layernorm.variance_epsilon,)融合算子架构图性能收益分析融合算子的核心收益来自两个方面显存带宽优化减少50%的显存读写次数降低内存带宽压力算子调度优化减少50%的CPU调度开销降低延迟核心优化五注意力算子升级——从npu_fusion_attention到npu_fused_infer_attention_score两种算子的本质差异在昇腾NPU平台上存在两种注意力计算算子​npu_fusion_attention​对应flash_attention_score​是一个通用FlashAttention框架内核​npu_fused_infer_attention_score​FIA对应fused_infer_attention_score​是一个推理态专用路由器专用内核入口从源码层面分析两者存在根本性设计差异‍GQA原生支持移除interleave操作的关键Qwen3-Embedding采用GQAGrouped Query Attention架构Query头数为16Key/Value头数为8。传统实现中为了适配不支持GQA的注意力算子需要将K/V头数扩展至与Query头数一致# 传统实现需要interleave扩头ifself.num_key_value_groups1:kk.repeat_interleave(self.num_key_value_groups,dim1)# [T, 8, D] → [T, 16, D]vv.repeat_interleave(self.num_key_value_groups,dim1)# [T, 8, D] → [T, 16, D]interleave操作的性能开销分析以一个真实batch为例q_shape(2006, 16, 128)​,k_shape(2006, 8, 128)K额外复制量2006 × 8 × 128 × 2 bytes ≈ 3.9 MiBV额外复制量≈ 3.9 MiB合计约7.8 MiB/layer/batch的纯复制流量而FIA算子原生支持GQA可直接传入num_heads​和num_key_value_heads参数# FIA实现原生支持GQA无需扩头torch_npu.npu_fused_infer_attention_score(queryq,# [T, 16, D]keyk,# [T, 8, D] - 保持原始头数valuev,# [T, 8, D] - 保持原始头数num_heads16,num_key_value_heads8,# 原生GQA支持input_layoutTND,...)算子级性能对比通过昇腾Profiling工具对优化前后的算子耗时进行详细分析FIA核心收益总结收益来源具体表现GQA原生支持避免K/V扩头复制节省约7.8 MiB/layer/batch的显存带宽TND热路径直接命中FAInfer路径减少通用框架开销中间状态精简更少的GM写回降低显存压力控制面轻量化更简化的epilogue和调度逻辑核心优化六QKV投影融合——npu_grouped_matmul批量矩阵乘法优化优化背景在Transformer的Attention层中QKV投影是三个独立的线性变换操作将隐藏状态分别投影到Query、Key、Value三个空间。传统实现中这三个投影是串行执行的# 传统实现三次独立矩阵乘法qF.linear(hidden_states,self.q_proj_weight)# [T, H] → [T, num_heads * head_dim]kF.linear(hidden_states,self.k_proj_weight)# [T, H] → [T, num_kv_heads * head_dim]vF.linear(hidden_states,self.v_proj_weight)# [T, H] → [T, num_kv_heads * head_dim]这种实现存在以下问题算子调度开销三次独立的算子调用增加CPU调度负担内存访问开销hidden_states需要被读取三次增加显存带宽压力并行性不足三个投影操作本质上是独立的可以并行执行npu_grouped_matmul算子原理昇腾NPU提供了npu_grouped_matmul算子支持批量矩阵乘法操作。该算子可以将多个矩阵乘法合并为一次API调用减少算子调度开销和内存访问次数。核心参数说明​x输入张量列表每个张量代表一个矩阵乘法的左操作数​weight权重张量列表每个张量代表一个矩阵乘法的右操作数​group_type​分组类型-1表示不分组每个张量独立计算​split_item​输出分割模式0表示输出多个张量与weight数量相同QKV投影融合实现在flash_qwen3.py中实现了QKV投影的融合优化# Qwen3Attention.__init__ 中的权重预处理self._grouped_qkv_weights[self.q_proj_weight.transpose(0,1).contiguous(),# [H, num_heads * head_dim]self.k_proj_weight.transpose(0,1).contiguous(),# [H, num_kv_heads * head_dim]self.v_proj_weight.transpose(0,1).contiguous(),# [H, num_kv_heads * head_dim]]# Qwen3Attention._project_qkv 中的融合投影def_project_qkv(self,hidden_states): 使用npu_grouped_matmul融合QKV投影 将三次独立矩阵乘法合并为一次API调用 hidden_stateshidden_states.contiguous()qkv_outputstorch_npu.npu_grouped_matmul([hidden_states,hidden_states,hidden_states],# 三个相同的输入self._grouped_qkv_weights,# 三个权重矩阵group_type-1,# 不分组每个张量独立计算split_item0,# 输出多个张量)returnqkv_outputs[0],qkv_outputs[1],qkv_outputs[2]# q, k, vQKV投影优化架构图技术细节权重预处理​npu_grouped_matmul算子要求权重矩阵以特定格式组织。在实现中需要注意以下细节权重转置NPU的矩阵乘法算子期望权重为[out_features, in_features]​格式而PyTorch的nn.Linear​权重为[out_features, in_features]​需要转置为[in_features, out_features]连续内存使用.contiguous()确保权重张量在内存中连续存储避免非连续内存访问带来的性能损失权重列表将三个投影权重组织为Python列表作为npu_grouped_matmul的输入性能收益分析QKV投影融合带来的收益主要体现在优化维度收益描述算子调度从3次独立调用减少到1次批量调用减少67%的调度开销显存带宽hidden_states只需读取1次减少67%的输入数据读取并行性三个投影操作在算子内部并行执行提高硬件利用率性能对比优化效果一目了然经过多轮迭代优化最终实现了显著的性能提升。以下是完整的性能测试数据性能对比可视化图表关键性能指标指标最佳配置相比原生提升最高QPS全NPU算子融合 4线程69.9%最低总时间全NPU算子融合 4线程41.1%最低响应时间NPU融合注意力RMSNorm 1线程30.2%各优化模块贡献分析优化模块QPS提升主要收益来源npu_fusion_attention57.6%消除注意力计算的显存瓶颈支持TND格式npu_rms_norm2.5%减少数据类型转换融合方差计算npu_linear3.5%减少算子调度开销优化内存访问npu_rotary_mul2.0%融合旋转编码计算npu_swiglu4.3%MLP层从3次矩阵乘减少到2次npu_grouped_matmul调度优化QKV投影融合减少67%算子调用性能对比可视化QPS性能对比柱状图核心结论长文本场景表现优异1024上下文场景下NPU与A100性能差距缩小至3%以内高批量时NPU甚至反超中高批量接近持平Batch≥8时NPU达到A10095%以上性能小批量场景差距明显Batch1时A100仍有约20-37%优势性价比优势NPU在embedding推理场景具备商用可行性项目价值与收获项目价值推动国产AI生态成功将主流Embedding推理框架迁移到昇腾NPU为国产AI基础设施添砖加瓦性能显著提升相比原生实现QPS提升近70%为实际业务部署提供了强有力的性能保障技术方案可复用TND格式适配、NPU融合算子替换等方案可推广到其他Transformer模型项目收获底层算子优化的重要性一个融合算子可能带来10%以上的性能提升数据格式设计对性能的影响TND格式从架构层面解决了padding开销问题跨平台适配的挑战不同硬件平台的算子接口、性能特性差异巨大需要针对性优化未来展望算子库持续完善基于op-plugin开源项目探索更多NPU专用算子多模型支持将适配方案扩展到更多Embedding模型如BGE、M3E等分布式推理探索NPU集群上的分布式Embedding服务方案结语从GPU到NPU的迁移不仅是硬件平台的切换更是对模型推理全链路的深度优化。本项目通过TND格式适配、NPU融合算子替换、Pooling层重构、QKV投影融合等核心优化实现了近70%的性能提升证明了国产AI芯片在主流推理场景下的竞争力。技术迭代永无止境期待与更多开发者一起共同推动国产AI生态的繁荣发展参考资料op-plugin开源项目昇腾CANN开发文档text-embedding-inference原项目