Causal Attention Debugging

You’re implementing causal self-attention for a text-generation model. Each token should attend only to itself and earlier tokens, but testing shows that changing the last token’s value vector also changes the first token’s output.

The implementation uses the following pseudocode. Q, K, and V are the projected queries, keys, and values for one sequence and one attention head; dropout and padding are disabled.

scores = (Q @ K.T) / sqrt(d_k)
scores = scores * allowed_mask
weights = softmax(scores, axis=-1)
output = weights @ V

For a three-token sequence, allowed_mask is:

[[1, 0, 0],
 [1, 1, 0],
 [1, 1, 1]]

How would you correct this implementation so that future tokens cannot affect an earlier token’s output?