A single attention head·8 min
The causal mask
0 / 1
# grade this stage (from your course directory)
npx theinference test --stage causal-mask --watchA language model predicts the next token, so a position must never peek at the future. Before softmax, blank out every entry where the key comes after the query.
The mask (T = 4): "✓" = keep, "·" = set to -1e9
k0 k1 k2 k3
q0 [ ✓ · · · ] position 0 sees only itself
q1 [ ✓ ✓ · · ] position 1 sees 0,1
q2 [ ✓ ✓ ✓ · ] lower triangle (incl. diagonal) survives
q3 [ ✓ ✓ ✓ ✓ ] position 3 sees everything
Keep j ≤ i. Mask j > i. The diagonal — a token attending to itself — always survives.
What to build
Your task
Implement runCausalMask
Implement runCausalMask(scores). Given a [T, T] score matrix, return a copy where every future entry (j > i) is set to -1e9. Leave j ≤ i untouched.
tests/adapters.ts
export function runCausalMask(scores: number[][]): number[][] {
// [T, T] -> [T, T], entries with j > i set to -1e9
}
Off-by-one: `j > i`, not `j >= i`
Use >, not >=. Masking the diagonal (j == i) would forbid a token from attending to itself, corrupting every downstream value. The diagonal must always survive.
How it's graded
Compared with allclose, atol 1e-5. Masked entries are exactly -1e9 on both your side and the reference, so they match precisely.