昇腾Text-embedding-inference推理性能优化实践：大幅提升向量推理性能

在AI大模型蓬勃发展的今天，文本嵌入（Text Embedding）作为自然语言处理的基础能力，广泛应用于语义检索、文本聚类、相似度计算等场景。HuggingFace开源的text-embedding-inference框架因其高效的推理性能和灵活的模型支持，成为业界广泛采用的Embedding服务解决方案。

为了让text-embedding-inference框架更好地在NPU上运行，充分发挥NPU的算力，昇腾从底层算子到上层模型重构，全链路优化了框架，在embedding模型推理上取得了很好的效果。本文将以Qwen3-Embedding-0.6B模型为例，展示优化过程的实现细节与性能成果。

核心优化一：TND格式适配

什么是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 格式对比图

性能收益

TND格式适配带来的收益体现在两个方面：

计算效率提升：消除padding开销，实际有效计算占比从约60%提升至接近100%

调用亲和算子：引入的NPU融合注意力算子对TND格式有原生优化，减少数据格式转换带来的内存搬运耗时

核心优化二：

NPU亲和算子替换 —— 69.9%提升

昇腾NPU提供了丰富的融合算子库，针对硬件特性深度优化，相比细粒度算子可实现大幅度性能提升。本项目重点替换了以下核心算子：

NPU融合算子替换架构图

归一化层：npu_rms_norm

RMS Norm是Qwen3模型使用的归一化方式，原始实现需要多次数据类型转换：

# NPU融合实现return torch_npu.npu_rms_norm( hidden_states.to(input_dtype), self.weight, epsilon=self.variance_epsilon)[0]

线性层：npu_linear

将所有投影层的F.linear替换为npu_linear，减少算子调度开销：

# flash_qwen3.py Qwen3Attention.forwardq = self.q_norm.forward( torch_npu.npu_linear(hidden_states, self.q_proj_weight, bias=None) .view(*input_shape, self.num_heads, self.head_dim))k = self.k_norm.forward( torch_npu.npu_linear(hidden_states, self.k_proj_weight, bias=None) .view(*input_shape, self.num_key_value_heads, self.head_dim))

为什么npu_linear比nn.Linear更快？

实际实现中，两者最终都调用同一个底层CANN算子aclnnAddmm，但调用路径不同：

调用链对比：

在不同shape下性能测试对比如下：

原因分析：NPU异步执行模型

结论：

旋转位置编码：npu_rotary_mul

RoPE（Rotary Position Embedding）是Transformer中关键的位置编码方式。我们实现了完整的NPU版本：

def apply_rotary_pos_emb_npu(q, k, cos, sin, unsqueeze_dim=1): # 预处理：将3D [T, N, D] 转换为 4D [1, T, N, D] def _pre_process(x, cos, sin): origin_shape = x.shape if len(origin_shape) == 3: x = x.unsqueeze(0) # TND → BTND return x, 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] cos = cos.reshape(1, -1, 1, head_dim) sin = sin.reshape(1, -1, 1, head_dim) # 调用NPU融合算子 output_q = torch_npu.npu_rotary_mul(q_, cos, sin) output_k = torch_npu.npu_rotary_mul(k_, cos, sin) # 后处理：还原为3D格式 return output_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格式连贯性。

MLP层：npu_swiglu融合

Qwen3使用SwiGLU激活函数，原始实现需要三次矩阵乘法，计算负载较重。

MLP层优化前后对比图

# 优化实现：权重合并 + 融合算子self.gate_up_proj_weight = torch.cat( [self.gate_proj_weight, self.up_proj_weight], dim=0)def forward(self, hidden_state): # 一次linear得到gate和up的拼接结果 gate_up_states = torch_npu.npu_linear( hidden_state, self.gate_up_proj_weight, ) # npu_swiglu融合SiLU激活和门控乘法 hidden_states = torch_npu.npu_swiglu(gate_up_states, dim=-1) return torch_npu.npu_linear(hidden_states, self.down_proj_weight,)

这一优化将MLP层从3次矩阵乘法减少到2次，同时消除了中间结果的显存分配。

核心优化三：npu_add_rms_norm融合算子

优化背景

在Transformer架构中，每个Decoder层包含两个关键路径：注意力计算和前馈网络（MLP）。传统实现中，残差连接与层归一化是分离的两个操作：

# 传统实现（两次独立操作）residual = hidden_stateshidden_states =self.input_layernorm.forward(hidden_states)attn_output =self.attention.forward(...)hidden_states = residual + attn_output # 残差相加hidden_states =self.post_attention_layernorm.forward(hidden_states) # 归一化

这种分离实现存在两个问题：

显存访问开销：残差相加需要读取两个张量，归一化需要再次读取结果

算子调度开销：两次独立算子调用，增加CPU调度负担

融合算子实现

昇腾NPU提供了npu_add_rms_norm融合算子，将"残差相加 + RMSNorm"合并为单一操作：

# flash_qwen3.py Qwen3DecoderLayer.forwardhidden_states, _, residual = torch_npu.npu_add_rms_norm( residual, # 残差张量 attn_output, # 注意力输出 self.post_attention_layernorm.weight, self.post_attention_layernorm.variance_epsilon,)

融合算子架构图

性能收益分析

融合算子的核心收益来自两个方面：

显存带宽优化：减少显存读写次数，降低内存带宽压力

算子调度优化：减少CPU调度开销，降低延迟

核心优化四：注意力融合算子

本文采用npu_fused_infer_attention_score（FIA）算子。

GQA原生支持：移除interleave操作的关键

Qwen3-Embedding-0.6B采用GQA（Grouped Query Attention）架构，Query头数为16，Key/Value头数为8。传统实现中，为了适配不支持GQA的注意力算子，需要将K/V头数扩展至与Query头数一致：

# 传统实现（需要interleave扩头）if self.num_key_value_groups > 1: k = k.repeat_interleave(self.num_key_value_groups, dim=1) # [T, 8, D] → [T, 16, D] v = v.repeat_interleave(self.num_key_value_groups, dim=1) # [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 MiB

V额外复制量：2006 × 8 × 128 × 2 bytes ≈ 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( query=q, # [T, 16, D] key=k, # [T, 8, D] - 保持原始头数 value=v, # [T, 8, D] - 保持原始头数 num_heads=16, num_key_value_heads=8, # 原生GQA支持 input_layout="TND", ...)

算子级性能对比

对优化前后的算子耗时进行详细分析如下：

FIA核心收益总结

核心优化五：QKV投影融合

在Transformer的Attention层中，QKV投影是三个独立的线性变换操作，将隐藏状态分别投影到Query、Key、Value三个空间。传统实现中，这三个投影是串行执行的：

# 传统实现（三次独立矩阵乘法）q = F.linear(hidden_states, self.q_proj_weight) # [T, H] → [T, num_heads * head_dim]k = F.linear(hidden_states, self.k_proj_weight) # [T, H] → [T, num_kv_heads * head_dim]v = F.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_states = hidden_states.contiguous() qkv_outputs = torch_npu.npu_grouped_matmul( [hidden_states, hidden_states, hidden_states], # 三个相同的输入 self._grouped_qkv_weights, # 三个权重矩阵 group_type=-1, # 不分组，每个张量独立计算 split_item=0, # 输出多个张量 ) return qkv_outputs[0], qkv_outputs[1], qkv_outputs[2] # q, k, v

QKV投影优化架构图

性能收益分析

QKV投影融合带来的收益主要体现在：

性能对比

经过多轮迭代优化，最终实现了显著的性能提升。以下是完整的性能测试数据：

性能对比可视化图表

各优化模块对性能收益贡献分析

结语

通过对模型推理全链路的深度优化，取得了如下的成果：

性能显著提升：通过TND格式适配、NPU融合算子替换、QKV投影融合等核心优化，实现了近100%的性能提升。

技术方案可复用：沉淀当前的优化方案，可推广到其他embedding模型

参考资料

op-plugin开源项目：

https://gitcode.com/Ascend/op-plugin

昇腾CANN开发文档：

https://www.hiascend.com/document

text-embedding-inference原项目：

https://gitcode.com/Ascend/text-embeddings-inference

昇腾Text-embedding-inference推理性能优化实践：大幅提升向量推理性能

什么是TND格式？为什么要做格式转换？

融合算子架构图

性能对比可视化图表

评论 (0)

相关文章