Implement grouped-query attention over a preallocated KV cache, with no torch.cat, for both prefill chunks and decode.

TL;DR: Allocate K and V once as [B, Hkv, max_len, d], write each new chunk in place at the row's current length, and reshape queries to [B, Hkv, G, T, d] so every group of G query heads reads its KV head without copying it. The causal mask is offset: query t of a chunk sits at absolute position len - n_new + t. Growing the cache with concatenation instead copies about 1.0 TiB to generate 4,096 tokens of a Llama 3.1 8B sequence whose final cache is 512 MiB.

How to approach it

Pin down the shapes and the ownership before writing code: how many query heads and KV heads, whether rows in the batch have different lengths, and whether a call carries one decode token or a prefill chunk. Then say the three rules out loud: KV is written once and never moved, query head h reads KV head h // G, and a query at absolute position i sees keys 0 to i. Everything else is bookkeeping.

A strong answer

The cost of the naive version is the reason the question exists. Concatenating one token onto the cache each step copies the whole cache each step, so the bytes moved grow with the square of the length:

Llama 3.1 8B KV per token = 2 x 32 layers x 8 kv_heads x 128 x 2 B = 131,072 B
tokens copied over 4,096 steps = 0 + 1 + ... + 4,095 = 8,386,560
bytes copied = 8,386,560 x 131,072 B = 1.10e12 B = 1.0 TiB
final cache  = 4,096 x 131,072 B = 512 MiB
read + write at 3.35 TB/s (H100 SXM) = 2 x 1.10e12 / 3.35e12 = 0.66 s per sequence

Two thirds of a second of pure copying per sequence, plus an allocator that fragments as tensors of every size come and go. Preallocation fixes both: the write for a new token is 131,072 B regardless of length. The price is reserving max_len per row up front (32 rows at 8,192 tokens is 32 GiB for this model), which is exactly the waste paged KV caches remove later. This question stops at the contiguous version.

import numpy as np

class KVCache:
    def __init__(self, batch, kv_heads, max_len, head_dim, dtype=np.float32):
        self.k = np.zeros((batch, kv_heads, max_len, head_dim), dtype)  # [B, Hkv, S, d]
        self.v = np.zeros_like(self.k)
        self.lens = np.zeros(batch, dtype=np.int64)                     # tokens per row

    def append(self, k_new, v_new, n_new):
        """k_new, v_new: [B, Hkv, T, d]; n_new[b] <= T real tokens in row b."""
        S = self.k.shape[2]
        for b in range(k_new.shape[0]):
            s, n = self.lens[b], n_new[b]
            if s + n > S:
                raise ValueError(f"row {b}: {s + n} tokens > max_len {S}")
            self.k[b, :, s:s + n] = k_new[b, :, :n]                     # in-place write
            self.v[b, :, s:s + n] = v_new[b, :, :n]
        self.lens += n_new

def gqa_attention(q, cache, n_new):
    """q: [B, Hq, T, d] for the n_new[b] tokens just appended. Returns [B, Hq, T, d]."""
    B, Hq, T, d = q.shape
    Hkv = cache.k.shape[1]
    G = Hq // Hkv                                     # query heads per KV head
    L = int(cache.lens.max()) if B else 0
    if L == 0 or T == 0:                              # empty batch, cache or chunk
        return np.zeros_like(q)
    k, v = cache.k[:, :, :L], cache.v[:, :, :L]       # views, no copy
    qg = q.reshape(B, Hkv, G, T, d)                   # head h uses KV head h // G
    s = np.einsum("bhgtd,bhsd->bhgts", qg, k) / np.sqrt(d)
    q_pos = (cache.lens - n_new)[:, None] + np.arange(T)[None, :]       # [B, T]
    key = np.arange(L)
    ok = (key[None, None, :] <= q_pos[:, :, None]) \
        & (key[None, None, :] < cache.lens[:, None, None])              # [B, T, L]
    ok &= (np.arange(T)[None, :] < n_new[:, None])[:, :, None]          # padded queries
    s = np.where(ok[:, None, None], s, -np.inf)
    m = s.max(-1, keepdims=True)
    m = np.where(np.isfinite(m), m, 0.0)              # fully masked (padding) rows
    p = np.exp(s - m)
    den = p.sum(-1, keepdims=True)
    p = p / np.where(den > 0, den, 1.0)
    out = np.einsum("bhgts,bhsd->bhgtd", p, v)
    return out.reshape(B, Hq, T, d)

The decisions an interviewer reads for:

1. No KV repeat. np.repeat or repeat_interleave on K and V materializes G copies every step. For Llama 3.1 8B (G = 4) a 4,096-token row's 512 MiB of KV becomes 2 GiB written and read back per decode step, so a bandwidth-bound decode now moves several times the bytes it needs. Reshaping the query to [B, Hkv, G, T, d] puts the group on its own axis and lets the einsum broadcast over it.
2. The offset causal mask. For a chunk of T queries appended to a cache that now holds S keys, query t is at position S - T + t. A mask built as tril(ones(T, S)) aligns the chunk to the start of the cache and lets the first query see nothing but key 0. The bug is silent at T = S (plain prefill) and at T = 1 (decode sees everything), so it only fails on chunked prefill, which is where the tests aim.
3. Ragged rows. Each row has its own length, so keys past lens[b] are masked, padded query slots past n_new[b] are masked, and a fully masked row gets zeros instead of NaN from exp(-inf - -inf).

The code was run and checked against a direct reference that repeats KV heads, loops over every query and does a plain softmax over keys 0 to its position: 200 random trials with 1, 2 or 4 KV heads, group sizes 1 to 8, batches of 1 to 3 rows with final lengths from 0 to 40, fed in chunks of 1 to 8 tokens (so decode, chunked prefill with T < S, and rows that receive zero tokens in a call all occur). 2,026 per-row comparisons matched to 1e-10 in float64. Separate checks confirmed that appending past max_len raises and leaves the row's length unchanged, that an empty batch and an all-empty cache return zeros of the right shape, and that the K slice is a view of the preallocated buffer.

The reversal condition: if rows have wildly different lengths, say a 100-token chat beside a 100,000-token document, the dense [B, Hkv, max_len, d] buffer wastes most of its memory and the max-length einsum wastes most of its FLOPs. That is when you move to a block table and a paged kernel.

What interviewers probe next

"Now make it work for MQA and for full MHA." Nothing changes: MQA is Hkv = 1 with G = Hq, MHA is G = 1, and the tests above already cover both ends.

"How would this look in PyTorch on a GPU?" Use k_cache[b, :, s:s+n].copy_(k_new) or index_copy_ into a buffer allocated once, and call scaled_dot_product_attention with enable_gqa=True or a kernel that takes the KV head mapping, so the repeat never materializes.

"Your decode step is 0.3 ms slower than expected at batch 64." Check that the einsum is not iterating over max_len instead of the current maximum length; slicing to lens.max() is the difference between reading the live cache and reading the whole reservation.

How it's scored

Weak: grows the cache with concatenation, repeats K and V to Hq heads, or builds the causal mask from the chunk alone so chunked prefill attends to the wrong keys.

Hire: preallocates, writes in place at the current length, maps query heads to KV heads by integer division, and gets decode and full prefill right.

Strong Hire: quantifies the concatenation cost, offsets the causal mask for T < S, handles ragged rows and fully masked padding without NaN, and tests against a direct reference on chunked inputs.

Common mistakes

• Mapping query head h to KV head h % Hkv instead of h // G, which interleaves groups and matches no released checkpoint's layout.
• Building the causal mask as a T by T lower triangle and padding it, so a 512-token chunk appended to 3,000 cached tokens cannot see its own prefix.
• Letting a fully masked padded query produce NaN, which then propagates through the next layer's normalization.
• Incrementing the length before the write, so an overflow error leaves the row claiming tokens it never stored.