Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance

TL;DR

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, 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.

Abstract

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.

Figures (13)

• Figure 1: Token generation performance of Llama3.1-8B on DGX H100 and SN40L. Synchronization overheads at kernel call boundaries in DGX H100 inhibit utilizing HBM bandwidth efficiently. Kernel looping (two orange bars to the right) mitigates these overheads on SN40L, and have been deployed in production.
• Figure 2: Llama3.1-8B with 32 decoder layers, and the operators within one decoder layer. Labels $K1$ -- $K10$ (green) are kernel calls on the DGX H100, obtained by profiling the execution of Llama3.1-8B using TensorRT-LLM tensorrt-llm with the NVIDIA Nsight profiler nsight-systems. Label $K0$ (orange) corresponds to one kernel call on SN40L for the entire decoder layer. Collective communication operations like allreduce do not execute asynchronously with other operators on the DGX H100. On SN40L, allreduce is pipelined and overlapped with other operators.
• Figure 3: Pseudocode showing kernel call schedule to generate one output token on (i) DGX H100 with TensorRT-LLM tensorrt-llm, (ii) SN40L-8 with layer fusion, and (iii) this work. On DGX H100, a single output token requires 323 kernel calls, while SN40L-8 needs 35 kernel calls. In this work, we reduce the number of kernel calls per output token on SN40L to 4.
• Figure 4: Time-per-Output-Token (TPOT) breakdown for Llama3.1-8B on 2, 4, and 8 H100s. GEMMs scale better than other operators (although not ideally). Allreduce and other operators demonstrate poor scaling, thus accounting for a larger fraction of the overall token generation time.
• Figure 5: Time-per-Output-Token (TPOT) breakdown for Llama3.1-8B on SN40L-8. Synchronization overheads washes out the performance upside with a single kernel decoder. This paper significantly reduces synchronization overheads to give close to peak TPOT.