通过循环内核优化,大幅减少推理时的同步开销,提升生成速度。
Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance
- 将多个层调用合并为单个带流水循环的内核,消除同步开销。
- 在SN40L上实现最高2.2倍解码加速,多芯片扩展达2.5倍。
- 适用于大规模语言模型推理,已投入商业云服务生产。
令牌生成速度对下一代AI推理应用至关重要。由于内核边界处的同步开销,GPU在令牌生成阶段仅能利用其峰值内存带宽的21%。尽管近期数据流架构通过将解码器层激进融合为单一内核缓解了此类开销,但层间仍存在同步惩罚导致性能损失。本文提出内核循环(Kernel Looping),一种专门的全局优化技术,利用现代数据流架构支持的层级融合与语言模型固有的重复层结构相结合的优化机会。该技术通过将连续的相同内核调用转换为含流水外循环的单个内核调用,彻底消除层间同步成本。我们在SambaNova SN40L可重构数据流单元(RDU)——一款商用AI加速器上评估了该方法。实验表明,内核循环使一系列主流开源模型的解码阶段提速最高达2.2倍;支持多芯片扩展,速度提升最高达2.5倍;在8和16个SN40L Socket上实现超过90%的峰值性能,并相较DGX H100最高提速3.7倍。内核循环及所评测模型已在商业AI推理云中投入生产部署。
原文摘要 · Abstract (English)
Token generation speed is critical to power the next wave of AI inference applications. GPUs significantly underperform during token generation due to synchronization overheads at kernel boundaries, utilizing only 21% of their peak memory bandwidth. While recent dataflow architectures mitigate these overheads by enabling aggressive fusion of decoder layers into a single kernel, they too leave performance on the table due to synchronization penalties at layer boundaries. This paper presents kernel looping, a specialized global optimization technique which exploits an optimization opportunity brought by combining the unique layer-level fusion possible in modern dataflow architectures with the repeated layer structure found in language models. Kernel looping eliminates synchronization costs between consecutive calls to the same kernel by transforming these calls into a single call to a modified kernel containing a pipelined outer loop. We evaluate kernel looping on the SambaNova SN40L Reconfigurable Dataflow Unit (RDU), a commercial dataflow accelerator for AI. Experiments demonstrate that kernel looping speeds up the decode phase of a wide array of powerful open-source models by up to 2.2$\times$ on SN40L. Kernel looping allows scaling of decode performance over multiple SN40L sockets, achieving speedups of up to 2.5$\times$. Finally, kernel looping enables SN40L to achieve over 90% of peak performance on 8 and 16 sockets and achieve a speedup of up to 3.7$\times$ over DGX H100. Kernel looping, as well as the models evaluated in this paper, are deployed in production in a commercial AI inference cloud.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。