Tensor shapes in a real-ish Transformer

Let’s make the dimensions realistic enough to understand GPU memory, but small enough to follow.

Assume a decoder-only Transformer with: - hidden size \(d_{model}=4096\) - 32 attention heads - head dimension \(d_{head}=128\)

Since \(32\times128=4096\).

Assume a prompt length \(T=2048\) and batch size \(B=1\).

1. The input tensor¶

After token embedding:

so:

Each of the 2048 prompt tokens carries a 4096-dimensional representation.

2. Generate Q, K and V¶

For ordinary multi-head attention:

and:

Each initially has shape:

Then reshape 4096 into 32 heads × 128 dimensions:

3. Attention creates a 2048×2048 matrix¶

For one head:

so:

With 32 heads:

That is about:

for one layer.

Naïve attention work therefore grows approximately as \(O(T^2)\).

4. But implementations need not store the whole matrix¶

FlashAttention computes attention in blocks so Q, K, and V tiles can stay in faster on-chip memory. The mathematics is unchanged; the memory traffic is dramatically reduced.

5. What survives prefill?¶

Old Q values are not needed later. Old K and V are.

For one layer:

Number of K values:

K+V:

At BF16/FP16 (2 bytes/value):

Across 32 layers:

for one 2048-token sequence in this simplified MHA example.

6. Decode shapes¶

For one new token:

and:

The cache contains roughly:

Attention for one head becomes:

This is a very different hardware workload from the large prefill matrix multiply.

7. Add batching¶

If 64 users decode simultaneously:

The server combines many tiny next-token operations into a healthier GPU workload. This is the core purpose of continuous batching.

8. Add GQA¶

Suppose Q still has 32 heads, but K/V only 8 heads.

Then:

KV memory falls to one quarter of the 32-KV-head case:

while retaining 32 independent query heads.

The formula to remember¶

A useful KV-cache approximation is:

The factor 2 is K+V.

This explains why model architecture directly affects serving capacity.

Why these dimensions matter to performance¶

Those tensor shapes are not merely bookkeeping. They tell you which dimensions grow when you add users, context, heads, or model width.

During prefill, the expensive score tensor is conceptually:

so doubling context length roughly quadruples the number of Q/K score interactions. During decode, by contrast, the new query has only one token position:

and compares against a cache of length \(T\). That is why the attention portion of a single decode step grows approximately linearly with current context length.

The large linear projections are different again. A projection such as:

uses a weight matrix of roughly \(4096\times4096\) in this example. In prefill, many rows of \(X\) reuse those same weights at once. During decode there may be only one row per sequence, which is why batching helps so much.

A practical serving calculation¶

Suppose an inference GPU has 80 GB of usable HBM and the quantized model weights consume 45 GB. That leaves roughly 35 GB for KV caches, temporary buffers, and runtime overhead.

If a long-context request consumes 500 MB of KV state, the naïve upper bound is only:

simultaneous sequences, before allowing for fragmentation and other buffers.

If GQA cuts that cache to 125 MB, the capacity picture changes radically.

This is why a seemingly small architectural choice such as the number of KV heads can translate directly into serving economics.

The shape-reading habit¶

Whenever you encounter a new inference technique, ask:

1. Which tensor is being made smaller?
2. Which dimension is being parallelized?
3. Which values must persist between tokens?
4. Which values can be recomputed or discarded?

That habit makes many systems papers much easier to reason about.