QKV attention programming with einsum fluency
A short programming screen around sequence models and scaled dot-product attention, plus quick ML knowledge questions. Reading equations and expressing the math with einsum-style ops is the point.
Use NumPy. Shapes: b batch, h heads, t query positions, s key positions, d head width, m model width.
import numpy as np
def softmax(x, axis=-1): ...
def attention(q, k, v, mask=None): ... # mask: True = blocked
def causal_mask(t, s=None): ...
def multi_head_attention(x, wq, wk, wv, wo, causal=True): ... # x: (b, t, m); w*: (m, h, d); wo: (h, d, m)
def grouped_query_attention(q, k, v, mask=None): ... # q: (b, h, t, d); k, v: (b, g, s, d)
def decode_step(x_new, cache, wq, wk, wv, wo): ... # one token with a KV cache dict
Sign in to see the worked answer
The full answer, chapter by chapter, with what interviewers expect at each level. Free with an account: one full answer per company.