Overlap Scheduling吞吐优化【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub一、 背景介绍当前AI模型部署设备通常采用CPU与算力卡组成的异构架构。在AI推理业务中CPU主要负责任务调度算力卡负责执行大量计算任务。由于算力卡成本远高于CPU其计算执行又依赖CPU的任务下发一旦CPU调度不及时算力卡便会陷入空闲等待造成高成本算力资源被低成本CPU拖累的局面。Overlap Scheduling是解决这一问题的通用优化手段——它将CPU的任务调度过程与算力设备的计算执行过程并行推进从而隐藏CPU调度延迟提升算力设备利用率与系统吞吐量。本文将详细介绍Overlap Scheduling在昇腾设备上的实现原理与应用方式。二、Overlap Scheduling实现原理在未使能Overlap Scheduling时CPU 与 NPU 以严格的串行方式协作CPU完成一批任务的调度下发后NPU 才开始执行计算NPU 执行完毕后CPU 再调度下一批任务。 Overlap Scheduling 的核心思路是提前下发在 NPU 执行当前批次计算任务的同时CPU 并行完成下一批次的任务调度与下发使两者的执行时间窗口相互重叠。具体实现上通常借助任务队列与异步下发机制—— CPU 将调度好的任务预先写入队列NPU 从队列中持续取任务执行从而解耦两者的时序依赖。上图便给出了使能Overlap Scheduling前后的示意图可以看到使能Overlap Scheduling前NPU存在大量空泡等待时间在使能Overlap Scheduling后消除了NPU的等待空泡整体耗时缩短提升了吞吐量。三、Overlap Scheduling昇腾使用实例SGLang的Overlap Scheduling实现介绍3.1 背景介绍SGLang是一个高性能的大语言模型LLM和视觉-语言模型VLM服务框架支持从单卡到大规模分布式集群的多种部署场景提供低延迟、高吞吐的推理服务。SGLang已适配昇腾NPU设备并通过Overlap Scheduling方案显著提升了大模型在NPU上的吞吐性能。3.2 实现架构SGLang采用双线程异步架构实现Overlap Scheduling通过任务队列解耦请求处理和NPU执行线程1调度线程接收推理请求组装成batch后提交到任务队列input_queue从输出队列output_queue获取推理结果并返回线程2执行线程从任务队列获取batch任务下发到NPU执行并将结果放入输出队列3.3 执行流程步骤线程1调度线程线程2执行线程说明1接收请求组装batch1和batch2提交到任务队列等待batch1结果-初始化阶段预先提交2个batch2等待中执行batch1将结果放入输出队列开始执行batch2/3获取batch1结果提交batch3等待batch2结果执行batch2中CPU处理结果与NPU执行并行4等待中完成batch2将结果放入输出队列开始执行batch3/5获取batch2结果提交batch4等待batch3结果...完成batch3开始batch4...循环往复关键设计预提交机制启动时预先提交2个batch避免NPU后续需要等待任务下发。流水线并行线程1处理结果和组装新batch的同时线程2下发任务到NPU上执行计算队列解耦通过任务队列和输出队列实现线程间的异步通信避免相互阻塞通过这种设计SGLang实现了CPU侧的请求处理与NPU侧的模型推理的充分重叠最大化设备利用率。 完成后方案的调度耗时示意图如下由于任务调度以及结果获取步骤的耗时通常都显著低于NPU任务执行耗时因此在首次下发两个任务之后每次下发单个任务的循环中可以保持NPU始终有任务执行并且首个任务的结果输出也不会过于滞后。 完整的代码调用栈示意图如下开发者们也可以访问SGLang社区查看具体的代码实现https://github.com/sgl-project/sglang3.4 模型效果实测实测LongCat-Flash模型560B参数量初期优化时在batch 8输入长度4k的场景下使能Overlap Scheduling技术前后吞吐提升1.7x倍 | batch | 输入长度 | overlap | 端到端时间ms | TPS | |-------|----------|---------|-----------------|--------| | 8 | 4k | 否 | 79 | 101.27 | | 8 | 4k | 是 | 45 | 177.78 |四、Overlap Scheduling常见问题及解决方案4.1 问题描述在昇腾设备上实现Overlap Scheduling后如果模型中存在Tiling依赖输入值的算子会导致任务下发阻塞进而影响CPU和NPU的并行执行降低整体吞吐性能。Tiling机制说明Tiling是指将算子的输入数据切分到NPU的多个核上并行执行的逻辑。Tiling的计算通常在CPU侧完成且大多数情况下只依赖输入的Shape信息。但部分特殊算子的Tiling需要依赖具体的输入值。 典型例子 ReduceSum算子用于沿指定维度对Tensor的所有元素求和包含两个输入axes指定求和的维度x待求和的Tensor该算子的Tiling必须知道axes的具体值才能根据x的Shape确定该维度的元素数量进而完成切分策略的计算。性能瓶颈分析当算子的Tiling依赖具体输入值时会因以下流程导致吞吐下降数据传输算子的输入数据存储在NPU侧而Tiling计算在CPU侧执行因此需要将依赖的输入数据从NPU拷贝到CPU同步阻塞CPU侧调用aclrtSynchronizeStream接口同步等待数据拷贝完成该同步操作会阻塞当前线程调度停滞由于调度线程被阻塞后续的任务无法及时下发破坏了CPU和NPU的执行重叠导致整体吞吐下降被打断的流水图如下图中可以看到由于插入多个aclrtSynchronizeStream调度线程任务下发阻塞进而NPU的空闲等待时间也增长。4.2 问题解决方案问题的核心矛盾在于Tiling计算在CPU侧执行但依赖的数据存储在NPU侧。因此最直接的优化思路是将Tiling计算迁移到NPU侧执行避免跨设备的数据同步。目前有以下两种实现方式方案一Kernel内部计算Tiling在算子的Kernel实现中直接完成Tiling计算。由于Tiling主要涉及标量运算可以利用NPU的标量计算单元在Kernel内部完成。例如对于ReduceSum算子可在Kernel实现中通过GetValue接口获取axes的具体值然后基于该值和输入Shape计算数据切分策略整个过程在NPU侧完成无需CPU参与。方案二Tiling下沉到AICPU利用昇腾NPU的Tiling下沉功能将Tiling计算任务调度到NPU侧的AICPU执行避免CPU与NPU之间的数据传输和同步。 使用详情参考文档 https://www.hiascend.com/document/detail/zh/canncommercial/850/opdevg/Ascendcopdevg/atlas_ascendc_10_00014.html五.总结Overlap Scheduling是提升昇腾设备吞吐性能的重要优化方式开发者们可以尝试将此方案应用到更多的AI框架中不仅能从实践中更深入地掌握这一技术特性也有机会将自己的代码贡献到优秀的开源项目中成为自身能力的优秀证明。【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考