上下文窗口扩展：RoPE, ALiBi

# 绝对位置编码（正弦/余弦方案）
import numpy as np

def get_sinusoidal_pe(seq_len: int, d_model: int) -> np.ndarray:
    """生成正弦绝对位置编码矩阵 [seq_len, d_model]"""
    pe = np.zeros((seq_len, d_model))
    position = np.arange(seq_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) *
                      -(np.log(10000.0) / d_model))
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    return pe

# 可学习的位置嵌入
import torch
learned_pe = torch.nn.Embedding(num_embeddings=4096,
                                embedding_dim=4096)

# 相对位置编码（Shaw et al. 方案）
import torch

class RelativeAttention(torch.nn.Module):
    def __init__(self, d_model: int, max_rel_pos: int = 128):
        super().__init__()
        self.max_rel_pos = max_rel_pos
        # 相对位置的可学习嵌入
        self.rel_k_emb = torch.nn.Embedding(
            2 * max_rel_pos - 1, d_model)
        self.rel_v_emb = torch.nn.Embedding(
            2 * max_rel_pos - 1, d_model)

    def forward(self, q: torch.Tensor, k: torch.Tensor,
                v: torch.Tensor) -> torch.Tensor:
        # 裁剪相对位置到 [-max_rel_pos+1, max_rel_pos-1]
        rel_pos = self.get_relative_positions(q)
        rel_k = self.rel_k_emb(rel_pos)
        rel_v = self.rel_v_emb(rel_pos)
        # 内容-内容注意力 + 内容-位置注意力
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores += torch.matmul(q, rel_k.transpose(-2, -1))
        weights = torch.softmax(scores, dim=-1)
        return torch.matmul(weights, v + rel_v)

# RoPE 核心实现
import torch

def apply_rotary_emb(
    x: torch.Tensor,          # [batch, heads, seq_len, head_dim]
    freqs_cis: torch.Tensor   # [seq_len, head_dim//2, 2]
) -> torch.Tensor:
    """对 x 应用旋转位置编码"""
    # 将 x 的最后一个维度拆分为实部和虚部对
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
    # 复数乘法 = 旋转变换
    x_complex = torch.view_as_complex(x_)
    freqs_cis = freqs_cis.view_as_complex()
    x_out = torch.view_as_real(x_complex * freqs_cis)
    return x_out.flatten(3).type_as(x)

# 预计算旋转频率
def precompute_freqs_cis(dim: int, end: int,
                         base: float = 10000.0):
    freqs = 1.0 / (base ** (
        torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, dtype=freqs.dtype)
    freqs = torch.outer(t, freqs).float()
    return torch.polar(torch.ones_like(freqs), freqs)

# RoPE 与注意力结合的完整流程
class RotaryAttention(torch.nn.Module):
    def __init__(self, dim: int, n_heads: int,
                 base: float = 10000.0):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.q_proj = torch.nn.Linear(dim, dim)
        self.k_proj = torch.nn.Linear(dim, dim)
        self.v_proj = torch.nn.Linear(dim, dim)
        self.o_proj = torch.nn.Linear(dim, dim)
        self.freqs_cis = precompute_freqs_cis(
            self.head_dim, 4096, base)

    def forward(self, x, positions=None):
        bsz, seq_len, _ = x.shape
        q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
        k = k.view(bsz, seq_len, self.n_heads, self.head_dim)
        # 应用旋转
        q = apply_rotary_emb(q, self.freqs_cis[positions])
        k = apply_rotary_emb(k, self.freqs_cis[positions])
        # 标准缩放点积注意力
        scores = (q @ k.transpose(-2, -1)) / (
            self.head_dim ** 0.5)
        weights = torch.softmax(scores, dim=-1)
        out = (weights @ v).view(bsz, seq_len, -1)
        return self.o_proj(out)

# ALiBi 注意力偏置矩阵
import torch

def get_alibi_slopes(n_heads: int) -> torch.Tensor:
    """计算每个头的 ALiBi 斜率"""
    # 斜率按几何级数分配: 1/2^(8/n) 的幂次
    def get_slopes_power_of_2(n):
        start = 2  (-(2  -(
            torch.log2(torch.tensor(n)) - 3)))
        ratio = start
        return start * ratio ** torch.arange(n)

    if torch.log2(torch.tensor(n_heads)).is_integer():
        return get_slopes_power_of_2(n_heads)
    # 非 2 的幂时，从最近的 2 的幂中选取
    closest = 2 ** int(torch.log2(
        torch.tensor(n_heads)).floor())
    extra = n_heads - closest
    slopes = get_slopes_power_of_2(2 * closest)
    return slopes[1::2][:extra].tolist() + slopes[:closest]

# 构建偏置矩阵
def build_alibi_bias(seq_len: int, n_heads: int):
    slopes = get_alibi_slopes(n_heads)
    # 相对位置矩阵: j - i
    rel_pos = torch.arange(seq_len).unsqueeze(0) - \
              torch.arange(seq_len).unsqueeze(1)
    # 每个头乘以各自的斜率
    return rel_pos.unsqueeze(0) * slopes.view(-1, 1, 1)

# ALiBi + 注意力完整实现
class AlibiAttention(torch.nn.Module):
    def __init__(self, dim: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.scale = self.head_dim ** -0.5
        self.q_proj = torch.nn.Linear(dim, dim)
        self.k_proj = torch.nn.Linear(dim, dim)
        self.v_proj = torch.nn.Linear(dim, dim)
        self.register_buffer(
            "alibi_bias", None, persistent=False)

    def _ensure_alibi(self, seq_len: int):
        if (self.alibi_bias is None or
                self.alibi_bias.shape[-1] < seq_len):
            self.alibi_bias = build_alibi_bias(
                seq_len, self.n_heads).to(
                    self.q_proj.weight.device)

    def forward(self, x):
        bsz, seq_len, _ = x.shape
        self._ensure_alibi(seq_len)
        q = self.q_proj(x).view(
            bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(
            bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(
            bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        scores = (q @ k.transpose(-2, -1)) * self.scale
        # 关键：添加 ALiBi 偏置
        scores = scores + self.alibi_bias[:, :seq_len, :seq_len]
        weights = torch.softmax(scores, dim=-1)
        return (weights @ v).transpose(1, 2).reshape(bsz, seq_len, -1)

# 外推问题可视化：RoPE 频率混叠
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

def visualize_extrapolation(base=10000, d=128,
                            train_len=2048, test_len=8192):
    """展示 RoPE 在外推时的频率问题"""
    dim = 0  # 第一个维度（最高频）
    freq = 1.0 / (base ** (2 * dim / d))

    train_positions = np.arange(train_len)
    test_positions = np.arange(test_len)

    # 旋转角度
    train_angles = train_positions * freq * 2 * np.pi
    test_angles = test_positions * freq * 2 * np.pi

    plt.figure(figsize=(12, 4))
    plt.plot(train_positions, np.sin(train_angles),
             label="训练范围", alpha=0.7)
    plt.plot(test_positions, np.sin(test_angles),
             label="外推范围", alpha=0.7)
    plt.axvline(train_len, color="red", ls="--",
                label="训练边界")
    plt.legend()
    plt.title("RoPE 旋转角度外推示意")
    plt.savefig("rope_extrapolation.png")

# 注意力模式外推误差
def extrapolation_error(model, seq_len):
    """测量模型在外推长度上的性能衰减"""
    ppl_train = evaluate_ppl(model, seq_len=2048)
    ppl_test = evaluate_ppl(model, seq_len=seq_len)
    degradation = (ppl_test - ppl_train) / ppl_train
    print(f"长度 {seq_len}: 困惑度衰减 {degradation:.1%}")

# 外推性 benchmark：PassKey Retrieval
import torch

def passkey_retrieval_test(model, tokenizer,
                           context_len: int,
                           key_pos_ratio: float = 0.5):
    """测试模型在指定长度下检索隐藏 key 的能力"""
    noise = "The grass is green. " * (context_len // 6)
    key = f"PassKey = {torch.randint(0, 10000, (1,)).item()}"
    prompt = noise + key + noise

    inputs = tokenizer(prompt, return_tensors="pt",
                       truncation=True,
                       max_length=context_len)
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=10,
            do_sample=False)
    return tokenizer.decode(outputs[0])

# 评估不同长度
for length in [2048, 4096, 8192, 16384, 32768]:
    result = passkey_retrieval_test(model, tok, length)
    print(f"长度 {length}: 检索准确率 = {'PassKey' in result}")

# NTK-aware 插值实现
import torch
from typing import Tuple

def get_ntk_aware_freqs_cis(
    dim: int, end: int, base: float = 10000.0,
    scale_factor: float = 1.0
) -> torch.Tensor:
    """NTK-aware 缩放的频率计算"""
    # 动态调整 base 参数
    # s_base = base * scale^(dim/(dim-2))
    adjusted_base = base * (scale_factor ** (dim / (dim - 2)))

    freqs = 1.0 / (adjusted_base ** (
        torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

    # 位置索引也被缩放
    t = torch.arange(end, dtype=freqs.dtype) / scale_factor
    freqs = torch.outer(t, freqs).float()

    return torch.polar(torch.ones_like(freqs), freqs)

# 使用示例：将 LLaMA 2048 扩展到 8192
scale = 4.0  # 4x 扩展
new_freqs = get_ntk_aware_freqs_cis(
    dim=128, end=8192, base=10000.0, scale_factor=scale)
print(f"原始 base=10000, 调整后 base={10000.0 * 4**(128/126):.0f}")

# NTK 插值 vs 线性插值对比
def compare_interpolation_methods(
    model, tokenizer, test_seq_len: int = 8192
):
    """比较不同插值方法的性能"""
    results = {}

    # 方法1: 线性插值（所有维度均匀缩放）
    linear_freqs = get_linear_interp_freqs_cis(
        dim=128, end=test_seq_len,
        base=10000.0, scale=4.0)
    results["linear"] = evaluate_ntk(model, linear_freqs)

    # 方法2: NTK-aware 插值
    ntk_freqs = get_ntk_aware_freqs_cis(
        dim=128, end=test_seq_len,
        base=10000.0, scale_factor=4.0)
    results["ntk_aware"] = evaluate_ntk(model, ntk_freqs)

    # 方法3: Dynamic NTK（根据序列长度动态调整）
    for s in [2.0, 4.0, 6.0, 8.0]:
        dyn_freqs = get_ntk_aware_freqs_cis(
            dim=128, end=test_seq_len,
            base=10000.0, scale_factor=s)
        results[f"dynamic_s={s}"] = evaluate_ntk(model, dyn_freqs)

    for method, ppl in results.items():
        print(f"{method:20s}: PPL = {ppl:.2f}")
    return results

# Position Interpolation 实现
def apply_position_interpolation(
    positions: torch.Tensor,
    scale_factor: float,
    max_train_len: int = 2048
) -> torch.Tensor:
    """简单但有效的位置索引缩放"""
    # 将实际位置线性映射到训练范围
    interpolated = positions / scale_factor
    # 截断到训练范围内
    interpolated = torch.clamp(interpolated,
                                0, max_train_len - 1)
    return interpolated

# 在推理中使用
original_positions = torch.arange(8192)
interp_positions = apply_position_interpolation(
    original_positions, scale_factor=4.0)
# 然后使用 interp_positions 查询预计算的 freqs_cis
q = apply_rotary_emb(q, freqs_cis[interp_positions])
k = apply_rotary_emb(k, freqs_cis[interp_positions])

# YaRN 频率计算（完整实现）
import torch
import math

def get_yarn_freqs_cis(
    dim: int, end: int, base: float = 10000.0,
    scale_factor: float = 1.0,
    original_max_position: int = 2048,
    beta_fast: float = 32,
    beta_slow: float = 1,
    mscale: float = 1.0,
) -> torch.Tensor:
    """YaRN: Yet another RoPE extensioN"""
    # 计算每个维度的频率
    pos_freqs = base ** (
        torch.arange(0, dim, 2).float() / dim)

    # 维度感知的动态缩放
    # 高频维度 (大 index) → 大缩放
    # 低频维度 (小 index) → 小缩放
    extrapolation_factor = scale_factor
    dim_factors = torch.maximum(
        torch.log(torch.ones(1) * scale_factor),
        torch.zeros(1))

    # 注意力感知缩放
    yarn_factor = torch.where(
        pos_freqs < beta_slow,
        torch.ones_like(pos_freqs),
        torch.where(
            pos_freqs > beta_fast,
            extrapolation_factor * torch.ones_like(pos_freqs),
            (pos_freqs - beta_slow) / (beta_fast - beta_slow) *
            (extrapolation_factor - 1) + 1
        )
    )

    inv_freq = 1.0 / (pos_freqs * yarn_factor)
    t = torch.arange(end, dtype=inv_freq.dtype)
    freqs = torch.outer(t, inv_freq)

    # 幅度缩放（YaRN 的温度修正）
    freqs = freqs * mscale
    return torch.polar(torch.ones_like(freqs), freqs)

# 实战：修改 Transformers 的 LLaMA RoPE
from transformers import LlamaConfig, LlamaForCausalLM
from transformers.models.llama.modeling_llama import (
    LlamaRotaryEmbedding,
    LlamaLinearScalingRotaryEmbedding,
    LlamaDynamicNTKScalingRotaryEmbedding,
)

# 方案1: 线性缩放（PI）
config = LlamaConfig.from_pretrained("meta-llama/Llama-2-7b")
config.rope_scaling = {
    "type": "linear",
    "factor": 8.0,  # 4096 → 32768
}
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b", config=config)

# 方案2: Dynamic NTK（推荐）
config.rope_scaling = {
    "type": "dynamic",
    "factor": 8.0,
    "original_max_position_embeddings": 4096,
}
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b", config=config)

# 实战：KV Cache 内存估算与优化
def estimate_kv_cache_gb(
    n_layers: int = 32, n_heads: int = 32,
    head_dim: int = 128, seq_len: int = 32768,
    batch_size: int = 1, dtype_bytes: int = 2
) -> float:
    """估算 KV Cache 的显存占用（GB）"""
    # 每层: 2(K+V) × batch × heads × seq × head_dim × dtype
    per_layer = 2 * batch_size * n_heads * seq_len * head_dim
    total_bytes = per_layer * n_layers * dtype_bytes
    return total_bytes / (1024 ** 3)

# LLaMA-7B @ 32K
ppl_mem = estimate_kv_cache_gb(seq_len=32768)
print(f"LLaMA-7B @ 32K: KV Cache ≈ {ppl_mem:.1f} GB")
# 输出: ~8.0 GB（FP16, batch=1）

# KV Cache 量化到 INT8 可减半
int8_mem = estimate_kv_cache_gb(seq_len=32768, dtype_bytes=1)
print(f"LLaMA-7B @ 32K (INT8): KV Cache ≈ {int8_mem:.1f} GB")
# 输出: ~4.0 GB

# Needle-in-Haystack 验证
def needle_in_haystack_test(
    model, tokenizer, context_len: int = 32768,
    n_rounds: int = 10
) -> float:
    """验证扩展后的长上下文检索能力"""
    correct = 0
    for _ in range(n_rounds):
        needle = f"{{'answer': {torch.randint(0,100,(1,)).item()}}}"
        haystack = "The sky is blue. " * (context_len // 4)
        prompt = haystack + needle + haystack
        # ... 生成并检查是否找到 needle
    return correct / n_rounds