Fix the attention layer, then add a KV cache

Sparrow Code (fictional) runs a small code-completion model. A teammate wrote its attention layer in NumPy, and the outputs look wrong. You get one layer: causal self-attention with grouped-query attention (GQA) and rotary position embeddings (RoPE), over a batch of right-padded sequences. Python 3, NumPy and the standard library only. No AI tools.

Conventions:

• x has shape (B, T, D). lengths[b] is the number of real tokens in row b, with 1 ≤ lengths[b] ≤ T. Positions at or after lengths[b] are padding.
• Weights: Wq is (D, H × hd); Wk and Wv are (D, Hkv × hd); Wo is (H × hd, D).
• Query head h reads KV head h // (H / Hkv).
• RoPE uses the half-split layout, and the position is the token's index in its own sequence.
• Output rows at padded positions are zeros.
• Use float64.

Part 1 (about 20 minutes). The code below has 5 bugs. Fix them, and name the rule each one breaks.

import numpy as np


def rope(x, positions, base=10000.0):
    """Rotary embedding, half-split layout. x: (..., T, hd); positions: (..., T) ints."""
    hd = x.shape[-1]
    half = hd // 2
    inv_freq = base ** (-np.arange(half) / half)              # (half,)
    ang = positions[..., None] * inv_freq                     # (..., T, half)
    cos, sin = np.cos(ang), np.sin(ang)
    x1, x2 = x[..., :half], x[..., half:]
    return np.concatenate([x1 * cos + x2 * sin, x1 * sin + x2 * cos], axis=-1)


def attention(x, w, n_heads, n_kv_heads, lengths):
    B, T, D = x.shape
    hd = w["Wq"].shape[1] // n_heads
    q = (x @ w["Wq"]).reshape(B, T, n_heads, hd).transpose(0, 2, 1, 3)
    k = (x @ w["Wk"]).reshape(B, T, n_kv_heads, hd).transpose(0, 2, 1, 3)
    v = (x @ w["Wv"]).reshape(B, T, n_kv_heads, hd).transpose(0, 2, 1, 3)
    pos = np.broadcast_to(np.arange(T), (B, 1, T))
    q, k = rope(q, pos), rope(k, pos)
    group = n_heads // n_kv_heads
    k = np.repeat(k, group, axis=1)
    v = np.repeat(v, group, axis=1)
    scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(D)
    causal = np.tril(np.ones((T, T), dtype=bool))
    real = np.arange(T)[None, :] < lengths[:, None]
    keep = causal[None, None] & real[:, None, None, :]
    scores = np.where(keep, scores, 0.0)
    probs = np.exp(scores) / np.exp(scores).sum(-1, keepdims=True)
    out = probs @ v
    out = out.reshape(B, T, n_heads * hd)
    out = out @ w["Wo"]
    return out * real[:, :, None]

Part 2 (about 25 minutes). Add a KV cache for generation.

class KVCache:
    def __init__(self, batch, n_kv_heads, head_dim, max_len): ...
    def nbytes(self): ...   # bytes held by K and V

def prefill(x, w, n_heads, n_kv_heads, lengths, cache):
    """Same output as attention(); also stores the prompts' keys and values."""

def decode_step(x_new, w, n_heads, n_kv_heads, cache):
    """x_new: (B, D), one new token per sequence at position cache.lengths[b]. Returns (B, D)."""

Rows have different lengths, so each sequence's new token has its own position. Each decode output must match a full recompute (attention over the whole sequence so far) to within 1e-10.

Part 3 (about 15 minutes). Add sliding-window attention with window W: a query at position i sees keys at positions i − W + 1 to i. Add window=None to attention. Then add a ring-buffer cache that holds O(W) memory per sequence: RingKVCache(batch, n_kv_heads, head_dim, window), prefill_window(...) and decode_step_window(...).

Example for Part 2: B = 3 and lengths = [5, 2, 4]. After prefill, cache.lengths is [5, 2, 4]. The first decode step writes positions 5, 2 and 4, and its 3 output rows equal rows 5, 2 and 4 of a full recompute with lengths [6, 3, 5].

Before you run anything, read the code in Part 1. Write each bug you can see and the rule it breaks. Then write the shape of every tensor in the function.