: NPU与GPU算子切分与流水线并行)
1. 背景与目标骁龙X2 Elite平台集成了三类计算单元Oryon CPU 负责通用逻辑与控制流Adreno GPU 擅长并行计算Hexagon NPU 专为张量运算优化。当前多数端侧 AI 部署仅利用 NPU 单一计算单元CPU 与 GPU 处于空闲状态算力浪费显著。异构计算的核心目标是使三类计算单元协同工作各自承担擅长的算子类型从而提升整体吞吐量并降低延迟。本系列分三篇展开介绍X2 Elite平台的异构计算的工程实践第一篇本篇NPU 与 GPU 算子切分与流水线并行第二篇CPU↔NPU 零拷贝与内存共享优化第三篇动态算子调度与能效最优策略本篇聚焦于算子切分策略将一个完整的模型推理流水线拆分至 NPU 与 GPU 两个计算单元使二者并行执行缩短端到端延迟。2. 异构计算架构基础2.1 三类计算单元的算子适配性不同计算单元对各类算子的执行效率存在显著差异选型时需依据算子特性匹配计算单元算子类型NPU 执行效率GPU 执行效率CPU 执行效率推荐计算单元矩阵乘GEMM优良差NPU卷积优良差NPU逐元素运算良优良GPU激活函数良优良GPUSoftmax良优良GPUReduce中优良GPU条件分支差差优CPU数据重组差中优CPU2.2 X2 Elite 异构计算软件栈X2 Elite 上调用三类计算单元的软件接口如下各层职责如下QNN SDK调用 Hexagon NPU 的官方接口支持 context binary 加载与推理OpenCL调用 Adreno GPU 的通用计算接口支持自定义 kernel 编写原生线程调用 Oryon CPU 执行控制逻辑与数据预处理统一内存CPU、NPU、GPU 共享 LPDDR5X无需跨设备数据搬运3. 算子切分策略3.1 切分原则算子切分需遵循三项原则算子匹配依据上表将算子分配至最适合的计算单元数据局部性相邻算子尽量分配至同一计算单元减少跨单元同步开销负载均衡各计算单元的执行时间尽量均衡避免长尾效应3.2 以 Transformer 模型为例的切分方案以一个 6 层 Transformer 模型为切分对象。每层包含注意力计算GEMM Softmax与前馈网络GEMM 激活函数切分方案如下子图算子组成分配计算单元依据QKV 投影GEMMNPUGEMM 为 NPU 最优算子注意力分数GEMM SoftmaxNPU GPUGEMM 在 NPUSoftmax 在 GPU注意力输出GEMMNPUGEMM 为 NPU 最优算子FFN 中间层GEMM GELUNPU GPUGEMM 在 NPUGELU 在 GPUFFN 输出GEMMNPUGEMM 为 NPU 最优算子3.3 切分实现importpyopenclasclimportqnnclassHeteroInference:def__init__(self,model_path):# NPU 上下文self.npuqnn.Model(model_path,backendlibQnnHtp.dll)self.npu.load()# GPU 上下文OpenCLself.ctxcl.create_some_context()self.queuecl.CommandQueue(self.ctx)self._compile_kernels()def_compile_kernels(self):# Softmax kernelsoftmax_kernel __kernel void softmax(__global const float* input, __global float* output, const int rows, const int cols) { int row get_global_id(0); if (row rows) return; float max_val -1e30f; for (int i 0; i cols; i) { max_val max(max_val, input[row * cols i]); } float sum 0.0f; for (int i 0; i cols; i) { output[row * cols i] exp(input[row * cols i] - max_val); sum output[row * cols i]; } for (int i 0; i cols; i) { output[row * cols i] / sum; } } # GELU kernelgelu_kernel __kernel void gelu(__global const float* input, __global float* output, const int size) { int i get_global_id(0); if (i size) return; float x input[i]; output[i] 0.5f * x * (1.0f tanh(0.7978845608f * (x 0.044715f * x * x * x))); } self.softmax_prgcl.Program(self.ctx,softmax_kernel).build()self.gelu_prgcl.Program(self.ctx,gelu_kernel).build()defforward(self,input_tensor):# 1. NPU 执行 QKV 投影qkvself.npu.execute({input:input_tensor})[output]# 2. NPU 执行注意力 GEMMattn_scoresself.npu.execute({input:qkv})[output]# 3. GPU 执行 Softmax与 NPU 下一层 GEMM 可并行softmax_outself._gpu_softmax(attn_scores)# 4. NPU 执行注意力输出投影attn_outputself.npu.execute({input:softmax_out})[output]# 5. NPU 执行 FFN 中间层 GEMMffn_interself.npu.execute({input:attn_output})[output]# 6. GPU 执行 GELUffn_actself._gpu_gelu(ffn_inter)# 7. NPU 执行 FFN 输出投影outputself.npu.execute({input:ffn_act})[output]returnoutputdef_gpu_softmax(self,tensor):# OpenCL 执行 Softmaxmfcl.mem_flags input_bufcl.Buffer(self.ctx,mf.READ_ONLY|mf.COPY_HOST_PTR,hostbuftensor)output_bufcl.Buffer(self.ctx,mf.WRITE_ONLY,tensor.nbytes)self.softmax_prg.softmax(self.queue,tensor.shape[0],None,input_buf,output_buf,np.int32(tensor.shape[0]),np.int32(tensor.shape[1]))resultnp.empty_like(tensor)cl.enqueue_copy(self.queue,result,output_buf)returnresult4. 流水线并行4.1 串行执行的瓶颈若按上文forward方法的实现NPU 与 GPU 算子交替执行某一时刻仅一个计算单元在工作。以单层 Transformer 为例串行执行的耗时约为 NPU 时间与 GPU 时间之和。4.2 流水线并行设计流水线并行的核心思路是当 NPU 执行第 N 层的 GEMM 时GPU 同时执行第 N-1 层的 Softmax。二者并行执行总耗时趋近于 NPU 与 GPU 中较长者。时序说明NPU 依次执行 Layer1-GEMM、Layer2-GEMM、Layer3-GEMMGPU 在 NPU 执行 Layer2-GEMM 的同时执行 Layer1-Softmax二者重叠执行总耗时从 6 个单位缩短至 4 个单位importthreadingimportqueueclassPipelineInference:def__init__(self,model_path,num_layers6):self.npuqnn.Model(model_path,backendlibQnnHtp.dll)self.npu.load()self.gpuOpenCLBackend()self.num_layersnum_layers self.task_queuequeue.Queue()defforward_pipeline(self,input_tensor):# 启动 GPU 消费线程gpu_threadthreading.Thread(targetself._gpu_worker)gpu_thread.start()# NPU 主循环逐层执行 GEMMcurrentinput_tensorforlayerinrange(self.num_layers):qkvself._npu_gemm(current,layer,qkv)attn_scoresself._npu_gemm(qkv,layer,attn)# 将 Softmax 任务交给 GPU异步self.task_queue.put((softmax,attn_scores,layer))# NPU 继续执行下一算子不等 GPUattn_outputself._npu_gemm(attn_scores,layer,proj)ffn_interself._npu_gemm(attn_output,layer,ffn1)# 将 GELU 任务交给 GPUself.task_queue.put((gelu,ffn_inter,layer))# 等待 GELU 结果仅此处需同步ffn_actself._wait_gpu_result(layer,gelu)currentself._npu_gemm(ffn_act,layer,ffn2)gpu_thread.join()returncurrentdef_gpu_worker(self):whileTrue:taskself.task_queue.get()iftaskisNone:breakop,data,layertaskifopsoftmax:resultself.gpu.softmax(data)self._store_result(layer,softmax,result)elifopgelu:resultself.gpu.gelu(data)self._store_result(layer,gelu,result)5. 性能实测5.1 算子级延迟对比对 Softmax 与 GELU 两个算子分别测量 NPU 与 GPU 的执行延迟算子输入形状NPU 延迟GPU 延迟CPU 延迟Softmax[12, 128, 128]0.42 ms0.18 ms0.85 msGELU[12, 128, 768]0.38 ms0.15 ms0.72 ms结论对于逐元素类算子Softmax、GELUGPU 延迟仅为 NPU 的 40%。将此类算子从 NPU 卸载至 GPU 具备明确的性能收益。5.2 端到端延迟对比对 6 层 Transformer 模型测量三种执行策略的端到端延迟执行策略延迟较纯 NPU 提升纯 NPU所有算子28.5 ms基准串行异构NPUGPU 交替31.2 ms-9.5%劣化流水线并行NPUGPU 重叠19.8 ms30.5%关键结论串行异构反而劣化因跨单元同步开销大于算子加速收益简单地将算子拆分至两个计算单元并不能带来性能提升流水线并行实现 30.5% 加速通过 NPU 与 GPU 重叠执行消除了同步等待端到端延迟从 28.5 ms 降至 19.8 ms5.3 功耗表现执行策略功耗说明纯 NPU7 W仅 NPU 工作流水线并行9 WNPU GPU 同时工作纯 CPU15 W全部在 CPU 执行流水线并行功耗较纯 NPU 增加 2 W但延迟降低 30.5%能效比吞吐量/功耗提升约 22%。6. 本篇小结本篇在 Snapdragon X2 Elite 上完成了 NPU 与 GPU 算子切分及流水线并行的工程实践主要成果如下基于算子特性将 Transformer 模型切分为 NPU 子图GEMM与 GPU 子图Softmax/GELU实现流水线并行NPU 与 GPU 重叠执行端到端延迟降低 30.5%验证了串行异构的陷阱——单纯拆分算子反而导致 9.5% 性能劣化流水线并行是关键当前实现中NPU 与 GPU 之间的数据传递仍通过显式拷贝完成存在额外开销。第二篇将引入统一内存架构下的零拷贝机制消除跨单元数据搬运成本。