从GPU到NPU：Qwen3-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对齐到相同长度，存在两个明显缺陷：

1. 内存浪费：短序列需要大量padding，无效计算占比高
2. 带宽压力：padding位置的数据搬运消耗宝贵的显存带宽

TND格式（Total_tokens × Num_heads × Head_dim）采用变长序列压缩策略，将批次内所有token展平为连续存储，通过cu_seqlens（累积序列长度）记录每个序列的边界。这种格式天然适合NPU的融合算子，可显著减少无效计算。

BNSD vs TND 格式对比图

代码实现要点

在flash_qwen3.py中，TND格式适配贯穿整个模型前向传播流程：

# embed函数中的TND格式处理
def embed(self, batch):
    if isinstance(batch, FlashBatch):
        cu_seqlens = batch.cu_seqlens  # 累积序列长度
        # 模型输出为 [T, H] 格式（T=总token数，H=隐藏维度）
        output = self.model.forward(input_ids=batch.input_ids, ...)
        hidden_states = output.last_hidden_state  # shape: [T, H]
        
        # 直接通过cu_seqlens索引取最后token，无需重组
        last_token_indices = cu_seqlens[1:] - 1
        embedding = hidden_states[last_token_indices]

在flash_attn.py中，注意力算子同样需要适配TND格式：

def npu_attn(q, k, v, num_heads, out, seqlen_q, ...):
    # npu_fusion_attention 原生支持TND layout
    out_ = torch_npu.npu_fusion_attention(
        query=q, key=k, value=v,
        head_num=num_heads,
        input_layout="TND",  # 关键：指定TND格式
        actual_seq_qlen=seqlen_q.tolist(),
        actual_seq_kvlen=seqlen_k.tolist(),
    )[0]

技术细节：cu_seqlens设备放置

在TND格式适配过程中，一个容易被忽视但影响性能的细节是cu_seqlens​的设备放置。npu_fusion_attention​算子要求actual_seq_qlen​参数为Python list类型。如果cu_seqlens​存储在NPU设备上，调用.tolist()方法时会触发D2H（Device to Host）数据传输，带来额外的延迟开销。

解决方案：在types.py中将cu_seqlens强制放在CPU上：

cu_seqlens = torch.tensor(pb.cu_seq_lengths, dtype=torch.int32, device="cpu")

这样.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 核心实现
def npu_attn(q, k, v, num_heads, out, seqlen_q, seqlen_k, ...):
    if is_causal:
        # 构造causal mask
        attn_mask_npu = torch.triu(
            torch.ones((2048, 2048), dtype=torch.bool, device=q.device), 
            diagonal=1
        )
        out_ = torch_npu.npu_fusion_attention(
            query=q, key=k, value=v,
            head_num=num_heads,
            input_layout="TND",
            scale=softmax_scale,
            actual_seq_qlen=seqlen_q.tolist(),
            actual_seq_kvlen=seqlen_k.tolist(),
            sparse_mode=3,  # causal模式
            atten_mask=attn_mask_npu
        )[0]

2.2 归一化层：npu_rms_norm

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

# 原始实现（已弃用）
# hidden_states = hidden_states.to(torch.float32)
# variance = hidden_states.pow(2).mean(-1, keepdim=True)
# hidden_states = hidden_states * torch.rsqrt(variance + eps)

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

2.3 线性层：npu_linear

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

# flash_qwen3.py Qwen3Attention.forward
q = 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下进行了详细测试：

Config (Batch, IN, OUT)	nn.Linear (µs)	npu_linear (µs)	性能差异
B=4, IN=128, OUT=64	37.34	5.24	npu_linear快 7.1×
B=32, IN=1024, OUT=512	24.85	6.99	npu_linear快 3.6×
B=128, IN=4096, OUT=2048	103.21	103.36	基本持平
B=1024, IN=4096, OUT=2048	759.41	759.40	基本持平

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

结论：何时使用npu_linear

场景	建议	原因
小batch推理	✅ 优先用npu_linear	dispatch开销占比大，可节省大量时间
大矩阵计算密集	两者无差异	kernel时间远大于dispatch，用nn.Linear可读性更好
Embedding服务	✅ 推荐npu_linear	典型小batch场景，性能提升显著

2.4 旋转位置编码：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格式连贯性。

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_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次，同时消除了中间结果的显存分配。

---

核心优化三：Pooling层重构——去除第三方依赖

原始代码依赖sentence_transformers库的Pooling实现，在NPU环境存在兼容性问题。我们完全重写了Pooling层，实现纯PyTorch原生版本：

# pooling.py 核心实现
class LastTokenPooling(_Pooling):
    """取每个序列最后一个真实token（适合decoder-only模型）"""
    def forward(self, model_output, attention_mask) -> Tensor:
        token_embeddings = model_output[0]
        attention_mask = attention_mask.to(dtype=torch.bool)
        # 计算每个序列的实际长度
        last_indices = attention_mask.sum(dim=1, keepdim=True) - 1
        last_indices = last_indices.clamp(min=0)
        # 使用gather提取最后token
        last_indices_expanded = last_indices.unsqueeze(-1).expand(
            -1, -1, token_embeddings.shape[-1]
        )
        return token_embeddings.gather(1, last_indices_expanded).squeeze(1)

class DefaultPooling(_Pooling):
    """工厂路由类，根据pooling_mode派发到对应实现"""
    def __init__(self, hidden_size, pooling_mode="last"):
        if pooling_mode == "mean":
            self.pooling = MeanPooling(hidden_size, pooling_mode)
        elif pooling_mode == "max":
            self.pooling = MaxPooling(hidden_size, pooling_mode)
        elif pooling_mode == "cls":
            self.pooling = CLSPooling(hidden_size, pooling_mode)
        elif pooling_mode in ("last", "last_token"):
            self.pooling = LastTokenPooling(hidden_size, pooling_mode)

新增的LastTokenPooling​类专为Qwen3等decoder-only模型设计，通过attention_mask计算真实序列长度，精确提取最后一个有效token的embedding。

---

核心优化四：npu_add_rms_norm融合算子——残差与归一化的深度合并

优化背景

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

# 传统实现（两次独立操作）
residual = hidden_states
hidden_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)  # 归一化

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

1. 显存访问开销：残差相加需要读取两个张量，归一化需要再次读取结果
2. 算子调度开销：两次独立算子调用，增加CPU调度负担

融合算子实现

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

# flash_qwen3.py Qwen3DecoderLayer.forward
hidden_states, _, residual = torch_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采用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额外复制量：≈ 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",
    ...
)

算子级性能对比

通过昇腾Profiling工具，对优化前后的算子耗时进行详细分析：

FIA核心收益总结

收益来源	具体表现
GQA原生支持	避免K/V扩头复制，节省约7.8 MiB/layer/batch的显存带宽
TND热路径	直接命中FAInfer路径，减少通用框架开销
中间状态精简	更少的GM写回，降低显存压力
控制面轻量化	更简化的epilogue和调度逻辑

---

核心优化六：QKV投影融合——npu_grouped_matmul批量矩阵乘法优化

优化背景

在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]

这种实现存在以下问题：

1. 算子调度开销：三次独立的算子调用，增加CPU调度负担
2. 内存访问开销：hidden_states需要被读取三次，增加显存带宽压力
3. 并行性不足：三个投影操作本质上是独立的，可以并行执行

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投影优化架构图

技术细节：权重预处理

​npu_grouped_matmul算子要求权重矩阵以特定格式组织。在实现中，需要注意以下细节：

1. 权重转置：NPU的矩阵乘法算子期望权重为[out_features, in_features]​格式，而PyTorch的nn.Linear​权重为[out_features, in_features]​，需要转置为[in_features, out_features]
2. 连续内存：使用.contiguous()确保权重张量在内存中连续存储，避免非连续内存访问带来的性能损失
3. 权重列表：将三个投影权重组织为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_attention	+57.6%	消除注意力计算的显存瓶颈，支持TND格式
npu_rms_norm	+2.5%	减少数据类型转换，融合方差计算
npu_linear	+3.5%	减少算子调度开销，优化内存访问
npu_rotary_mul	+2.0%	融合旋转编码计算
npu_swiglu	+4.3%	MLP层从3次矩阵乘减少到2次
npu_grouped_matmul	+调度优化	QKV投影融合，减少67%算子调用

性能对比可视化

QPS性能对比柱状图

核心结论

1. 长文本场景表现优异：1024上下文场景下，NPU与A100性能差距缩小至3%以内，高批量时NPU甚至反超
2. 中高批量接近持平：Batch≥8时，NPU达到A100 95%以上性能
3. 小批量场景差距明显：Batch=1时，A100仍有约20-37%优势
4. 性价比优势：NPU在embedding推理场景具备商用可行性

---

项目价值与收获

项目价值

1. 推动国产AI生态：成功将主流Embedding推理框架迁移到昇腾NPU，为国产AI基础设施添砖加瓦
2. 性能显著提升：相比原生实现，QPS提升近70%，为实际业务部署提供了强有力的性能保障
3. 技术方案可复用：TND格式适配、NPU融合算子替换等方案可推广到其他Transformer模型

项目收获

• 底层算子优化的重要性：一个融合算子可能带来10%以上的性能提升
• 数据格式设计对性能的影响：TND格式从架构层面解决了padding开销问题
• 跨平台适配的挑战：不同硬件平台的算子接口、性能特性差异巨大，需要针对性优化

未来展望

1. 算子库持续完善：基于op-plugin开源项目，探索更多NPU专用算子
2. 多模型支持：将适配方案扩展到更多Embedding模型（如BGE、M3E等）
3. 分布式推理：探索NPU集群上的分布式Embedding服务方案

---

结语

从GPU到NPU的迁移不仅是硬件平台的切换，更是对模型推理全链路的深度优化。本项目通过TND格式适配、NPU融合算子替换、Pooling层重构、QKV投影融合等核心优化，实现了近70%的性能提升，证明了国产AI芯片在主流推理场景下的竞争力。

技术迭代永无止境，期待与更多开发者一起，共同推动国产AI生态的繁荣发展！

---

参考资料：

• op-plugin开源项目
• 昇腾CANN开发文档
• text-embedding-inference原项目