core/nn/rotary library

Rotary Positional Embeddings (RoPE) — see https://arxiv.org/abs/2104.09864.

This is the "GPT-NeoX / LLaMA" flavour of RoPE (also called rotate_half): given a tensor of shape [N, headDim], split the last axis in half and rotate:

    q_out = q * cos + rotate_half(q) * sin
  where rotate_half([a | b]) = [-b | a]

The cos/sin tables of shape [maxCtx, headDim] are precomputed once. rotate_half is implemented as a headDim x headDim permutation-with-sign matrix P so that rotate_half(q) == q @ P — this lets us stay inside the existing tensor op surface (matmul + elementwise mul/add) with no new kernels.

Used by GPT-NeoX / Pythia / LLaMA / Mistral / Falcon etc.

Classes

RopeCache