X = token_embeddings

Q = X @ W_Q
K = X @ W_K
V = X @ W_V

scores = Q @ K.T

scaled_scores = scores / sqrt(d_k)

weights = softmax(scaled_scores)

output = weights @ V