FlashEVA: Accelerating LLM inference via Efficient Attention

TL;DR

FlashEVA introduces an efficient attention mechanism for Transformers, significantly reducing memory usage and increasing inference throughput while maintaining performance across multiple tasks.

Contribution

It presents FlashEVA, a novel implementation of EVA attention, enabling faster and more memory-efficient Transformer inference with minimal fine-tuning data.

Findings

Up to 6.7x higher throughput during inference

5x reduction in peak GPU memory usage

Effective across various downstream tasks

Abstract

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 FlashEVA, an efficient implementation of EVA (Efficient Attention via Control Variates), and demonstrate how to finetune transformers to adapt to FlashEVA attention. Our method enables fine-tuning of Transformer models with as few as 1.5B tokens while preserving effectiveness across various downstream tasks. Notably, FlashEVA 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.