P was already normalized in softmax step. PV = P_norm @ V gives the correct attention output. Dividing by row_sum again in the epilogue produces O = O_correct / row_sum (128x too small for uniform data).
P was already normalized in softmax step. PV = P_norm @ V gives the correct attention output. Dividing by row_sum again in the epilogue produces O = O_correct / row_sum (128x too small for uniform data).