AI Engineer Exhaustive
EasyPlusLLaMA

RoPE Frequency Table

LLaMA

Easy

Precompute cos/sin frequency tables

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

RoPE needs a phase (m\theta_i) for every position (m) and every pair-index (i). Llama 3 materializes those phases once as a complex table freqs_cis and reuses it on every forward. The herd paper specifies only the base (\theta=500{,}000) (Table 3). The table construction — which axis is position, how (\theta_i) is spaced, complex polar packing, and the extra factor-of-two length — is official code (precompute_freqs_cis), not a paper equation.

This note is the table, not the rotate. Application is apply_rotary_emb.

How it works

precompute_freqs_cis(dim, end, theta) builds a matrix in (\mathbb{C}^{\texttt{end}\times(\texttt{dim}/2)}). Here dim is head width (d_h=d/n_q), not the residual width. Transformer.__init__ calls it as

precompute_freqs_cis(params.dim // params.n_heads,
                     params.max_seq_len * 2,
                     params.rope_theta)

Inverse frequencies. For (i = 0,2,4,\ldots,d_h-2) (then truncated to (d_h/2) values):

\theta_i = \theta^{-i/d_h} = \theta^{-2k/d_h}\quad(k=0,\ldots,d_h/2-1).

In code: 1.0 / (theta ** (arange(0, dim, 2)[: dim//2].float() / dim)). Default theta on the helper is (10{,}000) (Su’s original); Llama 3 passes rope_theta=500000 from ModelArgs / Table 3. Larger (\theta) stretches the wavelengths so distant positions remain distinguishable — the reason the paper raised the base for long context.

Positions. (t = (0,1,\ldots,\texttt{end}-1)) in float32. The outer product (t\otimes(\theta_i)_i) is the real phase table (m\theta_i).

Complex cis. torch.polar(ones_like(freqs), freqs) yields

\mathrm{cis}[m,i] = e^{\mathrm{i}\, m\theta_i} = \cos(m\theta_i) + \mathrm{i}\sin(m\theta_i)

as complex64. Storing cis rather than separate cos/sin tensors is why apply_rotary_emb can be a single complex multiply.

Length. end = max_seq_len * 2, so the table is twice the configured context. The extra rows are unused on a strictly in-range call; they exist so a slice freqs_cis[start_pos:start_pos+S] still has room near the end of a full-length sequence. The paper does not mention this doubling — it is a repo choice.

Use. Transformer.forward moves the buffer to the activation device and slices rows [start_pos : start_pos+S]. reshape_for_broadcast then checks that this slice’s shape is (S, d_h/2) against the complex view of (Q).

Specified vs inferred: (\theta=500{,}000) is paper. The (2i/d_h) schedule is Su et al. The * 2 length, polar, and dim = d/n_heads argument are official-code details. generation.py’s max_seq_len <= 8192 assert is the April 2024 inference wrapper; the July paper’s 128K window would require a longer table built with a larger end.

Official code

llama/model.py — precompute_freqs_cis, the call in Transformer.__init__, and the slice in Transformer.forward. ModelArgs.rope_theta defaults to 500000. There is no other RoPE table builder in this repository.

Watch-outs

Sources