Implement causal multi-head attention, add a KV cache for token-by-token decoding, then extend it to grouped-query attention.

Write it once as GQA: plain MHA is the case where key-value heads equal query heads, and multi-query is one key-value head.

import torch, torch.nn as nn, torch.nn.functional as F

class GQA(nn.Module):
    def __init__(self, d, n_q, n_kv):
        super().__init__()
        self.hd, self.nq, self.nkv = d // n_q, n_q, n_kv
        self.q = nn.Linear(d, d, bias=False)
        self.kv = nn.Linear(d, 2 * n_kv * self.hd,
                            bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, cache=None):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.nq, self.hd)
        k, v = self.kv(x).view(
            B, T, 2, self.nkv, self.hd).unbind(2)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))
        if cache is not None:
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        cache = (k, v)  # n_kv heads: the memory saving
        r = self.nq // self.nkv
        k = k.repeat_interleave(r, dim=1)
        v = v.repeat_interleave(r, dim=1)
        S = k.shape[2]  # cached + new positions
        # new query i sits at position S - T + i
        m = torch.ones(T, S, dtype=torch.bool,
                       device=x.device).tril(S - T)
        a = q @ k.transpose(-2, -1) / self.hd ** 0.5
        a = a.masked_fill(~m, float("-inf"))
        y = F.softmax(a, dim=-1) @ v
        y = y.transpose(1, 2).reshape(B, T, -1)
        return self.o(y), cache

Interviewers check three things: the mask is offset by the cached length, the cache stores keys and values before repeating them, so it shrinks by the group factor, and cached decoding reproduces a full forward pass token for token.