在NPU单算子（torch_npu ）执行时如何进行性能优化？以MinerU为例

1 MinerU介绍

在AI技术快速发展的今天，大量非结构化数据的处理成为亟待解决的问题。尤其是PDF文档，作为最常见的文件格式之一，如何高效准确地提取其中的信息，成为了许多企业和研究机构的痛点。上海人工智能实验室（上海AI实验室）大模型数据基座OpenDataLab团队开源了全新的智能数据提取工具——MinerU，旨在解决这一问题。

MinerU能够将混合了图片、公式、表格、脚注等复杂元素的PDF文档转化为Markdown和JSON格式，大幅提升了AI语料的准备效率。凭借快速准确、开源易用的能力特性，MinerU受到广大用户及大模型开发者青睐，上线八个月，GitHub星标数已接近3万，被开发者誉为"大模型时代的文档提取、转换神器"。

MinerU是由上海人工智能实验室OpenDataLab团队开发的开源文档解析工具，致力于解决大模型（LLM）训练和RAG（检索增强生成）应用中高质量结构化数据的提取难题。自2024年7月开源以来，GitHub星标数迅速突破2.5万，成为开发者社区的热门选择。其核心价值在于将复杂文档（如PDF、网页、电子书）转换为机器可读的Markdown、JSON格式，同时保留原始文档的语义逻辑与多模态元素。

获取MinerU 官方PPT：PDF 解析神器
MinerU.pdf

MinerU官网：MinerU

MinerU代码地址：GitHub - opendatalab/MinerU: A high-quality tool for
convert PDF to Markdown and
JSON.一站式开源高质量数据提取工具，将PDF转换成Markdown和JSON格式。

MinerU 线上demo入口：MinerU

本篇文章通过介绍其在NPU的单算子模式下的一些性能优化的手段。这些手段也可以给其他模型在torch_npu下的使用提供一些优化指导。

torch_npu上的性能调优指南，大家可以关注
PyTorch模型在NPU上的调优

2 torch_npu性能优化

2.1 绑核优化

在PyTorch的训练或推理场景，可以通过设置环境变量CPU_AFFINITY_CONF来控制CPU端算子任务的处理器亲和性，即设定任务绑核。该配置能够优化任务的执行效率，避免跨 NUMA（非统一内存访问架构）节点的内存访问，减少任务调度开销。

可选的绑核方案如下：

粗粒度绑核：将所有任务绑定在NPU对应NUMA的CPU核心上，避免跨NUMA节点的内存访问，并支持粗粒度绑核上的自定义绑核。
细粒度绑核：在粗粒度绑核的基础上进一步优化，将主要任务锚定在NUMA节点的某固定CPU核心上，减少核间切换的开销。

配置示例
示例一：粗粒度绑核

export CPU_AFFINITY_CONF=1

示例二：细粒度绑核

export CPU_AFFINITY_CONF=2

示例三：自定义多张NPU卡的绑核范围

export CPU_AFFINITY_CONF=1,npu0:0-1,npu1:2-5,npu3:6-6

具体可参考：torch_npu绑核优化

2.2 流水优化

一级流水优化
一级流水优化是一种通用有效的优化方法，主要是将部分算子适配任务迁移至二级流水，使两级流水负载更均衡，并减少dequeue唤醒时间。
使能说明
该特性可以通过环境变量设置，一般用于以下场景：host-bound严重的网络场景。

使用方法如下：

export TASK_QUEUE_ENABLE=2

该环境变量的实现细节请参考《环境变量参考》中的“TASK_QUEUE_ENABLE”章节。

说明
ASCEND_LAUNCH_BLOCKING设置为“1”时，task_queue算子队列关闭，TASK_QUEUE_ENABLE设置不生效。
TASK_QUEUE_ENABLE配置为“2”时，由于内存并发，可能导致运行中NPU内存峰值上升。

2.2.1 TASK_QUEUE_ENABLE

通过此环境变量可配置task_queue算子下发队列是否开启和优化等级。

配置为“0”时：关闭task_queue算子下发队列优化，算子下发任务如图1所示。
关闭task_queue
配置为“1”或未配置时：开启task_queue算子下发队列Level 1优化，算子下发任务如图2所示。

Level 1优化：使能task_queue算子下发队列优化，将算子下发任务分为两段，一部分任务（主要是aclnn算子的调用）放在新增的二级流水上，一、二级流水通过算子队列传递任务，相互并行，通过部分掩盖减少整体的下发耗时，提升端到端性能。
配置为“2”时：开启task_queue算子下发队列Level 2优化，算子下发任务如图3所示。
Level 2优化：包含Level 1的优化并进一步平衡了一、二级流水的任务负载，主要是将workspace相关任务迁移至二级流水，掩盖效果更好，性能收益更大。该配置仅在二进制场景生效，建议配置值为Level 2优化。
图3 Level 2优化

此环境变量默认配置为“1”。

注意
ASCEND_LAUNCH_BLOCKING设置为“1”时，task_queue算子队列关闭，TASK_QUEUE_ENABLE设置不生效。

TASK_QUEUE_ENABLE配置为“2”时，由于内存并发，可能导致运行中NPU内存峰值上升。

2.2.2 ASCEND_LAUNCH_BLOCKING

通过此环境变量可控制算子执行时是否启动同步模式。

由于在昇腾NPU上进行模型训练时默认算子异步执行，导致算子执行过程中出现报错时，打印的报错堆栈信息并不是实际的调用栈信息。当设置为“1”时，强制算子采用同步模式运行，这样能够打印正确的调用栈信息，从而更容易地调试和定位代码中的问题。设置为“0”时则会采用异步方式执行。

默认配置为0。

注意 ASCEND_LAUNCH_BLOCKING设置为“1”时，强制算子采用同步模式运行会导致性能下降。
ASCEND_LAUNCH_BLOCKING设置为“1”时，task_queue算子队列关闭，TASK_QUEUE_ENABLE设置不生效。
ASCEND_LAUNCH_BLOCKING设置为“0”时，会增加内存消耗，有导致OOM的风险。配置示例 export
ASCEND_LAUNCH_BLOCKING=1

2.3 其他常用优化操作

配置ACLNN_CACHE_LIMIT环境变量
ACLNN_CACHE_LIMIT用于配置单算子在执行API时，在Host侧缓存的算子信息条目个数。ACLNN_CACHE_LIMIT的单位为个，取值范围为[1,10000000]，默认值为10000。缓存的算子信息包含workspace大小、算子计算的执行器以及tiling信息等。动态shape场景下，如算子shape范围较大时，用户可通过配置此环境变量适当增加算子缓存条目，提升调度性能。当模型遇到下发瓶颈时（尤其是问题出现在一级流水时），可以尝试此配置。配置方法如下：

export ACLNN_CACHE_LIMIT=100000

说明增加算子信息缓存条目会增加Host内存开销，需根据实际情况适当调整。

其他优化手段还有很多，大家可以根据网络当前的性能状况，进行定向选择优化
torch_npu上其他优化手段

3 MinerU性能优化

3.1 通过torch_npu自带配置项优化

export CPU_AFFINITY_CONF=1  # 新版本可设置export CPU_AFFINITY_CONF=2export TASK_QUEUE_ENABLE=2export ACLNN_CACHE_LIMIT=100000

关键优化手段主要包括如上三个环境变量，800T A2单卡优化约30%，其他手段优化效果不明显，投入产出比不高。

3.2 其他优化

3.2.1 OM优化

端到端推理过程中，rapid_table 使用了onnx。默认执行方式会运行在CPU侧，可以使用OM 静态图（TorchAir） 加速推理。

如下onnx图中，在靠近结尾处（图中A处）包含了Loop算子，无法直接转换成OM格式：

在这里插入图片描述
有两种可行的方案：

提取Loop算子的子图，手动进行循环控制，工作量较大
对onnx图进行粗粒度切分，Loop算子包括子图直接运行在CPU上，工作量较小

当前选取方案2，在上图中的B处对onnx进行切分。打点分析可知onnx输入是静态shape，因此可以成静态图，使用om+onnx混合推理。

切分代码如下：

import onnx
onnx.utils.extract_model("slanet-plus-sim.onnx","slanet-plus-front.onnx", ["x"], ["hardswish_72.tmp_0"], check_model=False)onnx.utils.extract_model("slanet-plus-sim.onnx","slanet-plus-rear.onnx",["hardswish_72.tmp_0"], ["save_infer_model/scale_0.tmp_0", "save_infer_model/scale_1.tmp_0"], check_model=False)

onnx.utils.extract_model是用于子模型的提取，可以参考如下
Onnx子模型提取

模型切割后，在推理过程中，即可根据输入的参数进行判断，如下以shape进行判断，如

# 调用前先判断shape，避免其他shape的输入
def infer(self, x: List):if x[0].shape == (xx,xxx):front_out = self.static_infer(x)else:front_out = self.single_op_infer(x)

当遇到x的时候才进行静态shape的优化。

由于rapid_table时间占比有限，在100个pdf上进行的测试结果显示可以优化1~2%

3.2.2 高性能内存库

tcmalloc（即Thread-Caching Malloc）是一个通用的内存分配器，通过引入多层次缓存结构、减少互斥锁竞争、优化大对象处理流程等手段，在保证低延迟的同时也提升了整体性能表现。这对于需要频繁进行内存操作的应用来说尤为重要，尤其是在高并发场景下能够显著改善系统响应速度和服务质量。
内存分配高性能库
性能可以提升1%~2%