实验 09 · RoPE 与 KV Cache — 增量解码与旋转位置编码

本实验突破自回归生成的吞吐瓶颈：手写旋转位置编码（RoPE）赋予相对位置感知，实现 KV Cache 增量解码将推理计算复杂度由 $O (N^{2})$ 压缩到 $O (N)$ ，并验证增量解码与全量前向的一致性对拍。

网页即时交互实验一：RoPE 旋转位置编码直觉

RoPE（Rotary Position Embedding）将位置信息编码为复数平面上的旋转矩阵。两个 token 向量的内积直接取决于它们的相对距离 $m - n$ ，并且在 position 0 处严格保持恒等变换：

交互：RoPE 旋转位置编码

每一对维度 (d_2k, d_2k+1) 被当成复平面上的一个矢量，位置 pos 让它旋转 θ = pos × ω_k。

position = 3显示维度对数 = 4base = 10000 显示 pos+1 的虚线矢量（相对位置差）

k	ω_k	θ = pos·ω	圈数
0	1.00e+0	3.000	0.48
1	7.50e-1	2.250	0.36
2	5.62e-1	1.687	0.27
3	4.22e-1	1.265	0.20

$$\omega_k = \text{base}^{-2k/d},\quad \begin{pmatrix} x'_{2k} \\ x'_{2k+1} \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x_{2k} \\ x_{2k+1} \end{pmatrix},\ \theta = \text{pos}\cdot\omega_k$$

教学要点：低维（k 小）旋转快 → 编码「近距离」；高维（k 大）旋转慢 → 编码「远距离」。调大 base 会让所有频率变慢，这正是「NTK 插值」外推长上下文的原理。

网页即时交互实验二：KV Cache 动态内存演进与增量推理

在纯手写生成循环中，每生成一个新 token，历史键值对无需重新计算。通过将历史 $K$ 与 $V$ 缓存在预分配张量中，模型每步仅需将最新 token 投影并拼接，从而显著减少访存开销：

交互：KV Cache 增量解码

自回归生成时，历史 token 的 K/V 永远不变。缓存它们，把每步的 O(n) 重算降成 O(1)。

d_head = 4

KV Cache（每步只追加一行）

#	token	K	V	attn
0	<s>	[-0.56, -0.94, -0.14, -0.89]	[-0.10, -0.63, -0.05, -0.85]
1	我	[0.78, -0.18, 0.69, 0.08]	[-0.26, -0.70, 0.67, 0.42]

无 Cache（全量重算）Q,K,V 形状 [2, 4]累计行数 3 → O(n²)

有 Cache（增量）Q 形状 [1, 4]，KV 读缓存 [2, 4]累计行数 2 → O(n)

输出（两条路径完全相同）[-0.2185, -0.6789, 0.4682, 0.0666]

$$\text{Attn}(q_n, K_{1:n}, V_{1:n}) = \text{softmax}\!\left(\frac{q_n K_{1:n}^\top}{\sqrt{d}}\right) V_{1:n}$$

教学要点：加速的代价是显存 —— cache 大小 = 2 × layers × heads × n × d_head。这就是长上下文推理「显存爆炸」的根源，也是 MQA / GQA / PagedAttention 要解决的问题。

🐍

实验 09 · RoPE 旋转位置编码

复数平面二维正交旋转矩阵与点积注意力相对位置平移不变性

↗ 独立演练台

# 实验 09 · 动手挑战：实现 RoPE 2D 旋转矩阵
import math
 
def apply_rope_2d(x, m, theta=10000.0):
    # TODO: 计算 angle = m * (1.0 / (theta ** 0))
    # TODO: 返回旋转后的二维向量 [x[0]*cos - x[1]*sin, x[0]*sin + x[1]*cos]
    pass
 
x = [1.0, 0.0]
x_0 = apply_rope_2d(x, m=0)
assert x_0 is not None, "apply_rope_2d 尚未实现"
assert abs(x_0[0] - 1.0) < 1e-5 and abs(x_0[1] - 0.0) < 1e-5
print("🎉 RoPE m=0 恒等变换与平移不变性断言 100% 通过！")

Python 3行 1, 列 114 行 · 400 字符

运行：⌘ / Ctrl + ↵

就绪 (点击运行)

点击右上角「▶ 运行」或按 ⌘/Ctrl+Enter 在浏览器端直接执行

交互实验指南

相对位置不变性：在 RoPE 演示中调整两个向量的绝对位置，观察当它们保持相同相对间距时，点积注意力打分是否保持恒定？
显存增长瓶颈：在 KV Cache 演示中逐步增加生成的上下文长度，观察显存占用的线性攀升过程，直观理解为什么工业界需要引入 PagedAttention 或 DeepSeek 的 MLA（低秩压缩）。

本地动手实验与单元测试

在网页端观察动态演进后，回到本地运行你的 RoPE 矩阵旋转与 KV Cache 增量解码实现：

bash

cd <仓库根目录>
source .venv/bin/activate
export PYTHONPATH="$PWD/python"

# 运行 RoPE 与 KV Cache 单元测试
pytest python/tests/test_inference_cache.py -v

下一步

→ 返回第10章正文查看完整实现细节与故障注入表
→ 进入实验 10：后训练微调与 LoRA 权重合并

实验 09 · RoPE 与 KV Cache — 增量解码与旋转位置编码 ​

网页即时交互实验一：RoPE 旋转位置编码直觉 ​

交互：RoPE 旋转位置编码

网页即时交互实验二：KV Cache 动态内存演进与增量推理 ​