FlashFormer: Whole-Model Kernels for Efficient Low-Batch Inference

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Nrusimha, Aniruddha, Brandon, William, Mishra, Mayank, Shen, Yikang, Panda, Rameswar, Ragan-Kelley, Jonathan, Kim, Yoon
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908692292567040
author Nrusimha, Aniruddha
Brandon, William
Mishra, Mayank
Shen, Yikang
Panda, Rameswar
Ragan-Kelley, Jonathan
Kim, Yoon
author_facet Nrusimha, Aniruddha
Brandon, William
Mishra, Mayank
Shen, Yikang
Panda, Rameswar
Ragan-Kelley, Jonathan
Kim, Yoon
contents The size and compute characteristics of modern large language models have led to an increased interest in developing specialized kernels tailored for particular training and inference workloads. Existing kernels primarily optimize for compute utilization, targeting the large-batch training and inference settings. However, low-batch inference, where memory bandwidth and kernel launch overheads are significant factors, remains important for many applications of interest such as in edge deployment and latency-sensitive applications. This paper describes FlashFormer, which fuses the entire transformer forward pass into a single kernel for accelerating low-batch inference of large language models. Across various model sizes and quantizations settings, FlashFormer achieves nontrivial speedups compared to existing inference kernels.
format Preprint
id arxiv_https___arxiv_org_abs_2505_22758
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle FlashFormer: Whole-Model Kernels for Efficient Low-Batch Inference
Nrusimha, Aniruddha
Brandon, William
Mishra, Mayank
Shen, Yikang
Panda, Rameswar
Ragan-Kelley, Jonathan
Kim, Yoon
Machine Learning
Computation and Language
The size and compute characteristics of modern large language models have led to an increased interest in developing specialized kernels tailored for particular training and inference workloads. Existing kernels primarily optimize for compute utilization, targeting the large-batch training and inference settings. However, low-batch inference, where memory bandwidth and kernel launch overheads are significant factors, remains important for many applications of interest such as in edge deployment and latency-sensitive applications. This paper describes FlashFormer, which fuses the entire transformer forward pass into a single kernel for accelerating low-batch inference of large language models. Across various model sizes and quantizations settings, FlashFormer achieves nontrivial speedups compared to existing inference kernels.
title FlashFormer: Whole-Model Kernels for Efficient Low-Batch Inference
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2505.22758