Support YaRN models (RoFormer implementation in rotary_embedding kernel)

[casper-hansen]

The YaRN model with a context size of 64k and 128k was recently released and pre-trained by people from Nous Research and EleutherAI. It uses the RoFormer type of embeddings that seem different from GPT-NeoX and GPT-J. It is based on the LLaMa 2 model, so support is mostly there, just need some small adjustments.

The original YaRN module uses the flash attention rotary embedding implementation and seems similar in functionality. You may also be interested in the original RoFormer implementation from Huggingface.

Models catalog:
https://huggingface.co/NousResearch/Yarn-Llama-2-7b-64k
https://huggingface.co/NousResearch/Yarn-Llama-2-7b-128k
https://huggingface.co/NousResearch/Yarn-Llama-2-13b-64k
https://huggingface.co/NousResearch/Yarn-Llama-2-13b-128k

[viktor-ferenczi]

Copied from #1027:

• YaRN paper: YaRN: Efficient Context Window Extension of Large Language Models
• YaRN code: YaRN Github