Qwen3-Embedding模型在昇腾NPU的优化实践

1. 项目背景与挑战

在自然语言处理领域，文本嵌入（Text Embedding）技术已经成为语义检索、文本聚类和相似度计算等任务的基础能力。HuggingFace开源的text-embedding-inference框架因其高效的推理性能和灵活的模型支持，成为业界广泛采用的解决方案。然而，该框架原生仅支持GPU和HPU平台，这为国产AI芯片生态的发展带来了挑战。

作为一名长期从事AI模型优化的工程师，我最近完成了将Qwen3-Embedding模型完整适配到华为昇腾NPU平台的工作。这个项目不仅涉及简单的平台迁移，更需要对模型架构和计算流程进行深度优化，以充分发挥NPU硬件的性能潜力。在这个过程中，我们遇到了几个关键挑战：

1. 计算格式适配：传统Transformer推理使用的BNSD格式在NPU上效率低下，需要转换为更适合硬件特性的TND格式
2. 算子兼容性：原框架依赖的CUDA专用算子（如Flash Attention）无法在NPU上运行
3. 性能瓶颈：NPU与GPU的架构差异导致直接迁移的性能无法满足生产需求

2. 核心优化方案与技术实现

2.1 TND格式适配：突破传统计算瓶颈

2.1.1 BNSD格式的局限性

在传统的Transformer推理中，数据通常以BNSD格式（Batch × Sequence_length × Num_heads × Head_dim）组织。这种格式存在两个明显缺陷：

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

以一个典型场景为例（batch_size=32，max_seq_len=512，平均seq_len=256）：

• 有效计算占比仅为50%
• 显存带宽浪费约50%

2.1.2 TND格式的优势

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

技术实现上，我们在flash_qwen3.py中完成了TND格式的全面适配：

python复制def embed(self, batch):
    if isinstance(batch, FlashBatch):
        cu_seqlens = batch.cu_seqlens  # 累积序列长度
        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]

2.1.3 关键优化：cu_seqlens设备放置

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

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

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

2.2 NPU融合算子替换策略

2.2.1 注意力层优化

原始代码依赖CUDA的Flash Attention，在NPU上无法运行。我们通过torch_npu.npu_fusion_attention实现了等效功能：

python复制def npu_attn(q, k, v, num_heads, out, seqlen_q, seqlen_k, ...):
    if is_causal:
        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.2 归一化层优化

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

python复制# 原始实现（已弃用）
# 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.2.3 MLP层优化

Qwen3使用SwiGLU激活函数，原始实现需要三次矩阵乘法。我们通过权重合并和npu_swiglu融合算子优化：

python复制self.gate_up_proj_weight = torch.cat(
    [self.gate_proj_weight, self.up_proj_weight], dim=0
)

def forward(self, hidden_state):
    gate_up_states = torch_npu.npu_linear(
        hidden_state, self.gate_up_proj_weight,
    )
    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次，同时消除了中间结果的显存分配。

2.3 残差与归一化的深度合并

传统实现中，残差连接与层归一化是分离的两个操作：

python复制# 传统实现（两次独立操作）
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)  # 归一化

我们使用npu_add_rms_norm融合算子将这两个操作合并：

python复制hidden_states, _, residual = torch_npu.npu_add_rms_norm(
    residual,                    # 残差张量
    attn_output,                 # 注意力输出
    self.post_attention_layernorm.weight,
    self.post_attention_layernorm.variance_epsilon,
)

2.4 QKV投影融合优化

传统QKV投影是三个独立的线性变换操作，我们使用npu_grouped_matmul算子进行融合：

python复制def _project_qkv(self, hidden_states):
    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

这一优化将三次独立矩阵乘法合并为一次API调用，减少了67%的调度开销和显存带宽压力。

3. 性能优化成果

经过上述优化，我们实现了显著的性能提升：

1. 长文本场景表现优异：1024上下文场景下，NPU与A100性能差距缩小至3%以内
2. 中高批量接近持平：Batch≥8时，NPU达到A100 95%以上性能
3. 整体QPS提升：相比原生实现，QPS提升近70%

具体性能数据对比如下：

4. 实践经验与注意事项

在实际优化过程中，我们积累了一些宝贵经验：

1. 算子选择策略：

   • 小batch推理优先使用npu_linear
   • 大矩阵计算密集场景两者无差异，可用nn.Linear保持可读性
   • Embedding服务推荐使用npu_linear

2. 调试技巧：

   • 使用torch_npu.npu.synchronize()确保准确测量NPU算子耗时
   • 通过torch_npu.npu.profiler进行性能分析
   • 注意NPU特有的内存对齐要求

3. 常见问题：

   • 算子输入输出格式必须严格匹配文档要求
   • 混合精度训练时注意保持计算图一致性
   • NPU特有的内存管理策略可能导致与GPU不同的行为

5. 项目价值与展望

本次优化工作不仅实现了Qwen3-Embedding模型在昇腾NPU上的高效运行，更为国产AI生态建设积累了宝贵经验。未来我们计划：

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

从GPU到NPU的迁移不仅是硬件平台的切换，更是对模型推理全链路的深度优化。通过TND格式适配、NPU融合算子替换等核心优化，我们证明了国产AI芯片在主流推理场景下的竞争力。