Multi Query Attention Is All You Need

Multi-Query Attention is All You Need

GLM 5.2 is live! Opus-level intelligence at open-source rates. Pay per token on serverless. Try it today.

Blog

Multi Query Attention Is All You Need Multi-Query Attention is All You Need

PUBLISHED 7/12/2023

Table of Contents

Performance Results

Scaling: MQA + Model Parallel

Try Out Models with MQA on the Fireworks Gen AI Platform

Appendix: Performance Bug in Original Falcon Implementation

Table of Contents

This post explores how a nascent modeling technique called Multi-Query Attention (MQA) significantly improves machine performance and efficiency for language inference tasks such as summarization, question answering, and retrieval-augmented generation. By using MQA-based efficiency techniques, users can get 11x better throughput and 30% lower latency on inference. Models that use Multi-Query Attention include LLaMA-v2 and Falcon . Further, we explore a technique for executing MQA in a distributed fashion to further improve latency. Finally, we show how the Fireworks Gen AI Platform allows you to tune LLMs to solve your business tasks and efficiently serve these models using the described techniques. Note that literature released after this blog post refers to Multi-Query Attention with multiple KV heads as “Group-Query Attention”. Efficient LLM Inference

Large Language Models (LLMs) based on the Transformer architecture have emerged as an effective technique for language tasks, including summarization, question answering (Q&A), and retrieval-augmented generation. However, using these models comes at a very high computational cost, and their execution is primarily done via compute accelerators like NVIDIA GPUs. Input and outputs to LLMs are represented as sequences of tokens (e.g. words). Training or fine-tuning LLMs that can handle long sequences (i.e. that have a long context window ) is an actively evolving field. Most OSS LLM base models are pre-trained with a 2K context window. In more and more use cases like document summarization or context-based question answering, the sequence length processed by the LLM can grow quite large–into thousands to tens of thousands of tokens. In the future, we believe long sequence lengths will be the new norm for most LLM use cases. But long sequences also have significant efficiency implications for the cost of inference. System performance for inference can be improved without changing the model through several techniques, including: • Saving computed state between iterations of the inference process (KV-caching) • Batching multiple sequences together during inference to reuse computational resources (batching) and, as an extension, continuously batching concurrent requests (e.g. Yu et al. ) • Memory allocation strategies to reduce memory fragmentation and maximize batch size. (e.g. VLLM )

However, the most effective way to improve inference performance is to co-design the model architecture and the system architecture. In this article, we highlight one such joint technique, Multi-Query Attention (MQA) , which dramatically reduces both memory space and memory bandwidth needed for inference computation. Space savings are proportional to the number of tokens, so it’s particularly beneficial for long sequences. Optimizing for MQA can lead to 11x better throughput and 30% lower latency in our benchmarks compared to the best openly available baselines without MQA. Background - Multi-Head Attention

Much of LLMs' language expressivity comes from mixing context across sequences via the attention operation. Vaswani et al . propose Multi-Head Attention as the following mathematical operation:

Multi-Head Attention Definition Here h represents the number of “heads” in the operation, S and L represent input and output sequence lengths (respectively), and d_k represents the hidden dimensionality of the model architecture. Equivalently in PyTorch code, we can write (with an extra batch dimension N): 1 2 3 4 5 6 7 8 9 10 11 Q = torch . randn ( N , h , S , d_k ) K = torch . randn ( N , h , L , d_k ) V = torch . randn ( N , h , L , d_k )

#

logits = torch . matmul ( Q , K . transpose ( 2 , 3 ) ) # Output shape [N, h, S, L] softmax_out = torch . softmax ( logits / math . sqrt ( d_k ) , dim = - 1 ) # Output shape [N, h, S, L] attn_out = torch . matmul ( softmax_out , V ) # Output shape [N, h, S, d\_k\]

Note that we have two sequence lengths: one that applies to our Q value and one that applies to both K and V values. During inference, we typically use incremental generation , where we progressively feed values into the network a single token at a time (i.e. S = 1) and compute K and V across the tokens seen so far (i.e. L grows as generation proceeds). As a result, K and V grow progressively as the output sequence is generated, and a common optimization technique is to use a mutable KV-cache across iterations. The inner loop of multi-head attention then looks like this: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20

Cached K and V values across iterations

K = torch . randn ( N , h , . . . , d_k ) V = torch . randn ( N , h , . . . , d_k )

Single-step QKV values computed during sequence generation

Q_incr = torch . randn ( N , h , 1 , d_k ) K_incr = torch . randn ( N , h , 1 , d_k ) V_incr = torch . randn ( N , h , 1 , d_k )

#

Update KV-cache

K = torch . cat ( [ K , K_incr ] , dim = - 2 ) V = torch . cat ( [ V , V_incr ] , dim = - 2 )

Compute attention (L is sequence length so far)

logits = torch . matmul ( Q_incr , K . transpose ( 2 , 3 ) ) # Output shape [N, h, 1, L] softmax_out = torch . softmax ( logits / math . sqrt ( d_k ) , dim = - 1 ) # Output shape [N, h, 1, L] attn_out = torch . matmul ( softmax_out , V ) # Output shape [N, h, 1, d_k]

Multi-Query Attention

Shazeer (2019) proposed a refinement to the Multi-Head Attention (MHA) algorithm called Multi-Query Attention (MQA), which improves machine efficiency of attention while incurring minimal accuracy degradation. The idea is simple: remove (or otherwise greatly reduce) the heads dimension h from the K and V values. Intuitively, we can say that in multi- head attention, the entire attention computation is replicated h times, whereas in multi- query attention, each “head” of the query value Q has the same K and V transformation applied to it. The incremental generation case looks like this: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22

Cached K and V values across iterations

K = torch . randn ( N , . . . , d_k )...