Given Q, K, V each of shape (T, d), return the (T, d) attention output:

• scale Q @ K.T by 1/sqrt(d)
• set every position j > i to -inf
• softmax row-wise
• multiply by V

Input

Q =
[[0. 1.]
 [2. 3.]
 [4. 5.]]
K =
[[0. 1.]
 [2. 3.]
 [4. 5.]]
V =
[[1. 0.]
 [0. 1.]
 [0. 0.]]

Output

[[1.00000000e+00 0.00000000e+00]
 [8.48604963e-04 9.99151395e-01]
 [8.80371707e-12 2.96710144e-06]]