Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/itsubaki/gpt/llms.txt

Use this file to discover all available pages before exploring further.

function.RoPE implements Rotary Position Embeddings (RoPE), precomputing cosine and sine lookup tables for all positions and head dimensions up to contextLen. It returns a curried RoPEFunc that can be applied to query and key tensors in any attention layer. Because the tables are computed once at model construction, there is no per-step overhead during forward or backward passes.

RoPEFunc type

type RoPEFunc func(offset int) func(x ...*variable.Variable) *variable.Variable
RoPEFunc is a two-level curried function:
offset
int
Position offset for KV-cache-based generation — the number of tokens already present in the cache. During training (no cache) this is 0. During incremental decoding, pass the current cache length so that new tokens are embedded at the correct absolute positions.
Calling rope(offset) returns a standard autograd-compatible function that accepts a tensor of shape (B, H, C, D) and returns the rotated tensor of the same shape:
DimensionMeaning
BBatch size
HNumber of attention heads
CSequence length (context window)
DHead dimension (embedDim) — must be even

RoPE() constructor

func RoPE(theta float64, embedDim, contextLen int) RoPEFunc
Precomputes the cosine and sine tables and returns a RoPEFunc. Panics if embedDim is odd.
theta
float64
required
Base frequency for the geometric sequence of rotation frequencies. The default is 10000. Higher values extend the effective context length — LLaMA 3 uses 500000.
embedDim
int
required
Head dimension from the attention layer. Must be even; the dimension is split into embedDim/2 rotation pairs.
contextLen
int
required
Maximum sequence length. Determines the size of the precomputed tables: (contextLen, embedDim/2) for both cosine and sine.
Precomputed table construction: For each position pos in [0, contextLen) and each dimension pair i in [0, embedDim/2):
pow   = 2i / embedDim
freq  = 1.0 / theta^pow
angle = pos * freq

cos[pos * halfDim + i] = cos(angle)
sin[pos * halfDim + i] = sin(angle)
Each table has flat length contextLen × (embedDim/2), addressed as pos * halfDim + i.

Rotation formula

During the forward pass, each consecutive pair of dimensions (d, d+1) at every position is rotated by the precomputed angle for that (position, pair) index:
y[d]   = x[d] * cos − x[d+1] * sin
y[d+1] = x[d] * sin + x[d+1] * cos
This is equivalent to complex multiplication: interpreting (x[d], x[d+1]) as a complex number and multiplying by e^{iθ} = cos + i·sin.

Backward pass

The gradient of the rotation with respect to the input is the transpose of the rotation matrix, which is itself a rotation by the conjugate angle ():
gx[d]   =  gy[d] * cos + gy[d+1] * sin
gx[d+1] = −gy[d] * sin + gy[d+1] * cos
This means the backward pass has the same computational structure as the forward pass, with sin sign-flipped for the cross terms. Backward example:
x := variable.New(1, 2, 3, 4, 1, 2, 3, 4).Reshape(1, 1, 2, 4)

var offset int
y := F.RoPE(10000, 4, 4)(offset)(x)
y.Backward()

fmt.Println(x.Grad)
// variable[1 1 2 4]([1 1 1 1 1.3817732906760363 -0.30116867893975674 1.009949833750832 0.9899501670824986])

Usage in model

RoPE is constructed once at model initialisation and the same RoPEFunc is shared across all transformer blocks. Each block applies it independently to its query and key projections:
rope := function.RoPE(10000, embedDim, maxContextLen)

for i := range numOfBlocks {
    gpt.Add(newBlock(i, embedDim, numOfHeads, rope, useCache...))
}
Forward pass example:
x := variable.New(1, 2, 3, 4, 1, 2, 3, 4).Reshape(1, 1, 2, 4)

var offset int
y := F.RoPE(10000, 4, 4)(offset)(x)
fmt.Println(y)
// variable[1 1 2 4]([1 2 3 4 -1.1426396637476532 1.922075596544176 2.9598506679133294 4.029799501669161])

Offset clamping

If offset + C > contextLen, the offset is automatically clamped to contextLen - C:
func (f *RoPET) clamp(C int) int {
    if f.offset+C > f.ContextLen {
        return f.ContextLen - C
    }
    return f.offset
}
This prevents out-of-bounds accesses into the precomputed tables when generating sequences that approach or exceed contextLen. During normal training (offset = 0, C ≤ contextLen) the clamp is a no-op.
Increasing theta — for example to 500000 as used in LLaMA 3 — significantly improves model performance on sequences longer than the training contextLen. The higher base frequency slows the rotation rate of low-frequency dimensions, allowing the model to generalise to positions it has not seen during training.

Build docs developers (and LLMs) love