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