GitLab

260601376  by hongkuny<hongkuny@google.com>:

    reorder Q,K to make TPU faster.

--

PiperOrigin-RevId: 260601376