[PyTorch] Exercise MHA fast path in JIT
Tests previously did not exercise this; now they do.
Differential Revision: [D35945821](https://our.internmc.facebook.com/intern/diff/D35945821/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76416
Approved by: https://github.com/ezyang
