3.5 A deeper look at Transformer positional encoding

Positional encoding looks minor in the Transformer architecture, but it is essential. Without it, a Transformer cannot tell “cat eats fish” from “fish eats cat”. This article starts from the root question of why position information is needed, walks through the ideas and implementations of Sinusoidal encoding, learned encodings, RoPE, ALiBi, and other mainstream schemes, then digs into positional encoding extrapolation for long-context extension, and finally connects that to AI Infra practice: CUDA kernel fusion and KV Cache management.

📑 Contents

• 1. Why Transformers need position information
• 2. Absolute vs relative positional encoding
• 3. Sinusoidal positional encoding
• 4. Learned positional encoding
• 5. RoPE: rotary position embedding
• 6. ALiBi: attention with linear biases
• 7. Long-context extension techniques
• 8. AI Infra engineering practice
• Self-check list
• References

---

1. Why Transformers need position information

1.1 Word order decides meaning

Natural language is a symbol system that depends heavily on order. The same words, rearranged, can mean something completely different.

A few examples:

• “He lent me a book” vs “I lent him a book” — lender and borrower swap
• “Only you understand me” vs “Only I understand you” — subject and object flip
• “He doesn’t know she came” vs “She doesn’t know he came” — who knows and who doesn’t reverse

Even at a subtler level, word order changes modification and tone:

• “a very good result” vs “a good very result” — the latter is ungrammatical
• “I am especially happy today” vs “Today I am especially happy” — a slight shift in emphasis (the first stresses “I”, the second stresses “today”)

These examples point to a basic fact: word order is part of meaning. Any language model that ignores order cannot understand natural language correctly.

1.2 Attention’s permutation equivariance: a math view

Now let’s state carefully why raw Self-Attention cannot see position.

Recall the core Self-Attention computation. Given input matrix $X \in \mathbb{R}^{N \times d}$ (N tokens, each d-dimensional), the process is:

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$ $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

Now permute the input arbitrarily. Let $\Pi$ be an $N \times N$ permutation matrix (exactly one 1 in each row and column, zeros elsewhere). $\Pi X$ shuffles the rows of X. The permuted Q, K, V are:

$$Q' = \Pi X W_Q = \Pi Q, \quad K' = \Pi K, \quad V' = \Pi V$$

Then the Attention output after permutation is:

$$\text{softmax}\left(\frac{Q' K'^T}{\sqrt{d_k}}\right) V' = \text{softmax}\left(\frac{\Pi Q K^T \Pi^T}{\sqrt{d_k}}\right) \Pi V$$

Because softmax is row-wise, and $\Pi$ only reorders rows:

$$\text{softmax}\left(\frac{\Pi Q K^T \Pi^T}{\sqrt{d_k}}\right) = \Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) \cdot \Pi^T$$

The final result:

$$\Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) \cdot \Pi^T \cdot \Pi V = \Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

That is what permutation equivariance means: after you shuffle the input, the output is only shuffled the same way, and the information each token aggregates does not change. In other words, raw Attention treats the input as a set, not a sequence — the token at position 0 and the token at position 99 look the same to it.

That is why we must inject positional encoding from outside: Attention itself is position-blind. It needs an external position signal to tell different word orders apart.

---

2. Absolute vs relative positional encoding

Before the concrete schemes, it helps to separate two design ideas.

2.1 Absolute positional encoding

Absolute positional encoding assigns each position in the sequence a fixed encoding vector $p_i$, then adds it to (or concatenates it with) the corresponding token embedding:

$$x_i' = x_i + p_i$$

The idea is direct: position 0 gets an “ID card” $p_0$, position 1 gets another $p_1$, and so on. The model learns relationships among these IDs and thereby senses relative position indirectly.

Representative schemes: Sinusoidal encoding (fixed), and the learned encodings in BERT/GPT-2 (learned).

2.2 Relative positional encoding

Relative positional encoding does not give each position a fixed label. Instead, when computing Attention scores it injects how far apart two tokens are.

An analogy: absolute encoding is like giving everyone a room number (“I live in 301, you live in 305”) and inferring distance from the numbers; relative encoding tells the two people while they talk that “you are four rooms apart”.

Representative schemes: RoPE (rotary position embedding), ALiBi (linear bias).

2.3 Comparing the two ideas

Almost all current large models use relative positional encoding (mainly RoPE). The reasons are analyzed later.

---

3. Sinusoidal positional encoding

3.1 The formulas

“Attention Is All You Need” proposed a positional encoding with no learnable parameters. For position $pos$ and dimension index $i$, the values are:

$$PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right)$$ $$PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right)$$

Here $pos$ is the token’s position in the sequence (starting at 0), $i$ is the index of the dimension pair ($0 \leq i < d/2$), and $d$ is the model’s hidden dimension.

Each adjacent pair $(2i, 2i+1)$ shares the same frequency $\omega_i = 1 / 10000^{2i/d}$, encoded with sine and cosine. As $i$ goes from 0 to $d/2 - 1$, the frequency decays exponentially from 1 (when $i=0$, the denominator is $10000^0 = 1$) down to $1/10000$ (when $i = d/2 - 1$).

3.2 Multi-frequency intuition: a radio-tuning metaphor

The formula looks abstract at first, but the intuition behind it is natural.

Imagine a row of radio knobs from left to right. The leftmost knob is tuned to a very high frequency — the needle spins fast, almost a full turn per position. The rightmost knob is tuned very low — the needle barely moves even after a hundred positions. The knobs in between step down in frequency.

Now give yourself a position number (say pos=42) and turn each knob to the matching angle. Because the frequencies differ, the pattern of angles for pos=42 is unique. Change the number (say pos=43) and every needle shifts by a different amount, forming another unique pattern.

That is how Sinusoidal encoding works:

• High-frequency dimensions (small $i$) change quickly and separate nearby positions (pos=42 and pos=43 differ clearly here)
• Low-frequency dimensions (large $i$) change slowly and separate distant positions (pos=0 and pos=1000 differ clearly here, while pos=42 and pos=43 look almost the same)

All frequencies together give each position a unique “spectral fingerprint”.

3.3 Python: generate and visualize Sinusoidal positional encoding

The code below builds a Sinusoidal positional encoding matrix and plots it as a heatmap:

import numpy as np
import matplotlib.pyplot as plt

def sinusoidal_position_encoding(max_len, d_model):
    """
    Generate a Sinusoidal positional encoding matrix.

    Args:
        max_len: maximum sequence length
        d_model: model hidden dimension

    Returns:
        PE: positional encoding matrix of shape (max_len, d_model)
    """
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]          # (max_len, 1)
    dim_pair_index = np.arange(0, d_model, 2)[np.newaxis, :]  # (1, d_model/2)

    # Frequency: omega_i = 1 / 10000^(2i/d_model)
    # Equivalent to exp(-2i/d_model * ln(10000))
    div_term = np.exp(dim_pair_index * -(np.log(10000.0) / d_model))

    # Even dims use sin, odd dims use cos
    PE[:, 0::2] = np.sin(position * div_term)
    PE[:, 1::2] = np.cos(position * div_term)

    return PE

# Generate positional encodings
max_len = 128
d_model = 64
PE = sinusoidal_position_encoding(max_len, d_model)

# Visualize
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Heatmap: full encoding matrix
im = axes[0].imshow(PE, aspect='auto', cmap='RdBu_r', interpolation='nearest')
axes[0].set_xlabel('Dimension Index')
axes[0].set_ylabel('Position')
axes[0].set_title('Sinusoidal Position Encoding Matrix')
plt.colorbar(im, ax=axes[0])

# Curves: encoding values vs position for selected dims
dims_to_plot = [0, 1, 10, 11, 30, 31]
for dim in dims_to_plot:
    label = f'dim {dim} ({"sin" if dim % 2 == 0 else "cos"}, freq_idx={dim // 2})'
    axes[1].plot(PE[:, dim], label=label, alpha=0.8)
axes[1].set_xlabel('Position')
axes[1].set_ylabel('Encoding Value')
axes[1].set_title('Encoding Values at Selected Dimensions')
axes[1].legend(fontsize=8)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('sinusoidal_pe.png', dpi=150)
plt.show()

Run it and you will see: in the heatmap, left columns (high frequency) show dense stripes and right columns (low frequency) show wide bands, which matches the different frequencies. In the curves, low-index dims oscillate hard and high-index dims change slowly.

3.4 Relative position via a linear transform

Sinusoidal encoding has a neat math property: the offset between any two positions can be expressed by a fixed linear transform. So the model can, in principle, learn relative position.

Concretely, for any fixed offset $k$ there is a matrix $M_k$ that depends only on $k$ such that:

$$PE(pos + k) = M_k \cdot PE(pos)$$

The reason is the angle-addition formulas. For a single frequency $\omega_i$, consider the dimension pair $(2i, 2i+1)$:

$$\begin{bmatrix} \sin(\omega_i (pos + k)) \\ \cos(\omega_i (pos + k)) \end{bmatrix} =$$ $$\begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix} \begin{bmatrix} \sin(\omega_i \cdot pos) \\ \cos(\omega_i \cdot pos) \end{bmatrix}$$

The rotation matrix on the right depends only on the offset $k$ and frequency $\omega_i$, not on absolute position $pos$. Stacking all dimension pairs gives a block-diagonal matrix $M_k$.

So although Sinusoidal encoding is an “absolute” encoding, its trig structure lets the model extract relative position by learning a linear transform. That ability is indirect, though — the model has to discover and use the property from data, rather than being given an explicit relative-position signal.

3.5 Limitations

Despite the clever design, Sinusoidal encoding has been largely replaced in the large-model era, mainly because:

1. Signal dilution: Position is added into the token embedding at the input layer. After dozens of nonlinear layers, that additive signal can get drowned out, so deep layers struggle to use it.

2. Hard extrapolation: In theory you can generate encodings for any length, but the model only saw finite lengths in training. For positions beyond the training length the values are defined, yet the model never trained on them, and Attention patterns often generalize poorly.

3. Indirect relative position: As in §3.4, the model must learn to pull out relative position via a linear transform. That adds learning burden compared with injecting relative position explicitly into Attention.

---

4. Learned positional encoding

4.1 Basic idea

Given the limits of fixed trig encodings, a natural idea is: let the model learn an encoding vector for each position. That is Learned Positional Embedding.

You keep a learnable embedding table of shape $(L_{max}, d)$, where $L_{max}$ is the max supported sequence length and $d$ is the model dimension. The encoding at position $pos$ is the $pos$-th row of that table, added to the token embedding:

$$x_i' = x_i + E_{pos}[i]$$

It is the same mechanism as word embeddings — word embeddings look up by token ID; position embeddings look up by position ID.

4.2 Representative models

BERT: learned positional encoding, max length 512. The position embedding matrix has 512 x 768 = 393,216 parameters, tiny next to BERT-base’s ~110M total.

GPT-2: also learned positional encoding, max length 1024. The matrix is 1024 x 768 (GPT-2 Small) or 1024 x 1600 (GPT-2 XL).

4.3 Pros and cons

Pros:

• Extremely simple: one embedding lookup plus an add
• The model can learn any position pattern, unconstrained by a fixed functional form
• Within the training length, performance usually matches or slightly beats Sinusoidal encoding

Cons:

• Hard length limit: max length is fixed when the model is defined. Longer sequences need retraining or awkward interpolation
• No extrapolation: positions outside the table have no encoding at all
• Shared absolute-encoding flaws: deep layers can still dilute the signal, and relative position still has to be learned indirectly

Those limits pushed researchers toward schemes that encode relative position directly in Attention — and RoPE became the most successful answer.

---

5. RoPE: rotary position embedding

RoPE (Rotary Position Embedding) was proposed by Jianlin Su in 2021 and is now used by almost every mainstream large model: LLaMA, Mistral, Qwen, DeepSeek, and others. One sentence captures the idea: rotate the Q and K vectors to encode position, so Attention scores naturally carry relative position.

5.1 Start in 2D: a compass metaphor

Before the derivation, build geometric intuition.

Imagine a compass in your hand, the needle pointing some direction as a 2D vector $(x, y)$. Rotate the compass by angle $\theta$ and the needle points to $(x', y')$. Rotation keeps length (the norm) and only changes direction.

In math:

$$\begin{bmatrix} x' \\ y' \end{bmatrix} = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix}$$

RoPE’s core move: use the token’s position to set the rotation angle. Position 0 does not rotate, position 1 rotates by $\theta$, position 2 by $2\theta$, and so on. Position $m$ rotates by $m\theta$.

Now take two tokens: the Q at position $m$ is rotated by $m\theta$, the K at position $n$ by $n\theta$. Their inner product (Attention score) depends only on the angle difference $(m - n)\theta$, not on each absolute angle, because rotation is a rigid transform.

That is the geometric reason RoPE encodes relative position: the inner product of two rotated vectors depends only on the difference of their rotation angles.

5.2 Derivation in 2D

Now the strict derivation. In 2D, let Q and K be the query and key at positions $m$ and $n$. RoPE applies position-dependent rotations:

$$\tilde{q} = R(m\theta) \cdot q = \begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} q_0 \\ q_1 \end{bmatrix}$$ $$\tilde{k} = R(n\theta) \cdot k = \begin{bmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{bmatrix} \begin{bmatrix} k_0 \\ k_1 \end{bmatrix}$$

The inner product after rotation:

$$\tilde{q}^T \tilde{k} = q^T R(m\theta)^T R(n\theta) k$$

Using $R(\alpha)^T = R(-\alpha)$:

$$R(m\theta)^T R(n\theta) = R(-m\theta) R(n\theta) = R((n-m)\theta)$$

Finally:

$$\tilde{q}^T \tilde{k} = q^T R((n-m)\theta) k$$

The result depends only on $(n - m)$, the relative position difference, not on the concrete absolute positions $m$ and $n$. That is exactly what we want.

5.3 Extension to higher dimensions

In practice the head dimension $d_k$ is much larger than 2 (usually 64 or 128). RoPE extends naturally: split the $d_k$-dim vector into $d_k / 2$ 2D subspaces, rotate each independently with a different frequency.

For the $i$-th 2D subspace ($0 \leq i < d_k / 2$), the rotation angle is $m \cdot \theta_i$, with:

$$\theta_i = \frac{1}{10000^{2i / d_k}}$$

Note this frequency matches Sinusoidal encoding exactly.

Stack the subspace rotations into a block-diagonal matrix:

$$R_m = \begin{bmatrix} R(m\theta_0) & & & \\ & R(m\theta_1) & & \\ & & \ddots & \\ & & & R(m\theta_{d_k/2 - 1}) \end{bmatrix}$$

Each $R(m\theta_i)$ is a $2 \times 2$ rotation matrix.

The full RoPE operation is:

$$\tilde{q}_m = R_m \cdot q_m, \quad \tilde{k}_n = R_n \cdot k_n$$

5.4 Why the inner product depends only on relative position: high-dim proof

In high dimensions the proof is still short. The rotated inner product is:

$$\tilde{q}_m^T \tilde{k}_n = q_m^T R_m^T R_n k_n$$

Because $R_m$ is block-diagonal, so is $R_m^T$ (each block transposed):

$$R_m^T R_n = \begin{bmatrix} R(m\theta_0)^T R(n\theta_0) & & \\ & \ddots & \\ & & R(m\theta_{d_k/2-1})^T R(n\theta_{d_k/2-1}) \end{bmatrix}$$

Each diagonal block:

$$R(m\theta_i)^T R(n\theta_i) = R(-m\theta_i) R(n\theta_i) = R((n-m)\theta_i)$$

So:

$$R_m^T R_n = R_{n-m}$$

Finally:

$$\tilde{q}_m^T \tilde{k}_n = q_m^T R\_{n-m} k_n = f(q_m, k_n, n - m)$$

The inner product depends only on the vectors themselves and the relative difference $n - m$, not on absolute positions. Done.

5.5 Efficient implementation: skip building the rotation matrix

In production you do not build a $d_k \times d_k$ block-diagonal matrix and multiply. 2D rotation is element-wise multiply and add:

$$\tilde{q}_{2i} = q_{2i} \cos(m\theta_i) - q_{2i+1} \sin(m\theta_i)$$ $$\tilde{q}_{2i+1} = q_{2i} \sin(m\theta_i) + q_{2i+1} \cos(m\theta_i)$$

Here is a Python implementation of the core RoPE ops:

import numpy as np

def precompute_freqs(dim, max_seq_len, base=10000.0):
    """
    Precompute RoPE frequency parameters and trig values.

    Args:
        dim: attention head dimension (d_k)
        max_seq_len: maximum sequence length
        base: frequency base, default 10000

    Returns:
        cos_cached: cosines of shape (max_seq_len, dim/2)
        sin_cached: sines of shape (max_seq_len, dim/2)
    """
    # Frequency: theta_i = 1 / base^(2i / dim), i = 0, 1, ..., dim/2 - 1
    freq_indices = np.arange(0, dim, 2, dtype=np.float32)  # [0, 2, 4, ..., dim-2]
    freqs = 1.0 / (base ** (freq_indices / dim))            # (dim/2,)

    # Position sequence
    positions = np.arange(max_seq_len, dtype=np.float32)     # [0, 1, 2, ..., max_seq_len-1]

    # Outer product: angle for each (position, freq_pair)
    angles = np.outer(positions, freqs)  # (max_seq_len, dim/2)

    cos_cached = np.cos(angles)
    sin_cached = np.sin(angles)

    return cos_cached, sin_cached

def apply_rope(q, cos, sin):
    """
    Apply RoPE rotation to a query or key vector.

    Args:
        q: vector of shape (seq_len, dim); dim must be even
        cos: precomputed cosines of shape (seq_len, dim/2)
        sin: precomputed sines of shape (seq_len, dim/2)

    Returns:
        q_rotated: rotated vector, same shape as input
    """
    dim = q.shape[-1]
    assert dim % 2 == 0, "dimension must be even"

    # Split into even and odd dimensions
    q_even = q[:, 0::2]   # (seq_len, dim/2) -- take q_0, q_2, q_4, ...
    q_odd  = q[:, 1::2]   # (seq_len, dim/2) -- take q_1, q_3, q_5, ...

    # Rotate each 2D subspace
    q_even_rotated = q_even * cos - q_odd * sin
    q_odd_rotated  = q_even * sin + q_odd * cos

    # Interleave and merge
    q_rotated = np.zeros_like(q)
    q_rotated[:, 0::2] = q_even_rotated
    q_rotated[:, 1::2] = q_odd_rotated

    return q_rotated

# Usage example
dim = 128        # head dimension
seq_len = 64     # sequence length

# Precompute trig values (once at init)
cos_cached, sin_cached = precompute_freqs(dim, seq_len)

# Fake Q and K
q = np.random.randn(seq_len, dim).astype(np.float32)
k = np.random.randn(seq_len, dim).astype(np.float32)

# Apply RoPE
q_rotated = apply_rope(q, cos_cached, sin_cached)
k_rotated = apply_rope(k, cos_cached, sin_cached)

# Check: rotated inner product depends only on relative position
# Q at 10 and K at 15
dot1 = np.dot(q_rotated[10], k_rotated[15])

# Q at 20 and K at 25 (same relative gap of 5)
# Note: original Q and K values differ here, so the dots differ
# But for the same q, k vectors placed at different absolute positions,
# equal relative gaps give equal inner products
print(f"dot(q_rotated[10], k_rotated[15]) = {dot1:.4f}")

Note: two ways to pair dimensions. The implementation above pairs adjacent dims $(q_{2i}, q_{2i+1})$, which is Meta’s original LLaMA and GPT-J style. HuggingFace transformers’ LLaMA implementation pairs front and back halves, $(q_i, q_{i+d/2})$, via rotate_half, which is GPT-NeoX style:

def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return np.concatenate((-x2, x1), axis=-1)

# HF style: expand cos/sin to (seq_len, dim) by repeating the dim/2 table twice on the last axis
q_rotated = q * cos + rotate_half(q) * sin

The two are mathematically equivalent, but each requires $W_q$ and $W_k$ output dims ordered the matching way. HF’s conversion script (convert_llama_weights_to_hf.py) permutes $W_q$ and $W_k$ when converting Meta weights to HF format. When you port a model or write your own RoPE kernel, make sure weight layout and pairing match: a mismatch will not crash, but the model output will be completely wrong. This is a common pitfall when porting models.

5.6 RoPE vs Sinusoidal: links and differences

RoPE and Sinusoidal encoding are deeply related, but they work differently.

Links:

• Both use exactly the same frequency $\theta_i = 1 / 10000^{2i/d}$
• Both rely on trig periodicity and angle-addition
• The linear-transform property of relative position in Sinusoidal encoding (§3.4) is essentially a rotation; RoPE can be seen as making that property explicit

Differences:

One difference worth stressing: RoPE does not transform V at all. Position only affects “who should attend to whom” (Attention weights), not “what information is passed after attending” (V content). That design is cleaner — position controls the routing of information flow, not the information itself.

---

6. ALiBi: attention with linear biases

ALiBi (Attention with Linear Biases), proposed by Press et al. in 2022, takes a very different path from RoPE.

6.1 Core idea

ALiBi does not modify Q or K. After Attention scores are computed, it adds a negative bias proportional to distance:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + m \cdot \text{Bias}\right) V$$

Bias matrix entries are:

$$\text{Bias}[i][j] = -|i - j|$$

$m$ is a per-head scalar slope. Different heads use different $m$ values (usually a geometric series such as $1/2, 1/4, 1/8, \ldots$).

In plain terms: ALiBi puts a “distance penalty” on Attention scores — the farther two tokens are, the more is subtracted, so the model tends to attend nearby. Different heads decay at different rates: some almost only look at neighbors (large slope), others can still look far (small slope).

6.2 Pros and limits

Pros:

• Extremely simple: one add on Attention scores
• No learnable parameters
• Strong extrapolation: the paper shows a model trained at 1024 can extrapolate to 2048 or longer at inference
• Almost zero compute cost

Limits:

• Limited expressive power for position (only one pattern: linear distance decay)
• In very long contexts, the linear penalty may over-suppress distant tokens
• In real large-model practice, ALiBi is adopted far less than RoPE

ALiBi is mainly used by a few models such as BLOOM. Among positional encodings for large models, RoPE has become the de facto standard for its expressiveness and extrapolation extensibility.

---

7. Long-context extension techniques

Large models are usually trained with a fixed context window (for example 4096 or 8192 tokens), but users often need longer inputs — a novel, a long code file, a full meeting transcript. That raises a core engineering challenge: how can the model handle sequences longer than the training length at inference?

7.1 The challenge of extrapolating positional encoding

Take RoPE. If the training max length is 4096, then $m$ in $m \cdot \theta_i$ is at most 4095. At inference, position $m = 8000$ pushes some high-frequency dims far beyond anything seen in training, and Attention patterns can shift hard.

Concretely, high-frequency dims (small $i$, large $\theta_i$) already “spun many turns” in training, so a few more matter less; low-frequency dims (large $i$, small $\theta_i$) may have spun less than one turn in training, and suddenly asking for two turns puts them in unknown territory. Uneven extrapolation across frequency dims is the root of the problem.

7.2 Position Interpolation

The simplest fix is Position Interpolation: shrink position indices proportionally so the new max position maps to the original training max.

To extend context from $L_{train}$ to $L_{target}$, the scale factor is $s = L_{target} / L_{train}$, and the RoPE angle at position $m$ becomes:

$$m' \cdot \theta_i = \frac{m}{s} \cdot \theta_i$$

Intuitively, encodings that used to cover 0 to 4095 are “compressed” into 0 to $4095/s$, freeing room for a longer sequence.

Problem: uniform scaling treats every frequency the same, but high-frequency dims do not need scaling (they already extrapolate well); low-frequency dims need heavy scaling. Uniform interpolation hurts high-frequency resolution of nearby positions.

7.3 NTK-aware Scaling

The core insight of NTK-aware Scaling (Neural Tangent Kernel-aware Scaling) is: do not scale every frequency the same way; give different frequencies different strategies.

Concretely, change RoPE’s base (originally 10000) to:

$$\text{base}' = \text{base} \cdot s^{d/(d-2)}$$

where $s$ is the extension factor. The effect:

• High-frequency dims (small $i`): frequency almost unchanged, keeping nearby-position resolution
• Low-frequency dims (large $i`): frequency drops a lot, wavelength stretches, covering a larger position range

The “NTK” name comes from the theoretical motivation: under the Neural Tangent Kernel view, the frequency distribution of Fourier features affects generalization, and NTK-aware Scaling designs its frequency adjustments from that angle.

7.4 YaRN improvements

YaRN (Yet another RoPE extensioN) adds two important improvements on top of NTK-aware Scaling:

First: a segmented strategy. YaRN splits dimensions into three frequency groups and treats each differently:

• High-frequency dims (wavelength shorter than original training length / some factor): no scaling; keep original RoPE behavior
• Low-frequency dims (wavelength longer than training length × some factor): full linear interpolation scaling
• Mid band: smooth transition between “no scale” and “full scale”

This segmentation is finer than NTK-aware’s continuous adjustment and better extends long-range ability while keeping short-range resolution.

Second: temperature correction for Attention distribution. After context grows, Attention score distributions shift (more tokens in the sum). YaRN introduces a temperature tied to the extension ratio to correct that shift.

7.5 Dynamic NTK

The methods above need a target length and scale params before inference. In real deployments, input length varies — some requests are short, some long. Scaling short requests for length extension is unnecessary and can hurt quality.

Dynamic NTK’s idea: adjust the base from the current actual sequence length.

During inference, current length $L_{current}$ grows as tokens are generated. If $L_{current} \leq L_{train}$, use the original base (no scaling); if $L_{current} > L_{train}$, recompute base from $L_{current} / L_{train}$.

Benefits:

• Short sequences stay untouched at original precision
• Long sequences get an appropriate scale automatically
• No need to know the max length ahead of time

---

8. AI Infra engineering practice

Positional encoding choices and implementations affect AI Infra engineering directly. This section covers a few important links.

8.1 Fusing RoPE into CUDA kernels

RoPE’s core op is element-wise multiply-add on Q and K (§5.5). Compute-wise it is light ($O(N \cdot d)$), far smaller than Attention matmuls ($O(N^2 \cdot d)$). But as a standalone CUDA kernel, launch overhead and extra memory traffic become the bottleneck.

So in high-performance inference engines, RoPE is usually fused into the front of the Attention kernel:

1. Fuse with QKV projection: after the QKV linear projections, rotate Q and K in the same kernel so intermediate results are not written back to HBM and reread
2. Fuse with FlashAttention: in FlashAttention’s tiling loop, once each block’s Q and K sit in SRAM, apply RoPE in place, then run Attention immediately

Benefits of fusion:

• One fewer kernel launch (saves about 5-10 µs of launch overhead; in Decode, when each step handles one token, that overhead is a real fraction)
• One fewer HBM round trip (especially important for memory-bound Decode)

Also, RoPE trig values ($\cos(m\theta_i)$ and $\sin(m\theta_i)$) are usually precomputed and cached as a lookup table so trig is not recomputed every inference. In mainstream implementations this table is an ordinary tensor in GPU global memory (HBM) — for example vLLM’s cos_sin_cache, shape (max_position, rotary_dim); the kernel reads the matching row from global memory by each token’s position.

8.2 Long context and the KV Cache

How far positional encoding can extrapolate directly sets how long a context the model can support, and longer context puts heavy pressure on KV Cache management.

Take LLaMA-2-7B (numbers from the introductory Transformer article):

From 4K to 128K, KV Cache memory grows 32×. That means:

• A single request can exhaust memory: 128K context alone needs 64 GB of KV Cache, near a full A100
• Concurrency collapses: even with enough memory, simultaneous requests drop sharply
• PagedAttention matters more: with long context, dynamic KV Cache growth and fragmentation get worse and need finer memory management

So long-context positional encoding (§7.2–7.5) is not only a model-capability issue; it is a system-level challenge. Engineering has to solve all of:

1. Positional encoding: pick a suitable extrapolation/interpolation scheme (YaRN, Dynamic NTK, etc.) so Attention quality holds on long sequences
2. KV Cache management: paging, KV Cache quantization (FP8/INT4), eviction policies, and related techniques to control memory
3. Compute: FlashAttention’s $O(N)$ memory profile is critical for long context (otherwise the $O(N^2)$ Attention matrix eats enormous memory)

8.3 Engineering comparison of positional encoding schemes

From an engineering view, RoPE became the de facto standard not only for elegant theory but for engineering friendliness: no extra params (no impact on model parallelism), seamless fusion into existing Attention kernels, and a mature ecosystem of long-context extensions.

---

🎯 Self-check list

After reading this article, check whether you really understand positional encoding:

• Can explain why raw Self-Attention cannot tell word order apart, and briefly state what permutation equivariance means mathematically
• Can distinguish absolute vs relative positional encoding design, and name representative schemes for each
• Can write the Sinusoidal formulas from memory and explain what different frequency dims do
• Can use angle-addition formulas to show why relative position in Sinusoidal encoding comes from a linear transform
• Can state RoPE’s core op (2D rotation) and derive why the rotated inner product depends only on relative position difference
• Can write RoPE’s element-wise implementation and explain why cos/sin are precomputed
• Can state the core differences among RoPE, Sinusoidal, Learned PE, and ALiBi
• Can explain why NTK-aware Scaling beats uniform Position Interpolation (keep high frequency, stretch low frequency)
• Can estimate long-context impact on KV Cache memory and name at least two mitigation strategies
• Can explain why RoPE is usually fused into the Attention kernel

📚 References

Papers

• Attention Is All You Need (Vaswani et al., 2017)：https://arxiv.org/abs/1706.03762 — original Transformer paper; source of Sinusoidal encoding
• RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021)：https://arxiv.org/abs/2104.09864 — RoPE rotary position embedding
• Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (Press et al., 2022)：https://arxiv.org/abs/2108.12409 — ALiBi
• Extending Context Window of Large Language Models via Positional Interpolation (Chen et al., 2023)：https://arxiv.org/abs/2306.15595 — Position Interpolation (PI)
• YaRN: Efficient Context Window Extension of Large Language Models (Peng et al., 2023)：https://arxiv.org/abs/2309.00071 — YaRN long-context extension
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin et al., 2019)：https://arxiv.org/abs/1810.04805 — learned positional encoding
• Language Models are Unsupervised Multitask Learners (Radford et al., 2019)：https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf — GPT-2

Blogs and tutorials

• Jianlin Su: Rotary Position Embedding that drives researchers to despair：https://kexue.fm/archives/8265 — RoPE author’s Chinese blog, derivation from scratch
• Jianlin Su: The road of Transformer upgrades：https://kexue.fm/archives/8130 — positional encoding series
• Eleuther AI: Rotary Embeddings：https://blog.eleuther.ai/rotary-embeddings/ — detailed English explanation of RoPE
• The Illustrated Transformer (Jay Alammar)：https://jalammar.github.io/illustrated-transformer/ — illustrated Transformer intro (includes positional encoding visuals)