FlashEVA: Accelerating LLM inference via Efficient Attention

Transformer models have revolutionized natural language processing, achieving state-of-the-art performance and demonstrating remarkable scalability. However, their memory demands, particularly due to maintaining full context in memory, pose significant challenges for inference. In this paper, we present FlashEV A, an efficient implementation of EV A (Efficient Attention via Control V ariates), and demonstrate how to finetune transformers to adapt to FlashEV A attention. Our method enables fine-tuning of Transformer models with as few as 1.5B tokens while preserving effectiveness across various downstream tasks. Notably, FlashEV A achieves up to 6.7x higher throughput and 5x lower peak GPU memory usage during inference compared to standard Transformer implementations. Despite these improvements, we observe limitations in retrieval-focused tasks. Our implementation offers control over the trade-off between throughput and accuracy through adjustable hyperparameters, providing flexibility for diverse use cases. This work represents a significant step towards more efficient and adaptable Transformer-based models for inference. Transformer models have become ubiquitous in the field of natural language processing, achieving state-of-the-art performance across a wide range of tasks (Dosovitskiy et al., 2021; Wang et al., 2019; Radford et al., 2019; Dong et al., 2018). Their success can be attributed to their ability to scale effectively and the possibility of parallel training, which has led to significant improvements in model capabilities (Kaplan et al., 2020; Gadre et al., 2024).