allow customized head_dim (#7065)

Pull Request resolved: https://github.com/pytorch/executorch/pull/6872

This is for resolving the ask in this [post](https://fb.workplace.com/groups/pytorch.edge.users/permalink/1574875706716050/).

Similar change in HF: https://github.com/huggingface/transformers/pull/32502
ghstack-source-id: 255340016

Differential Revision: [D65974454](https://our.internmc.facebook.com/intern/diff/D65974454/)

Co-authored-by: Lunwen He <lwhecser@gmail.com>