So, starting from the top, this conversation will be really confusing if you don't have an intuitive handle on how big the various tensors in an LLM typically are.

In terms of components, the weights are generally matrices. This means an N*M sized tensor, and you can expect it to scale geometrically as you increase the size of the model.

Typically, 90+% of the weight matrices are in the FFN. Because matrix multiplication is memory bound, the FFN is the cause of LLMs being memory bound to start with. It's just a big block of matmuls.

The hidden state at any given point is only a vector (ie: size "M"). This means that loading weights to cache >> loading hidden state to cache.

Attention, on the other hand, has the minority of the weights in the model. What it does, is it compares every token against every other token. That's why it has an O(N\^2) scaling characteristic by default. First it multiplies the Key and Query weights by every input to build the K and Q matrices, and then it multiplies the transpose of one of those against the other, to give you the Attention matrix. Then it Softmaxes that matrix, and dot products it against the Value matrix and possibly a linear output transform.

If you were paying very close attention (no pun intended) you might notice there's a really interesting behavior here: What if you kept the old Keys (and even the old Value matrix)?

Well, an interesting consequence, is that you only add one new entry in the Key matrix per new token, and the Attention matrix itself is basically unchanged (just one new row and column per token). If you do a few funky things with this, you get KV caching, which makes the time to generate a new token linear. There's also a way of taking Attention's cost from quadratic to linear, called "Flash Attention" which works in a different way (it's more algorithmic), but long story short, they let you manage the computational and memory cost of Attention somewhat.

Anyway, the interesting thing about Attention, is that its cost is more like a CNN when you get a really big context window. The reason is that you use the same Q and K and V weights on every single token, so you're re-using the same weights in memory multiple times.

You can do this with the FFN, too, by batching (which is what I was referring to earlier), meaning that you calculate multiple tokens at the same time (for different contexts), and because loading the weights is so expensive, and loading a hidden state is so cheap comparatively, you end up with the same latency for around 4-8 requests as sending a single one.

This doesn't go forever (there are memory bottlenecks, real world bottlenecks, CUDA overhead, etc etc), but the end result is that batching gives you more total T/s.