Where we are in the block

X = E[ids] + PE
R1 = X + Attn(LN1(X))·Wo
out = R1 + FFN(LN2(R1))
FFN(x) = GELU(x·W1)·W2
p = softmax(LN(out)·E⊤) → append, repeat

Inside Attn, for each head h

Q = X̂Wq K = X̂Wk V = X̂Wv → h heads
Oh = softmax(QhKh⊤√dk + mask) Vh
O = [ O1 ‖ … ‖ Oh ]

Shapes of this operation