What is keepdim actually keeping?
The formula was right, but subtracting a per-row maximum either failed or broadcast in the wrong direction.
The intuition
Reducing (S, V) over the vocabulary gives (S, 1) with keepdim=True, or (S,) without it. Broadcasting aligns dimensions from the right. The singleton vocabulary axis tells the operation to reuse one number across each row.
shifted = logits - logits.amax(-1, keepdim=True)Equal dimension sizes can hide an incorrect broadcast. A small example with different sequence and vocabulary sizes is a better check than relying on a large tensor that happens to run.
Related reading · PyTorch: Broadcasting semantics