Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Koeplinger, David, Gandhi, Darshan, Nandkar, Pushkar, Sheeley, Nathan, Musaddiq, Matheen, Zhang, Leon, Goodbar, Reid, Shaffer, Matthew, Wang, Han, Wang, Angela, Wang, Mingran, Prabhakar, Raghu
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866929569779417088
author Koeplinger, David
Gandhi, Darshan
Nandkar, Pushkar
Sheeley, Nathan
Musaddiq, Matheen
Zhang, Leon
Goodbar, Reid
Shaffer, Matthew
Wang, Han
Wang, Angela
Wang, Mingran
Prabhakar, Raghu
author_facet Koeplinger, David
Gandhi, Darshan
Nandkar, Pushkar
Sheeley, Nathan
Musaddiq, Matheen
Zhang, Leon
Goodbar, Reid
Shaffer, Matthew
Wang, Han
Wang, Angela
Wang, Mingran
Prabhakar, Raghu
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_23668
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance
Koeplinger, David
Gandhi, Darshan
Nandkar, Pushkar
Sheeley, Nathan
Musaddiq, Matheen
Zhang, Leon
Goodbar, Reid
Shaffer, Matthew
Wang, Han
Wang, Angela
Wang, Mingran
Prabhakar, Raghu
Computation and Language
Artificial Intelligence
Hardware Architecture
D.3.4; C.1.3
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.
title Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance
topic Computation and Language
Artificial Intelligence
Hardware Architecture
D.3.4; C.1.3
url https://arxiv.org/abs/2410.23668