旋转位置编码 RoPE：从原理到长度外推

def apply_rope(q, positions, theta_base=10000.0):
    """
    q: (batch, seq_len, num_heads, head_dim)
    positions: (batch, seq_len) 位置索引
    """
    d = q.shape[-1]
    half = d // 2
    
    # 计算频率
    freqs = 1.0 / (theta_base ** (torch.arange(0, half).float() / half))
    # (batch, seq_len, half)
    angles = positions.unsqueeze(-1) * freqs  # (batch, seq_len, half)
    
    cos = angles.cos()  # (batch, seq_len, half)
    sin = angles.sin()  # (batch, seq_len, half)
    
    # 扩展到 full dim: [cos_0, cos_1, ..., cos_{d/2-1}, cos_0, ...]
    cos = torch.cat([cos, cos], dim=-1)  # (batch, seq_len, d)
    sin = torch.cat([sin, sin], dim=-1)
    
    # 旋转: [q0, q1, ..., q_{d/2}, q_{d/2+1}, ...]
    # → [q0*cos0 - q1*sin0, q1*cos0 + q0*sin0, ...]
    q_rot = q * cos + rotate_half(q) * sin
    return q_rot

def rotate_half(x):
    """将后半部分取负并交换: [a, b] → [-b, a]"""
    half = x.shape[-1] // 2
    return torch.cat([-x[..., half:], x[..., :half]], dim=-1)

def ntk_aware_freqs(d, base=10000.0, scale=4.0):
    """NTK-aware 频率缩放"""
    half = d // 2
    # 原始频率
    freqs = 1.0 / (base ** (torch.arange(0, half).float() / half))
    # NTK 缩放: base -> base * scale^(d/(d-2))
    scaled_base = base * (scale ** (d / (d - 2)))
    freqs_scaled = 1.0 / (scaled_base ** (torch.arange(0, half).float() / half))
    return freqs_scaled

def yarn_freqs(d, base=10000.0, original_len=4096, target_len=32768):
    """
    YaRN 频率计算
    """
    half = d // 2
    s = target_len / original_len
    
    # 计算每个频率的波长
    freqs = 1.0 / (base ** (torch.arange(0, half).float() / half))
    wavelengths = 2 * math.pi / freqs
    
    # 定义区间边界
    low_freq_wavelen = original_len
    high_freq_wavelen = original_len / (s ** 0.1)  # 经验值
    
    freqs_scaled = freqs.clone()
    for i in range(half):
        wl = wavelengths[i]
        if wl < high_freq_wavelen:
            # 高频: 不缩放
            pass
        elif wl > low_freq_wavelen:
            # 低频: 缩放 s 倍
            freqs_scaled[i] = freqs[i] / s
        else:
            # 中频: 线性插值
            alpha = (wl - high_freq_wavelen) / (low_freq_wavelen - high_freq_wavelen)
            freqs_scaled[i] = freqs[i] * (1 - alpha * (1 - 1/s))
    
    return freqs_scaled

def yarn_attention_scale(target_len, original_len):
    """YaRN 注意力温度缩放"""
    s = target_len / original_len
    if s <= 1:
        return 1.0
    return math.sqrt(1 + math.log(s) / math.log(original_len))

def longrope_search(d, original_len, target_len, eval_fn):
    """
    Longrope: 搜索最优的每组频率缩放因子
    """
    half = d // 2
    freqs = 1.0 / (10000 ** (torch.arange(0, half).float() / half))
    
    # 搜索空间: 每个维度独立缩放
    best_lambdas = torch.ones(half)
    best_score = eval_fn(freqs)
    
    for i in range(half):
        for lam in np.linspace(0.5, 8.0, 50):
            candidate = best_lambdas.clone()
            candidate[i] = lam
            score = eval_fn(freqs * candidate)
            if score > best_score:
                best_score = score
                best_lambdas = candidate
    
    return freqs * best_lambdas

# LLaMA 系列: 从 0 开始
positions = torch.arange(0, seq_len)

# 某些实现: 第一个 token 位置为 0 但有 padding 偏移
positions = torch.arange(0, seq_len) + padding_offset

# ✅ 正确: 在 Flash Attention kernel 内部应用 RoPE
from flash_attn import flash_attn_func
q_rot = apply_rope(q, positions)
k_rot = apply_rope(k, positions)
output = flash_attn_func(q_rot, k_rot, v, causal=True)

# ❌ 错误: 先做标准 attention 再应用 RoPE

def dynamic_ntk_scaling(seq_len, max_train_len=4096, base=10000.0):
    if seq_len <= max_train_len:
        return base
    scale = seq_len / max_train_len
    return base * (scale ** (d / (d - 2)))