RoPE 位置编码算法详解

# 第一个参数是绝对值（模），第二个参数是角度
torch.polar(abs, angle, *, out=None) → Tensor

# 假设 freqs = [x, y], 则 torch.polar(torch.ones_like(freqs), freqs) 
# = [cos(x) + sin(x)j, cos(y) + sin(y)j]
>>> angle = torch.tensor([np.pi / 2, 5 * np.pi / 4], dtype=torch.float64)
>>> z = torch.polar(torch.ones_like(angle), angle)
>>> z
tensor([ 6.1232e-17+1.0000j, -7.0711e-01-0.7071j], dtype=torch.complex128)
>>> a = torch.tensor([np.pi / 2], dtype=torch.float64) # 数据类型必须和前面一样
>>> torch.cos(a)
tensor([6.1232e-17], dtype=torch.float64)

# 第一个参数是输入张量
# 第二个参数是重复次数
# dim: 沿着该维度重复元素。如果未指定维度，默认会将输入数组展平成一维，并返回一个平坦的输出数组。
torch.repeat_interleave(input, repeats, dim=None, *, output_size=None) → Tensor

def compute_theta(dim: int, base: float = 10000.0, device: torch.device = torch.device('cpu')) -> torch.Tensor:
    """
    计算旋转位置编码中的 Theta 角度值。

    参数：
    - d (int): 嵌入向量的维度（必须为偶数）。
    - base (float): 基础频率参数, 默认为10000.0。
    - device (torch.device): 计算设备, 默认为CPU。

    返回：
    - torch.Tensor: 包含Theta值的1D张量, 形状为 [d/2]。
    """
    if dim % 2 != 0:
        print("嵌入维度 dim 必须为偶数")
    i = torch.arange(1, (dim//2) + 1, dtype=torch.float32, device=device)
    theta_i = base ** (-2*(i - 1) / dim)

    return theta_i

def precompute_freqs_cis(dim: int, seq_len: int, base: float = 10000.0, device: torch.device = torch.device('cpu')):
    theta = compute_theta(dim, base, device) # theta 角度值序列，向量, 大小为 dim // 2
    m = torch.arange(seq_len, device=device) # # token 位置值序列，向量，大小为 seq_len
    m_theta = torch.outer(m, theta) # 所有 token 位置的所有 Theta 值范围, 矩阵，尺寸为 [seq_len, dim // 2]
    freqs_cis = torch.polar(torch.ones_like(m_theta), m_theta) # e^{i*m*\theta}，本质上是旋转矩阵
    return freqs_cis

def reshape_for_broadcast(freqs_cis, x):
    ndim = x.ndim
    assert ndim >= 2
    assert freqs_cis.shape == (x.shape[1],x.shape[-1]), "the last two dimension of freqs_cis, x must match"
    shape = [d if i==1 or i==ndim-1 else 1 for i,d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, device: torch.device = torch.device('cpu')):
    """
    参数:
        - x_q(torch.Tensor): 实际上是权重 W_q * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]
        - x_k(torch.Tensor): 实际上是权重 W_k * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]
        - freqs_cis (torch.Tensor): 频率复数张量, 形状为 [max_seq_len, head_dim]
    返回:
        - Tuple[torch.Tensor, torch.Tensor]: 旋转编码后的查询和键张量
    """
    # 实数域张量转为复数域张量
    xq_reshape = xq.reshape(*xq.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] 
    xk_reshape = xk.reshape(*xk.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] 
    xq_complex = torch.view_as_complex(xq_reshape) # 复数形式张量
    xk_complex = torch.view_as_complex(xk_reshape) # 复数形式张量

    # 旋转矩阵（freqs_cis）的维度在序列长度（seq_len，维度 1）和头部维度（head_dim，维度 3）上需要与嵌入的维度一致。
    # 此外，freqs_cis 的形状必须与 xq 和 xk 相匹配，因此我们需要将 freqs_cis 的形状从 [max_seq_len, head_dim] 调整为 [1, max_seq_len, 1, head_dim]。
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_complex) # [max_seq_len, 1, 1, dim // 2]

    # 应用旋转操作，并将结果转回实数域。# flatten(2) 将后面两个维度压成一个维度
    xq_out = torch.view_as_real(xq_complex * freqs_cis).flatten(3) 
    xk_out = torch.view_as_real(xk_complex * freqs_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

import torch

def build_rope_cache(
    batch_size: int,
    seq_len: int,
    head_dim: int,
    base: int = 10000,
    max_position_embeddings: int = 32768,
    dtype: torch.dtype = torch.float32,
):
    """
    构建 RoPE 频率缓存.

    Args:
        batch_size (int): 批大小 B
        seq_len (int): 序列长度 S
        head_dim (int): 每个注意力头维度 H (必须是偶数)
        base (int, optional): 对数基数(常用 10000). Defaults to 10000.
        max_position_embeddings (int, optional): 支持的最长位置. Defaults to 32768.
        dtype (torch.dtype, optional): 输出张量 dtype. Defaults to torch.float32.

    Returns:
        positions_cpu (LongTensor): [B, S] 随机位置 id
        cos (Tensor): [max_position_embeddings, head_dim//2] 余弦频率
        sin (Tensor): [max_position_embeddings, head_dim//2] 正弦频率
    """
    # 1) 为每个 batch 生成位置 id（演示用随机，实际推理时应为 torch.arange）
    positions_cpu = torch.randint(0, seq_len, (batch_size, seq_len), dtype=torch.long)

    # 2) 计算逆频率 inv_freq: [H/2]
    inv_freq = 1.0 / (
        base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
    )

    # 3) 生成全局位置索引 t: [max_position_embeddings]
    t = torch.arange(max_position_embeddings, dtype=torch.float32)

    # 4) 外积得到每个位置对应的旋转角度 freqs: [max_position_embeddings, H/2]
    freqs = torch.einsum("i,j->ij", t, inv_freq)

    cos = freqs.cos().to(dtype)
    sin = freqs.sin().to(dtype)

    return positions_cpu, cos, sin

# ========== 简单单元测试 / demo ========== #
torch.manual_seed(0)
B, S, H = 2, 4, 8
pos, cos, sin = build_rope_cache_fixed(B, S, H)

print("positions shape:", pos.shape)
print("cos shape:", cos.shape, " dtype:", cos.dtype)
print("sin shape:", sin.shape, " dtype:", sin.dtype)
print("pos[0]:", pos[0])
print("cos[0, :4]:", cos[0, :4])

positions shape: torch.Size([2, 4])
cos shape: torch.Size([32768, 4])  dtype: torch.float32
sin shape: torch.Size([32768, 4])  dtype: torch.float32
pos[0]: tensor([0, 3, 1, 0])
cos[0, :4]: tensor([1., 1., 1., 1.])

def apply_rotary_pos_embedding(
    query: torch.Tensor,          # [B, S, num_heads, head_dim]
    key:   torch.Tensor,          # [B, S, num_heads, head_dim]
    cos:   torch.Tensor,          # [L_max, head_dim/2]  — 由 build_rope_cache 预先生成
    sin:   torch.Tensor,          # [L_max, head_dim/2]
    positions_ids: Optional[torch.Tensor] = None,  # [B, S] 或 [S]
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    将 RoPE 旋转应用到 query / key。
    """
    # 1) 若未显式传入位置 id，则默认按 0..S-1 顺序
    if positions_ids is None: 
        positions_ids = torch.arange(cos.size(0), device=cos.device)

    # 2) 取出对应行；reshape(-1) 兼容批维
    # positions_ids 可能是一维，也可能是二维 ⇒ 统一成 [...,1,-1]
    cos_sel = cos.index_select(0, positions_ids.reshape(-1))
    sin_sel = sin.index_select(0, positions_ids.reshape(-1))

    # 3) 恢复批次形状，并在 head 维前插入空维，便于广播
    target_shape = (*positions_ids.shape, 1, -1)    # [B,S,1,head_dim/2] 或 [S,1,head_dim/2]
    cos_sel = cos_sel.view(target_shape)
    sin_sel = sin_sel.view(target_shape)

    # 若缺少 batch 维，自动 unsqueeze 再广播
    if cos_sel.dim() == 3:                          # [S,1,H/2] 情况
        cos_sel = cos_sel.unsqueeze(0)              # → [1,S,1,H/2]
        sin_sel = sin_sel.unsqueeze(0)

    # 4) 内部旋转函数：偶/奇维度各自为实部/虚部
    def _rotate(t: torch.Tensor) -> torch.Tensor:
        """
        将 (x0,x1,x2,x3,...) 视为复数 (x0+i·x1 , x2+i·x3, ...)
        与 (cos+i·sin) 相乘，再拆成偶/奇分量。
        """
        t1, t2 = t[..., ::2], t[..., 1::2]        # 偶、奇
        rot1 = t1 * cos_sel - t2 * sin_sel        # 实部
        rot2 = t2 * cos_sel + t1 * sin_sel        # 虚部
        return torch.stack((rot1, rot2), dim=-1).flatten(-2)

    # 5) 分别旋转 Q / K
    return _rotate(query), _rotate(key)

B, S, num_heads, head_dim = 2, 4, 4, 8, 
# 简易 rope 缓存（完整实现见上一轮 build_rope_cache）
pos = torch.arange(S)
inv_freq = 1 / (10000 ** (torch.arange(0, head_dim, 2) / head_dim))
freqs = torch.outer(pos, inv_freq)
cos, sin = freqs.cos(), freqs.sin()

# 随机生成 q, k 并应用 RoPE
q = torch.randn(B, S, num_heads, head_dim)
k = torch.randn(B, S, num_heads, head_dim)
q_rot, k_rot = apply_rotary_pos_embedding(q, k, cos, sin, pos)

# 断言形状
assert q_rot.shape == (B, S, num_heads, head_dim)
assert k_rot.shape == (B, S, num_heads, head_dim)

# 慢速公式验证单点正确性
def naive_rot(t, i):
    c, s = cos[i], sin[i]
    t1, t2 = t[..., ::2], t[..., 1::2]
    rot = torch.stack((t1 * c - t2 * s, t2 * c + t1 * s), -1).flatten(-2)
    return rot

i = 1  # 第 1 个 token
gold = naive_rot(q[0, i], i)
assert torch.allclose(q_rot[0, i], gold, atol=1e-5)
print("通过！q_rot[0,1] ≈ gold，形状一致")