Transformer Inference Building Blocks · AI infra interview lesson

Why It Matters

When generating a third token, the first two tokens’ keys and values already exist. Draw Q/K/V shapes and visibility before caching. RoPE, GQA and sampling address separate concerns.

One concrete number: Llama-3-8B has 32 Q heads and 8 KV heads, a GQA group size of G=4. If decode used repeat_interleave to copy K/V up to 32 heads, every layer would move an extra 2 * 24 * 128 * 2 = 12,288 bytes per token; with a 4K context and 32 layers, that is about 1.5 GiB of temporary tensors per request per step, and decode is already memory bound. Broadcasting (a view with stride 0) needs no copy at all.

Walk the Shapes First

B=1, T=3, D=8, H=2, H_kv=1, hd=4, G=H/H_kv=2: prefill 3 tokens, then decode 1:

Rule read off the table: RoPE, GQA, and the KV cache never change which keys a query may see, only how the result is computed; when T differs from S (decode, chunked prefill) the causal mask must be bottom-right aligned, so query i sits at absolute position S - T + i.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: two RoPE layouts and the Meta → HF weight permute

What's new: two ways of pairing dimensions for the same rotation, and the weight reordering that makes them agree.

• Angles: inv_freq[i] = theta^(-2i/hd), angle[pos, i] = pos * inv_freq[i], for i from 0 to hd/2-1.
• Interleaved (Meta's original LLaMA, GPT-J): pairs (2i, 2i+1), equivalent to treating adjacent dims as a complex number multiplied by exp(j*angle).
• rotate_half (HF transformers, nano-vllm, GPT-NeoX): pairs (i, i+hd/2), y = x*cos + rotate_half(x)*sin, rotate_half([x1, x2]) = [-x2, x1].
• Relation: rope_half(x[..., perm]) == rope_inter(x)[..., perm] with perm = [0,2,4,...,1,3,5,...]. HF's conversion script reorders the output rows of Wq and Wk per head with this perm: w.view(H, hd//2, 2, D).transpose(1, 2).reshape(H*hd, D).

import torch
torch.manual_seed(0)

def rope_angles(pos, hd, theta=10000.0):
    inv_freq = 1.0 / theta ** (torch.arange(0, hd, 2, dtype=torch.float64) / hd)
    return pos.to(torch.float64)[:, None] * inv_freq[None, :]          # (T, hd/2)

def rope_interleaved(x, pos, theta=10000.0):          # x: (..., T, hd)
    ang = rope_angles(pos, x.shape[-1], theta)
    xc = torch.view_as_complex(x.reshape(*x.shape[:-1], -1, 2).contiguous())
    return torch.view_as_real(xc * torch.polar(torch.ones_like(ang), ang)).flatten(-2)

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def rope_rotate_half(x, pos, theta=10000.0):
    ang = rope_angles(pos, x.shape[-1], theta)
    cos, sin = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1)  # (T, hd)
    return x * cos + rotate_half(x) * sin

Match weight layout to RoPE pairing

The two RoPE conventions pair different coordinates. Switching the rotation without matching the weight layout pairs the wrong coordinates. This permutation changes layout, not learned values.

def permute_meta_to_hf(w, n_heads):
    out, D = w.shape
    return w.view(n_heads, out // n_heads // 2, 2, D).transpose(1, 2).reshape(out, D)

T, D, H, hd = 5, 16, 2, 8
pos = torch.arange(T)
perm = torch.cat([torch.arange(0, hd, 2), torch.arange(1, hd, 2)])
x = torch.randn(T, hd, dtype=torch.float64)
assert torch.allclose(rope_rotate_half(x[:, perm], pos), rope_interleaved(x, pos)[:, perm])

# relative position: <R(m)q, R(n)k> depends only on m - n
q, k = torch.randn(hd, dtype=torch.float64), torch.randn(hd, dtype=torch.float64)
def score(m, n):
    qm = rope_rotate_half(q[None], torch.tensor([m]))[0]
    kn = rope_rotate_half(k[None], torch.tensor([n]))[0]
    return qm @ kn
assert torch.allclose(score(3, 1), score(10, 8)) and torch.allclose(score(0, 0), q @ k)

# Meta weights + interleaved  vs  permuted HF weights + rotate_half: same QK^T
xs = torch.randn(T, D, dtype=torch.float64)
Wq, Wk = torch.randn(H * hd, D, dtype=torch.float64), torch.randn(H * hd, D, dtype=torch.float64)
def heads(y): return y.view(T, H, hd).transpose(0, 1)                   # (H, T, hd)
qm = rope_interleaved(heads(xs @ Wq.T), pos); km = rope_interleaved(heads(xs @ Wk.T), pos)
qh = rope_rotate_half(heads(xs @ permute_meta_to_hf(Wq, H).T), pos)
kh = rope_rotate_half(heads(xs @ permute_meta_to_hf(Wk, H).T), pos)
assert torch.allclose(qm @ km.transpose(-1, -2), qh @ kh.transpose(-1, -2))
assert torch.allclose(qh, qm[..., perm])
# rotate_half without the permute: scores do not match
qbad = rope_rotate_half(heads(xs @ Wq.T), pos); kbad = rope_rotate_half(heads(xs @ Wk.T), pos)
assert not torch.allclose(qm @ km.transpose(-1, -2), qbad @ kbad.transpose(-1, -2))

The permute applies only to Wq and Wk: Wv and Wo never see RoPE, and attention scores are dot products within one head, so reordering q and k with the same perm leaves them unchanged.

Step 2: GQA by broadcasting instead of copying, and the causal mask when T ≠ S

What's new: view the Q heads as groups (B, H_kv, G, T, hd), give K/V a size-1 dim with unsqueeze, and let matmul broadcast; align the mask bottom-right.

scores[b, g_kv, g, t, s] = q[b, g_kv*G + g, t] · k[b, g_kv, s] / sqrt(hd); query t may see key s iff s <= S - T + t.

import torch
import torch.nn.functional as F
torch.manual_seed(0)

def causal_mask(T, S):
    # True means visible; query t sits at absolute position S - T + t
    return torch.arange(S)[None, :] <= (torch.arange(T)[:, None] + S - T)

def gqa_attention(q, k, v, causal=True):
    B, H, T, hd = q.shape
    H_kv, S = k.shape[1], k.shape[2]
    G = H // H_kv
    qg = q.view(B, H_kv, G, T, hd)
    s = qg @ k.unsqueeze(2).transpose(-1, -2) / hd ** 0.5     # (B, H_kv, G, T, S)
    if causal:
        s = s.masked_fill(~causal_mask(T, S), float("-inf"))
    return (s.softmax(-1) @ v.unsqueeze(2)).reshape(B, H, T, hd)

B, H, H_kv, hd = 2, 8, 2, 4
for T, S in [(5, 5), (1, 7), (3, 7)]:            # square, decode, chunked prefill
    q = torch.randn(B, H, T, hd)
    k, v = torch.randn(B, H_kv, S, hd), torch.randn(B, H_kv, S, hd)
    G = H // H_kv
    k_rep, v_rep = k.repeat_interleave(G, dim=1), v.repeat_interleave(G, dim=1)  # reference: explicit copy
    ref = F.scaled_dot_product_attention(q, k_rep, v_rep, attn_mask=causal_mask(T, S))
    assert torch.allclose(gqa_attention(q, k, v), ref, atol=1e-6)
    ref2 = F.scaled_dot_product_attention(q, k, v, attn_mask=causal_mask(T, S), enable_gqa=True)
    assert torch.allclose(gqa_attention(q, k, v), ref2, atol=1e-6)

# the broadcast dim does not copy
e = k.unsqueeze(2).expand(B, H_kv, G, S, hd)
assert e.stride()[2] == 0 and e.data_ptr() == k.data_ptr()

# trap: with T != S, is_causal=True is top-left aligned, so the decode query only sees key 0
q1 = torch.randn(1, 2, 1, hd); k1 = torch.randn(1, 2, 4, hd); v1 = torch.randn(1, 2, 4, hd)
assert torch.allclose(F.scaled_dot_product_attention(q1, k1, v1, is_causal=True), v1[:, :, :1])

Q head h maps to KV head h // G (G consecutive Q heads share one KV head), the same as repeat_interleave(G, dim=1). Writing k.repeat(1, G, 1, 1) instead gives the mapping h % H_kv, which does not match how the weights are laid out.

Step 3: KV cache (concat / preallocated) checked against a full forward

What's new: combine Steps 1 and 2 into an attention layer with RoPE + GQA, run it incrementally with a concat cache and with a preallocated… (Pro lesson)

Step 4: sampling: temperature, top-k, top-p

What's new: filtering and sampling from logits to a token id. Order: divide by temperature → top-k → top-p → softmax → sample; (Pro lesson)

In Real Systems

nano-vllm nanovllm/layers/rotary_embedding.py : RotaryEmbedding precomputes cos_sin_cache ( (max_pos, 1, hd) , cos in the first half, sin in… (Pro lesson)

Common Pitfalls

Written: loading HF weights with interleaved RoPE (or the reverse). (Pro lesson)

Say It in the Interview

"There are two RoPE layouts: Meta's interleaved pairs (2i, 2i+1) and HF's rotate_half pairs (i, i+d/2); (Pro lesson)

Hands-on Notebook

Each exercise in the notebook has a problem statement and a function skeleton; write your version, then run the test cell below it. (Pro lesson)