RoPE - Rotary Position Embeddings

Implements Rotary Position Embeddings (RoPE) as used in LLaMA, GPT-NeoX, and BitNet.

Mathematical Definition

RoPE encodes position information by rotating query/key vectors in complex space:

NeoX split-half style (used by BitNet b1.58):
For position m and dimension i (0 ≤ i < headDim/2):
θᵢ = m × base^(-2i/headDim)  where base = 500000.0 for BitNet

Dimension pairs: (x[i], x[i + headDim/2])

Applied as:
x'[i]             = x[i] × cos(θ) - x[i + headDim/2] × sin(θ)
x'[i + headDim/2] = x[i] × sin(θ) + x[i + headDim/2] × cos(θ)

Why RoPE?

Traditional positional encoding (sinusoidal):

• Absolute positions → fixed vectors added to embeddings
• No relative position information in attention scores

RoPE advantages:

1. Relative position encoding: Attention scores depend on relative positions
2. Efficient: No learned parameters, just rotation
3. Extrapolation: Can handle longer sequences than training length
4. Mathematically elegant: Complex number rotation in frequency space

Attention Score Property

For queries Q and keys K at positions m and n:

(RoPE(Q, m) · RoPE(K, n)) = f(Q, K, m-n)

The dot product depends only on relative position (m-n), not absolute positions!

Implementation Strategy

Precomputed approach (llama.cpp):

1. Precompute cos(θ) and sin(θ) for all positions and dimensions
2. Store in lookup tables: cos_cache[pos][dim], sin_cache[pos][dim]
3. During forward pass: read from cache and apply rotation

On-the-fly computation (this implementation):

1. Compute θ = pos × base^(-2i/d) in shader
2. Compute cos(θ) and sin(θ) using GPU intrinsics
3. Apply rotation to adjacent dimension pairs

References

• RoFormer paper: "RoFormer: Enhanced Transformer with Rotary Position Embedding" (Su et al., 2021)
• LLaMA: https://github.com/facebookresearch/llama (rope_forward)
• llama.cpp: ggml/src/ggml.c (ggml_rope_impl)

Layer Configuration

structure Hesper.Layers.RoPE.Config : Type

RoPE configuration

• dim : Nat
• maxSeqLen : Nat
• base : Float

def Hesper.Layers.RoPE.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Layers.RoPE.instReprConfig : Repr Config

Equations

• Hesper.Layers.RoPE.instReprConfig = { reprPrec := Hesper.Layers.RoPE.instReprConfig.repr }

Helper Functions

def Hesper.Layers.RoPE.computeTheta (pos : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (dimPair : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.u32)) (totalDim : Nat) (base : Float) : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)

Compute rotation angle θ for position and dimension pair θᵢ = pos × base^(-2i/d)

Equations

• One or more equations did not get rendered due to their size.

GPU Kernel Implementation

def Hesper.Layers.RoPE.ropeKernel (config : Config) (batchSize seqLen numHeads : Nat) (headDimOverride posOffset : Nat := 0) : WGSL.Monad.ShaderM Unit

RoPE kernel: Apply rotary embeddings to query or key tensor

Input shape: [batch, seq_len, num_heads, head_dim] Output shape: [batch, seq_len, num_heads, head_dim]

Algorithm:

for each position pos in sequence:
  for each dimension pair (2i, 2i+1):
    θ = pos × base^(-2i/d)
    x'[2i]   = x[2i]   × cos(θ) - x[2i+1] × sin(θ)
    x'[2i+1] = x[2i]   × sin(θ) + x[2i+1] × cos(θ)

Workgroup strategy: Each thread handles one dimension pair for one token

@param config RoPE configuration @param batchSize Batch size @param seqLen Current sequence length @param numHeads Number of attention heads

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RoPE.ropeKernelDynamic (config : Config) (batchSize seqLen numHeads : Nat) (headDimOverride : Nat := 0) : WGSL.Monad.ShaderM Unit

RoPE kernel with dynamic posOffset from params buffer. Produces identical WGSL regardless of position, enabling pipeline caching. Params buffer layout: [posOffset: u32] (4 bytes minimum, but shares attention params buffer) For single-token inference: batchSize=1, seqLen=1, reads posOffset from params[0].

Equations

• One or more equations did not get rendered due to their size.

Optimized: Cached RoPE (Future)

def Hesper.Layers.RoPE.ropeCachedKernel (config : Config) (batchSize seqLen numHeads : Nat) : WGSL.Monad.ShaderM Unit

Precomputed RoPE: Use cached cos/sin values

This is more efficient for long sequences where we repeatedly apply RoPE to the same positions. Common in inference with KV caching.

Approach:

1. Precompute cos(θ) and sin(θ) for all positions and dimensions
2. Store in buffers: cos_cache[max_seq_len][head_dim/2], sin_cache[...]
3. During forward: lookup instead of compute

@param config RoPE configuration @param batchSize Batch size @param seqLen Current sequence length @param numHeads Number of attention heads

Equations

• One or more equations did not get rendered due to their size.

High-Level API

structure Hesper.Layers.RoPE.RoPE : Type

RoPE layer structure

• config : Config

def Hesper.Layers.RoPE.create (config : Config) : IO RoPE

Create RoPE layer (no learned parameters)

@param config RoPE configuration

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RoPE.forward {β : Type} [GPUBackend β] (ctx : β) (layer : RoPE) (inputBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen numHeads : Nat) (headDim posOffset : Nat := 0) : IO Unit

Apply RoPE to query or key tensor

@param device WebGPU device @param layer RoPE layer @param inputBuf GPU buffer [batch, seq_len, num_heads, head_dim] @param outputBuf GPU buffer for output (same shape) @param batchSize Batch size @param seqLen Current sequence length @param numHeads Number of attention heads

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RoPE.forwardDynamic {β : Type} [GPUBackend β] (ctx : β) (layer : RoPE) (inputBuf outputBuf paramsBuf : GPUBackend.Buf β) (batchSize seqLen numHeads : Nat) (headDim : Nat := 0) (preparedRef : Option (IO.Ref (Option (GPUBackend.CachedDispatch β))) := none) : IO Unit

Apply RoPE with dynamic posOffset from a params buffer. The params buffer must contain [posOffset: u32, ...] (posOffset at index 0). Produces identical WGSL across tokens → enables pipeline + bind group caching.

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RoPE.forwardCached {β : Type} [GPUBackend β] (ctx : β) (layer : RoPE) (inputBuf cosCacheBuf sinCacheBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen numHeads : Nat) : IO Unit

Apply cached RoPE (requires precomputed cos/sin buffers)

@param device WebGPU device @param layer RoPE layer @param inputBuf GPU buffer [batch, seq_len, num_heads, head_dim] @param cosCacheBuf Precomputed cos values [max_seq_len, head_dim/2] @param sinCacheBuf Precomputed sin values [max_seq_len, head_dim/2] @param outputBuf GPU buffer for output @param batchSize Batch size @param seqLen Current sequence length @param numHeads Number of attention heads

Equations

• One or more equations did not get rendered due to their size.

Cache Generation Utilities

def Hesper.Layers.RoPE.generateCache (config : Config) (numHeads : Nat) : IO (ByteArray × ByteArray)

Generate cos/sin cache on CPU for later GPU upload

This is typically done once during model initialization.

@param config RoPE configuration @param numHeads Number of attention heads @return (cos_cache, sin_cache) as ByteArrays

Equations

• One or more equations did not get rendered due to their size.