About Mask in your code

[lfchener]

The mask of query why be placed after softmax, but not before softmax?

[soobinseo]

The query mask does not affect softmax for key dimensions because it is only masking for query dimensions. Therefore, you do not need to apply query mask before softmax.

Thanks.