Decoder-Only Transformers & KV Cache

While the original 2017 Transformer was an encoder-decoder architecture designed for machine translation, modern foundation Large Language Models (LLMs) like GPT-4, LLaMA-3, and Mistral are strictly Decoder-Only.

Decoder-only architectures optimize auto-regressive next-token prediction by combining Causal Self-Attention, RMSNorm, Rotary Position Embeddings (RoPE), and inference-time Key-Value (KV) Caching.

1. Causal Masking & Pre-LN Transformer Blocks#

Auto-regressive language modeling factorizes the joint probability of token sequence $X = (x_1, \dots, x_T)$ into conditional next-token probabilities:

To prevent token $i$ from attending to future tokens $j > i$ during parallel training, an upper-triangular causal mask $M \in \mathbb{R}^{T \times T}$ is added before softmax:

Because $\exp(-\infty) = 0$, attention weights for all future positions evaluate strictly to zero.

Root Mean Square Normalization (RMSNorm)#

Zhang and Sennrich (2019) discovered that the scaling property of LayerNorm dominates its stabilizing effect, whereas mean-centering is computationally redundant. RMSNorm discards mean calculation:

Where $\gamma \in \mathbb{R}^d$ is a learnable gain parameter. RMSNorm reduces normalization runtime by 7% to 10% across deep architectures.

2. Rotary Position Embeddings (RoPE)#

Su et al. (2021) introduced Rotary Position Embeddings (RoPE) to encode relative positional geometry directly into Query and Key inner products.

We seek an embedding function $f(x, m)$ such that the inner product between query $q$ at position $m$ and key $k$ at position $n$ depends purely on relative offset $m - n$:

Derivation in the Complex 2D Plane#

Treating a 2D coordinate pair $(q_1, q_2)$ as a complex number $z = q_1 + i q_2 \in \mathbb{C}$:

Multiplying by $e^{i m \theta} = \cos(m \theta) + i \sin(m \theta)$ rotates the vector by angle $m \theta$:

Now, compute the inner product of rotated query and rotated key:

Because rotation matrices are orthogonal and commutative in 2D:

Therefore:

The inner product is mathematically determined solely by relative distance $m - n$, enabling natural length generalization.

3. KV Cache Architecture and Arithmetic Intensity#

During inference generation, tokens are emitted auto-regressively one token at a time.

The KV Cache Principle#

At timestep $t$, historical tokens $x_1, \dots, x_{t-1}$ have already been processed. Their Key and Value projections do not change:

1. Cache previous key tensors $K_{\text{cache}} \in \mathbb{R}^{(t-1) \times d_k}$ and value tensors $V_{\text{cache}} \in \mathbb{R}^{(t-1) \times d_v}$ in GPU memory.
2. At step $t$, compute Query $q_t \in \mathbb{R}^{1 \times d_k}$, new Key $k_t \in \mathbb{R}^{1 \times d_k}$, and new Value $v_t \in \mathbb{R}^{1 \times d_v}$ for the current token only.
3. Append $k_t$ to $K_{\text{cache}}$ and $v_t$ to $V_{\text{cache}}$.
4. Compute attention vector:

KV Cache Memory Footprint Formula#

For a model with $L$ layers, $H$ attention heads, head dimension $d_k$, context length $T$, batch size $B$, and 16-bit floating point precision ($2$ bytes):

For LLaMA-3-70B ($L = 80, H = 64, d_k = 128$) with context length $T = 8192$ and batch size $B = 4$:

This proves why LLM token generation is memory-bandwidth bound (high memory transfer, low arithmetic compute) rather than compute-bound.

4. From-Scratch Implementation: Zero Dependencies#

The following implementations construct an Autoregressive Causal Attention Layer with an Explicit KV Cache from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import math
from typing import List, Tuple, Optional

class CausalKVCacheAttention:
    def __init__(self, d_model: int = 4, d_k: int = 2):
        self.d_model = d_model
        self.d_k = d_k
        self.scale = 1.0 / math.sqrt(d_k)

        # Cache storage for keys and values across timesteps
        self.k_cache: List[List[float]] = []
        self.v_cache: List[List[float]] = []

    def reset_cache(self):
        self.k_cache.clear()
        self.v_cache.clear()

    @staticmethod
    def _dot(a: List[float], b: List[float]) -> float:
        return sum(x * y for x, y in zip(a, b))

    @staticmethod
    def _softmax(logits: List[float]) -> List[float]:
        max_val = max(logits)
        exps = [math.exp(x - max_val) for x in logits]
        sum_exp = sum(exps)
        return [e / sum_exp for e in exps]

    def forward_step(
        self,
        q_t: List[float],
        k_t: List[float],
        v_t: List[float]
    ) -> List[float]:
        # 1. Append current token key and value to cache
        self.k_cache.append(k_t)
        self.v_cache.append(v_t)

        seq_len = len(self.k_cache)

        # 2. Compute attention scores: q_t * K_cache^T / sqrt(d_k)
        scores = [
            self._dot(q_t, self.k_cache[pos]) * self.scale
            for pos in range(seq_len)
        ]

        # 3. Softmax over all past positions (causal by construction)
        weights = self._softmax(scores)

        # 4. Multiply weights by V_cache
        output = [0.0] * len(v_t)
        for pos in range(seq_len):
            for col in range(len(v_t)):
                output[col] += weights[pos] * self.v_cache[pos][col]

        return output

if __name__ == "__main__":
    attn = CausalKVCacheAttention(d_model=4, d_k=2)

    # Simulate 3 autoregressive generation steps
    steps = [
        ([1.0, 0.5], [1.0, 0.5], [10.0, 1.0]),  # Token 1
        ([0.8, 1.2], [0.8, 1.2], [20.0, 2.0]),  # Token 2
        ([0.2, 0.9], [0.2, 0.9], [30.0, 3.0]),  # Token 3
    ]

    print("Autoregressive Generation with KV Cache:")
    for step_idx, (q, k, v) in enumerate(steps):
        out_t = attn.forward_step(q, k, v)
        print(f"Step {step_idx + 1} (Cache Size: {len(attn.k_cache)}):")
        print("  Emitted Representation:", [round(x, 4) for x in out_t])

5. Active Knowledge Checks#