AI Infra
SwiGLU
Swish 门控线性单元是一种带有门控机制的前馈神经网络结构，将输入分别送入两个不同的线性变换分支，其中一个分支经过 SiLU 激活函数产生门控信号，再将两个分支的输出元素相乘，最后通过一个线性层将结果映射回原始维度。数学表达式为
其中 \(W_{\mathrm{gate}},W_{\mathrm{up}},W_{\mathrm{down}}\) 分别表示三个线性层的权重矩阵。计算流程如下
graph LR
A(输入x
维度：hidden_size)
B(gate_proj
线性投影)
C(up_proj
线性投影)
D(SiLU
非线性激活)
E(逐元素相乘
维度：intermediate_size)
F(down_proj
映射回hidden_size)
G(输出
维度：hidden_size)
A-->B-->D-->E
A-->C-->E-->F-->G- 门控：让网络根据当前输入，动态调节不同特征的作用
- 分支特征：两个分支权重参数不同，从同一个输入中学习到不同的特征组合。一条路径生成待处理的特征，另一条生成对这些特征的调节信号。
- 逐元素相乘：一个分支产生的数值能够直接控制另一个分支对应特征的幅度。
RoPE
注意力里的点积本身并不知道 token 的先后顺序。旋转位置编码（Rotary Position Embedding，RoPE）的做法是：在计算注意力之前，按 token 的位置旋转 Query 和 Key。每个位置使用不同角度；两个旋转后向量的点积因此能够体现它们的相对距离。
从二维旋转理解相对位置
先看二维向量。位置为 \(m\) 时，将向量旋转 \(m\theta\)：
注意力分数来自 Query 和 Key 的点积，因此：
虽然 \(q_m\) 和 \(k_n\) 各自使用绝对位置 \(m\)、\(n\)，点积中的位置关系却以差值 \(n-m\) 出现。这不意味着注意力分数只由距离决定：原始 \(q\)、\(k\) 的内容仍然参与计算。RoPE 改变的是位置进入点积的方式。
Qwen2.5-0.5B-Instruct
一个 Attention Head 不止二维。如果所有维度都使用同一个 \(\theta\)，信息表达能力有限， RoPE 为不同维度定义不同频率。把维度为 \(d\) 的向量分成 \(\frac{d}{2}\) 对，每对使用不同的旋转频率。第 \(i\) 对的角度为 \(m\theta_i\)，其中（步进为 \(\frac{2}{d}\)）：
位置 \(m\) 对应的角度为 \(m\theta_i\)，所以第 \(i\) 个二维平面的旋转矩阵为
整个 RoPE 可以理解为一个大矩阵
当 \(i\) 很小时，\(\theta_i\) 比较大，位置增加，角度变化明显，容易区分近距离变化，表示高频特征；反之，\(i\) 很大时，位置增加角度变化微小，表示低频位置特征。
在本文对应的 Qwen2.5-0.5B-Instruct 配置 中，hidden_size 为 \(896\)、Query Head 数为 \(14\)、KV Head 数为 \(2\)，所以 head_dim 为 \(896/14=64\)；rope_theta（即上式的 base）为 \(1\,000\,000\)。每个 Head 有 \(32\) 个频率。GQA 让多个 Query Head 共享一个 KV Head，但不改变各 Head 内的旋转公式：RoPE 分别作用于 Q 和 K，不作用于 V。Hugging Face 的 Qwen2 实现 也采用这一顺序。
input X
│
├── W_Q ── Q ── RoPE ── Q'
│
├── W_K ── K ── RoPE ── K'
│
└── W_V ── V
│
Q' × K'^T ── softmax
│
× V
│
output
注：Qwen2.5 所用实现则把向量分成前后两半，配对 \((x_i,x_{i+32})\)。例如一个 \(64\) 维 Head 的第 \(0\) 维与第 \(32\) 维配对，而不是与第 \(1\) 维配对。这两种布局都能表达二维旋转，但对同一组模型权重不可随意互换。
对照代码实现
下面是项目中 RotaryEmbedding 的核心代码。约定 Q、K 的形状分别为 (batch, seq, q_heads, head_dim) 和 (batch, seq, kv_heads, head_dim)，positions 为 (batch, seq)。单批次打包 token 时也可分别使用 (tokens, heads, head_dim) 与 (tokens,)。
import torch
from torch import nn
def rotate_half(x: torch.Tensor) -> torch.Tensor:
left, right = x.chunk(2, dim=-1)
return torch.cat((-right, left), dim=-1)
class RotaryEmbedding(nn.Module):
"""Qwen2.5 RoPE for Q/K shaped (..., seq, heads, head_dim)."""
def __init__(self, head_dim: int, base: float):
super().__init__()
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
freqs = (
positions.to(device=q.device, dtype=torch.float32)[..., None]
* self.inv_freq
)
angles = torch.cat((freqs, freqs), dim=-1)
cos = angles.cos().to(q.dtype).unsqueeze(-2)
sin = angles.sin().to(q.dtype).unsqueeze(-2)
return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
inv_freq 保存 \(32\) 个 \(\theta_i\)，它由配置直接计算得出，因此注册为不写入 checkpoint 的 buffer。位置 \(m\) 与 inv_freq 相乘得到每对维度的 \(m\theta_i\)；torch.cat((freqs, freqs), dim=-1) 将角度扩展为 \(64\) 维，使前后两半使用相同的频率。rotate_half 把每对 \((a,b)\) 变成 \((-b,a)\)；若以 \(c\)、\(s\) 表示余弦和正弦值，\(x\odot c+\operatorname{rotate\_half}(x)\odot s\) 正好实现二维旋转。
对批量输入而言，cos、sin 的形状是 (batch, seq, 1, 64)。那个大小为 \(1\) 的 Head 维使其同时广播到 \(14\) 个 Query Head 和 \(2\) 个 KV Head，不需要把 K 复制成 \(14\) 个 Head。位置与三角函数先用 FP32 计算，再转回 Q 的 dtype；在 BF16 推理时，这能避免直接用低精度计算角度。RoPE 这里只负责旋转，缩放点积、causal mask 和 KV Cache 读写属于 Attention 的其他部分。
Prefill、Decode 与 KV Cache
位置从 \(0\) 开始。假设 prompt 长度为 \(16\)，prefill 的位置是 \(0\) 到 \(15\)；下一次 decode 输入新 token 时，其位置是 \(16\)，也就是 进入这一步之前已有的上下文 token 数。对于批量请求，每条序列应使用自己的上下文长度，不能统一从 batch 内的 \(0\) 重新编号。
当位置为 \(0\) 时，角度全为 \(0\)，所以 \(\cos(0)=1\)、\(\sin(0)=0\)，旋转结果应与输入相同。使用 KV Cache 时，写入缓存的 K 应是 已旋转的 K；后续 decode 只旋转当前新 token 的 Q、K，读取旧 K 时不再旋转一次。否则，同一段上下文在完整前向与增量前向中会得到不同的注意力分数。
验证实现
项目的测试以 Hugging Face 的 Qwen2RotaryEmbedding 和 apply_rotary_pos_emb 为基线，比较 FP32、BF16 以及 CPU、CUDA 上的 Q/K 输出。测试包含位置 \(0\)、\(1\)、\(15\)、\(16\)、\(1024\)、不同序列使用不同位置、packed token 形状，还检查 decode 的位置 \(16\) 与完整 prefill 中同一位置相等。运行命令为 python -m pytest tests/test_rotary_embedding.py -q。