第1章：长上下文技术 (Long Context)#

import torch

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """
    预计算旋转角度（复数形式）

    Args:
        dim: head_dim (注意不是 hidden_size)
        end: 最大序列长度 max_seq_len
        theta: 基频 (LLaMA 1用10000, LLaMA 3用500000)
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end)  # 位置索引 [0, 1, ..., end-1]

    # 外积计算所有位置的所有频率
    freqs = torch.outer(t, freqs).float()  # [end, dim//2]

    # 转为复数 e^{i*freqs} = cos(freqs) + i*sin(freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    """
    应用 RoPE 旋转

    Args:
        xq: Query [batch, seq_len, n_heads, head_dim]
        xk: Key   [batch, seq_len, n_kv_heads, head_dim]
        freqs_cis: 预计算的复数频率
    """
    # 将 Q, K 重塑为复数张量 (把最后一维 dim 拆成 dim/2 个复数)
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # 广播形状以匹配
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)

    # 复数乘法（即旋转）: (a+bi)(c+di) = (ac-bd) + i(ad+bc)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

# Linear Scaling
scale = 8.0  # 4k -> 32k
t = torch.arange(end) / scale

def get_ntk_base(scale: float, dim: int, base: float = 10000.0):
    """
    计算 NTK 修正后的 Base

    Args:
        scale: 扩展倍数 (e.g., 8)
        dim: head_dim
    """
    # 核心公式：base = base * scale ^ (dim / (dim-2))
    new_base = base * (scale ** (dim / (dim - 2)))
    return new_base

# 使用新的 base 计算 freqs
freqs = 1.0 / (new_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

import torch
import torch.nn.functional as F

# 启用 FlashAttention
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    output = F.scaled_dot_product_attention(
        query, key, value,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True
    )

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class LlamaRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=32768, base=10000, device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base

        # 预计算 cos/sin
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)

        self.update_freqs(max_position_embeddings, device)

    def update_freqs(self, seq_len, device):
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.cos_cached = emb.cos()[None, None, :, :]
        self.sin_cached = emb.sin()[None, None, :, :]

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_position_embeddings:
            # 动态扩展（简单线性外推）
            self.update_freqs(seq_len, x.device)

        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    # 简单的实现版本，没有使用复数
    # q, k: [bs, num_heads, seq_len, head_dim]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class LongContextAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = self.hidden_size // self.num_heads

        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)

        self.rotary_emb = LlamaRotaryEmbedding(self.head_dim)

    def forward(self, x, attention_mask=None):
        bsz, seq_len, _ = x.shape

        # 1. 投影
        q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 2. 应用 RoPE
        cos, sin = self.rotary_emb(v, seq_len=seq_len)
        q, k = apply_rotary_pos_emb(q, k, cos, sin, None)

        # 3. FlashAttention
        # 自动选择最优实现（FlashV2 > MemEfficient > Math）
        context_layer = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attention_mask,
            dropout_p=0.0 if not self.training else 0.1,
            is_causal=True
        )

        # 4. 输出
        context_layer = context_layer.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size)
        output = self.o_proj(context_layer)

        return output

# 测试代码
device = "cuda" if torch.cuda.is_available() else "cpu"
config = type('Config', (), {'hidden_size': 4096, 'num_heads': 32})()
attn = LongContextAttention(config).to(device)

x = torch.randn(1, 1024, 4096).to(device)
with torch.no_grad():
    out = attn(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {out.shape}")  # 应该是 [1, 1024, 4096]
    print("长上下文 Attention 计算成功！")