Thread by @amanrsanger

With a 256K token prompt, a 7b model can generate tokens as quickly as codellama-7b with an 8K prompt.

How? The model must use multi-query attention.

Here's why...
(1/10)

For large context windows and large batch inference, the bottleneck for generation speed is bottlenecked by KV cache size.

To illustrate this, let’s look at our 8K context window with multi-head attention at a batch size of 16.
(2/10)

To generate each token, we need to spend 7e9*2*16 = 224 GFLOPs. [1]

We won’t be compute bound as a single A100 has 300TFLOPs, meaning it could sustain >1000 tokens/s without the memory bottleneck.
(3/10)

We’ll also need to read the model weights from memory, costing us 7e9*2 = 14GB of memory bandwidth [2]

Finally, we need to read the KV cache on each token. This is:

128 attn dim * 32 kv heads *32 layers * 2 key + value *16 batch size *8K tokens * 2 bytes = 67GB!
(4/10)

Now, how fast should we be decoding in theory? Well, using tp=2 on two A100s, we’re reading 81GB on each token or 40GB/A100.

Assuming we can achieve 70% of the 2TB/s bandwidth, that’s 35 tok/s.
(5/10)

Let’s swap out codellama for our multi-query model. Notice that we can directly trade context length for heads in our KV cache memory.

We multiply our context window by 32, and divide the num_heads by 32 giving us the exact same size.
(6/10)

This means for our 256k token prompt, we're decoding at 35 tok/s on a reasonable batch size
(7/10)

Now, prefill is a whole different story, and actually getting the model to pay attention to 256K tokens from training is another can of worms...
(8/10)

But MQA models give such massive inference wins for long context that they may be worth it despite slightly worse scaling perf than MHA or GQA.

Plus, we’re already working with an MQA model that beats codellama 7b on code evals.
(9/10)

Finally, 256K is not the limit of how far you can take things with vanilla transformers for reasonable inference perf.

There are a few more long context tricks that, in theory, should be able to preserve that same perf up to 1M+ ;)

(10/10)

[1] Simple model flops calc
We can also account for attention flops, but notice they'll also be insignificant compared to KV cache size:

For each of the attention heads and sequences in the batch, we’re doing a 1 x 128 dim (query) times a 128 x 256k (keys) matmul.

That’s 2* 128*256k FLOPs, then * 16 bs * 32 kv heads * 32 layers.

Then the attention weights are a 1 x 256K vector, multiplying a 256K x 128 vector (the values), giving 2 * 256K * 128 FLOPs, again * 16bs * 32 kv heads * 32 layers. This is 2 TFLOPs.

On our 2 GPUs, we should be able to do 300TPS (probs closer to 200TPS) if not bottlenecked by compute

But we're still bottnecked by mem bw, which keeps us at 35 TPS

* Actually we need to spend 2TFLOPs/token when accounting for attention, but on our 2 A100s, that would be 200-300TPS, which the memory bw bottleneck prevents us from hitting

[2] - We need to read model weights to generate each token, and we're assuming bf16 or fp16, hence the term there