Llama 3.1 8B is the closest thing the field has to a reference decoder: 32 layers, hidden size 4096, grouped-query attention with 32 query heads over 8 KV heads, SwiGLU, RMSNorm, RoPE at theta 500,000, vocabulary 128,256. That makes it the right model to attach the one piece of inference machinery everyone uses and few people write by hand: the KV cache. Mine is the only model in the repository that carries one, and it is about twenty lines.

Grouped queries without repeat_kv

The usual implementation repeats each KV head four times to match the 32 query heads. I skip the copy: reshape queries to five dimensions, grouped by KV head, leave keys and values with a singleton group axis, and let broadcasting do the rest. No repeat_interleave, no materialized copies of K and V, and the grouping is visible in the shapes.

repeat_factor = num_attention_heads // num_key_value_heads    # 32 // 8 = 4
q = q.reshape(B, num_key_value_heads, repeat_factor, T, head_size)
k = k.reshape(B, num_key_value_heads, 1, T, head_size)
v = v.reshape(B, num_key_value_heads, 1, T, head_size)
scores = q @ k.transpose(-1, -2)    # B, n_kv, 4, T, T via broadcasting

Condensed from llama_3.1_8B/model.py

RoPE has to know where it is

A cache changes positional encoding in one subtle way: when the model processes token 4097 alone, its rotary angle must correspond to position 4097, not position 0. So the frequency table takes the past length as an offset, and everything else stays vanilla RoPE.

@staticmethod
def get_rope_params(theta, length, head_dim, past_len=0):
    freq = theta ** ((-torch.arange(0, head_dim, 2)) / head_dim).unsqueeze(0)
    pos = torch.arange(past_len, past_len + length).unsqueeze(1)   # T, 1
    angles = pos * freq                                            # T, head_dim//2
    return torch.cos(angles), torch.sin(angles)

llama_3.1_8B/model.py

The cache itself is concatenation

Each attention layer receives its past keys and values, concatenates the new ones along the time axis, attends over the joined sequence, and hands the joined tensors back as the new past. The score matrix goes rectangular, new queries against all keys, so the mask does too.

past_length = 0 if past_k is None else past_k.size(2)
cos, sin = self.get_rope_params(rope_theta, T, head_size, past_len=past_length)
q, k = self.apply_rope(q, cos, sin), self.apply_rope(k, cos, sin)

if past_k is not None:
    k_all = torch.cat([past_k, k], dim=2)    # B, n_kv, T + past, head_size
    v_all = torch.cat([past_v, v], dim=2)
else:
    k_all, v_all = k, v

scores = q @ k_all.transpose(-1, -2)         # B, n_kv, 4, T, T + past
mask = torch.tril(torch.ones(T, T + past_length, device=x.device))
scores = scores.masked_fill(mask == 0, float("-inf"))
scores = F.softmax(scores / math.sqrt(head_size), dim=-1)
ctx = scores @ v_all

present = (k_all, v_all) if use_cache else None
return x, present

Condensed from the attention forward

The model's forward threads a list of per-layer pasts down the stack and collects the presents on the way back up. That is the entire mechanism behind fast autoregressive decoding: prefill once, then each step pays for one token of attention instead of the whole sequence.

Training details that matter

The loss does the causal shift explicitly, logits against next tokens, with ignore_index=-100 for padding, and the SwiGLU feed-forward is a one-liner: fc3(silu(fc2(x)) * fc1(x)). The config carries Llama 3.1's long-context RoPE scaling fields, the 8x factor and the low and high frequency cutoffs; this implementation keeps rotary vanilla and treats the frequency-dependent scaling as a documented extension, the same config-as-map convention as the rest of the repository. The file also keeps hand-written RMSNorm and LayerNorm classes at the bottom as reference implementations, spare parts for teaching.