长上下文推理：RoPE 扩展、位置插值与 Ring Attention

注意力计算的复杂度：
□ 标准注意力：O(N² · d)
  - N = 序列长度，d = head dim
  - 8K → 64M 次点积；128K → 16G 次（256 倍）
□ 显存：注意力矩阵本身 O(N²)
  - 128K 的注意力矩阵（FP16）= 128K² × 2 bytes = 32 TB
  - 不可能物化 → 必须分块（FlashAttention）

KV Cache 的复杂度：
□ 随序列线性增长：O(N · layers · kv_heads · d · 2)
□ 128K 上下文的 KV Cache 常常比模型权重还大
□ 长上下文的第一瓶颈往往是 KV Cache 而非注意力

长上下文的三个拦路虎：
□ 位置编码外推：训练在 8K，推理到 128K，位置信号失真
□ 计算量：注意力 O(N²) 增长
□ 显存：KV Cache 线性增长 + 注意力矩阵（分块后缓解）

对应解法：
□ 位置外推 → RoPE 扩展 / 位置插值（PI、NTK、YaRN）
□ 计算 → FlashAttention 分块 + 稀疏/滑窗注意力
□ 显存 → PagedAttention + KV 量化 + 序列并行

RoPE 的做法：
□ 把 query/key 向量两两分组（成对）
□ 按位置 m 旋转每对：(x1, x2) → 旋转 θ_m 角度
□ 旋转角随位置线性增长 → 相对位置编码
□ 不同维度用不同频率：θ_i = base^(-2i/d)
  - 低维（i 小）：高频，转得快
  - 高维（i 大）：低频，转得慢

# RoPE 的频率与旋转
import torch

def rope_freqs(dim, base=10000.0):
    # inv_freq: 每个维度对的角速度
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    return inv_freq

def apply_rope(x, pos, inv_freq):
    freqs = torch.outer(pos, inv_freq)      # (seq, dim/2)
    cos, sin = freqs.cos(), freqs.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin,
                        x1 * sin + x2 * cos], dim=-1).flatten(-2)

为什么长上下文会失效：
□ 训练在 8K → 见过的最大位置是 8192
□ 推理到 128K → 位置 100000 从未训练过
□ 高频维度转了好几圈（还行）
□ 低频维度没转够一圈（外推到未训练的角度）→ 失真
□ 结果：困惑度飙升、长文档理解崩坏

核心思路对比：
□ 位置插值 PI（Position Interpolation）：
  - 把长位置「压缩」回训练范围
  - pos' = pos × (L_train / L_target)
  - 128K 压回 8K：pos' = pos / 16
  - 缺点：高频维度被压缩 → 局部区分度下降
□ NTK 感知插值：
  - 不线性压缩，而是按频率「非线性」缩放 base
  - 高频维度几乎不动，低频维度多缩放
  - 保住局部精度 + 扩展全局范围
□ YaRN：
  - 在 NTK 基础上引入温度因子（attention scaling）
  - 分频段处理：高频不动、中频插值、低频外推
  - 少量微调即可外推到很长上下文

PI 与 NTK 的直觉：
□ PI：把 128K 的尺子压缩成 8K → 刻度更密但精度降
□ NTK：把尺子「拉长」，低频刻度稀疏、高频刻度不变
  - base 从 10000 提到 500000 甚至 1000000
  - 周期变长 → 低频维度覆盖更长范围

YaRN 的三段式：
□ 高频（周期 < 原训练长度）：不插值，直接外推
□ 中频：线性插值（PI）
□ 低频（周期 > 目标长度）：保持外推
□ 再加 1/sqrt(t) 温度缩放补偿注意力熵

# NTK-aware 缩放 base 的直觉公式
import math

def ntk_base(orig_base, orig_len, target_len, dim):
    # 按维度比例放大 base
    scale = target_len / orig_len
    return orig_base * (scale ** (dim / (dim - 2)))

# 例：base=10000, 8K → 128K, dim=128
# → base ≈ 10000 × 16^(128/126) ≈ 170000 量级

标准注意力的显存灾难：
□ S = Q·Kᵀ  → N×N 矩阵
□ 128K 下 = 128K × 128K × 2 bytes = 32 TB
□ 根本无法物化 → 必须分块

FlashAttention 的分块思路：
□ 把 Q 切成块，逐块与 K/V 计算
□ 用 online softmax 累积（不物化完整 S）
□ 显存从 O(N²) 降到 O(N)（只存 Q/K/V 块）
□ 计算仍是 O(N²)，但访存高效 → 实际更快

长上下文的额外优化：
□ 滑窗注意力（Sliding Window）：只看最近 W 个 token
  - 复杂度 O(N·W)，适合局部依赖
  - Mistral 的 SWA 是代表
□ 稀疏注意力：只算重要位置对
□ 分块 + 重计算：用算力换显存
□ 注意力 sink：保留开头若干 token 的注意力

FlashAttention 版本差异：
□ FA-1：分块 + online softmax，显存 O(N)
□ FA-2：优化并行与工作划分，速度提升
□ FA-3：面向 Hopper（H100）的 TMA/WGMMA 优化
□ 长上下文首选 FA-2/FA-3（显存与速度双优）

Ring Attention 的核心：
□ 把长序列按卡切成块（每卡一块 Q、K、V）
□ 各卡环形传递 K/V 块（ring 拓扑）
□ 每卡用自己的 Q 块与「流经」的所有 K/V 块计算
□ 计算与通信重叠：算当前块时，传下一块
→ 注意力计算被分布到多卡，显存线性下降

环形通信示意（4 卡，序列切 4 块）：
step 0: 卡0 用 Q0 算 K0/V0，同时把 K0/V0 发给卡1
step 1: 卡0 收到 K3/V3（来自卡3），算 Q0×K3/V3
       同时转发 K3/V3 给卡1 ...
□ 每卡每步只存一块 K/V → 显存 O(N/卡数)
□ 通信量 ∝ K/V 大小 × 卡数（可重叠隐藏）

与 FlashAttention 的结合：
□ 每步内部的块间计算仍用 FlashAttention
□ online softmax 跨步累积（因为 K/V 分多步到达）
□ 结果是「Ring Attention = 序列并行 + FlashAttention」
□ 支持训练（反向需要反向 ring 通信）

其他序列并行：
□ DeepSpeed-Ulysses：按 head 切分（all-to-all）
□ Megatron 序列并行：配合 TP 切序列
□ 选择：Ring 按序列切（省显存）、Ulysses 按 head 切（通信少）

KV Cache 大小估算：
□ 公式：2 × layers × kv_heads × head_dim × seq_len × batch × bytes
□ Llama-3-8B（GQA，kv_heads=8，head_dim=128，32 层）：
  - 每 token KV = 2 × 32 × 8 × 128 × 2 bytes = 128 KB
  - 128K 序列 × 1 batch = 16 GB（仅 KV！）
  - 8 并发 → 128 GB → 远超单卡
□ 长上下文的显存瓶颈是 KV Cache，不是权重

KV Cache 优化手段：
□ PagedAttention：分块管理，消除碎片（见 vLLM）
□ KV 量化：FP16 → FP8/INT8 → 直接砍半
□ GQA/MQA：减少 kv_heads → 从源头减少
□ 滑窗 + sink：只保留部分 KV
□ 序列并行（Ring）：KV 分布到多卡
□ 前缀共享：公共前缀的 KV 复用（见前缀缓存篇）

显存预算（128K 上下文，A100 80 GB）：
□ 权重（FP16）     : 16 GB
□ KV Cache（单请求）: 16 GB
□ 激活与工作区     : 8 GB
→ 单请求已用 40 GB，并发能力极其有限
→ 必须靠 KV 量化 + 分页 + 序列并行才能上并发

两条路线：
□ 纯外推（Training-free）：
  - 直接改 RoPE base / 用 NTK/YaRN 缩放
  - 无需训练，但外推长度有限（通常 2~4 倍）
  - 质量随外推倍数衰减
□ 继续训练（Continued Pretraining）：
  - 用长文档数据继续训（从 8K → 32K → 128K）
  - 质量最好，但需要数据与算力
  - 通常配合位置插值，分阶段扩展

分阶段扩展的常见做法：
□ 阶段 1：8K → 32K，用 PI + 少量长数据
□ 阶段 2：32K → 128K，用 YaRN + 更多长数据
□ 每阶段验证长文任务（检索、摘要、QA）
□ 避免一步到位（一步跳到 128K 通常崩）

训练侧的配合手段：
□ 长文档数据构造（书籍、代码库、长对话）
□ 位置编码的课程学习（逐步加长）
□ 注意力 sink 的训练（让开头 token 常驻）
□ 数据打包（packing）要避免跨样本注意力污染

性能特征：
□ 首 token 延迟（TTFT）随长度显著上升（prefill 是 O(N²)）
□ 解码延迟相对稳定（每步只算 1 个 token）
□ 长上下文的服务瓶颈在 prefill 与 KV 显存
□ 吞吐随长度下降（KV 占用挤占并发）

生产配置示例（vLLM 类引擎）：
--max-model-len 131072
--rope-scaling '{"type":"yarn","factor":16.0,"original_max_position_embeddings":8192}'
--kv-cache-dtype fp8
--enable-chunked-prefill
--max-num-seqs 8

成本权衡：
□ 长度 × 成本：128K 请求的成本远高于 8K（KV + 计算）
□ 分层服务：短请求走标准实例，长请求走长上下文实例
□ 前缀缓存：公共前缀复用 KV → 大幅省 prefill
□ 上下文压缩：RAG 只注入相关片段而非全文
□ 滑窗 + 摘要：超长文档滚动摘要而非全量注入

高频踩坑：
□ 直接改 max_position_embeddings 不改 RoPE → 位置失真，输出崩坏
□ RoPE scaling 配置与训练时不一致 → 长文理解错误
□ 忘记 KV 量化就开 128K → OOM
□ 长上下文压测只用短请求 → 上线后 TTFT 爆炸
□ 注意力 sink 未保留 → 开头信息被滑窗挤出，质量下降
□ 数据打包时跨样本注意力污染 → 训练侧质量受损
□ 位置插值后未验证「中间位置」（lost in the middle）
□ chunked prefill 分块不当 → 首 token 延迟不降反升

排查清单：
□ 长文检索任务（needle-in-haystack）：验证位置编码是否生效
□ 中间位置测试：验证 lost-in-the-middle 程度
□ TTFT 随长度曲线：是否符合预期
□ KV 显存占用曲线：是否线性
□ 与短上下文输出对比：位置扩展是否损伤短文能力

needle-in-haystack 示意：
□ 在长文档不同位置插入「针」（特定事实）
□ 让模型检索 → 记录不同位置的命中率
□ 健康曲线：全位置高命中
□ 病态曲线：开头/结尾高、中间低（lost in the middle）