01 - Transformer深入理解（全面版）¶

import torch
import torch.nn as nn
import math

class TokenEmbedding(nn.Module):
    """
    Token嵌入层完整实现
    """
    def __init__(self, vocab_size, d_model, padding_idx=None):
        super().__init__()  # super()调用父类方法
        self.vocab_size = vocab_size
        self.d_model = d_model

        # 嵌入矩阵
        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=d_model,
            padding_idx=padding_idx  # 用于mask填充位置
        )

        # 缩放因子（Transformer原论文使用）
        self.scale = math.sqrt(d_model)

        # 初始化（重要！）
        self._init_weights()

    def _init_weights(self):
        """
        嵌入层初始化策略
        使用N(0, 1/d_model)初始化
        """
        nn.init.normal_(self.embedding.weight, mean=0, std=1/math.sqrt(self.d_model))

        # 如果有padding_idx，将该位置嵌入置零
        if self.embedding.padding_idx is not None:
            with torch.no_grad():  # 禁用梯度追踪，避免初始化操作被记入计算图
                self.embedding.weight[self.embedding.padding_idx].fill_(0)

    def forward(self, x):
        """
        Args:
            x: [batch_size, seq_len] token IDs
        Returns:
            embeddings: [batch_size, seq_len, d_model]
        """
        # 查找嵌入并缩放
        # 缩放原因：后续要与位置编码相加，需要平衡量级
        return self.embedding(x) * self.scale

# 使用示例
vocab_size = 32000
d_model = 512
batch_size = 2
seq_len = 10

embedding_layer = TokenEmbedding(vocab_size, d_model)
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
embeddings = embedding_layer(input_ids)

print(f"输入形状: {input_ids.shape}")
print(f"嵌入形状: {embeddings.shape}")
print(f"嵌入范围: [{embeddings.min():.3f}, {embeddings.max():.3f}]")

# 使用Hugging Face Tokenizer示例
from transformers import AutoTokenizer

# 加载GPT-2的分词器
tokenizer = AutoTokenizer.from_pretrained("gpt2")

text = "Transformer架构 revolutionized NLP"
tokens = tokenizer.tokenize(text)
print(f"分词结果: {tokens}")
# 输出: ['Trans', 'former', '架构', 'Ġre', 'volution', 'ized', 'ĠN', 'LP']

# 转换为ID
input_ids = tokenizer.encode(text)
print(f"Token IDs: {input_ids}")

# 解码
decoded = tokenizer.decode(input_ids)
print(f"解码结果: {decoded}")

class SinusoidalPositionalEncoding(nn.Module):
    """
    正弦位置编码实现
    """
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 创建位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)

        # 位置索引 [max_len, 1]
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # unsqueeze增加一个维度

        # 维度索引的除数项
        # 10000^(2i/d_model) = exp(2i * -log(10000) / d_model)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() *
            (-math.log(10000.0) / d_model)
        )

        # 偶数维度用sin
        pe[:, 0::2] = torch.sin(position * div_term)

        # 奇数维度用cos
        pe[:, 1::2] = torch.cos(position * div_term)

        # 注册为buffer（不参与训练）
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: [batch_size, seq_len, d_model]
        """
        seq_len = x.size(1)

        # 添加位置编码
        x = x + self.pe[:, :seq_len, :]

        return self.dropout(x)

# 可视化位置编码
import matplotlib.pyplot as plt

def visualize_positional_encoding():
    d_model = 128
    max_len = 100

    pe = SinusoidalPositionalEncoding(d_model, max_len)
    encoding = pe.pe[0].numpy()  # [max_len, d_model]

    plt.figure(figsize=(12, 6))
    plt.imshow(encoding, aspect='auto', cmap='viridis')
    plt.colorbar()
    plt.xlabel('Dimension')
    plt.ylabel('Position')
    plt.title('Sinusoidal Positional Encoding')
    plt.show()

    # 观察不同位置的编码
    plt.figure(figsize=(12, 4))
    for pos in [0, 10, 20, 50]:
        plt.plot(encoding[pos], label=f'Pos {pos}')
    plt.legend()
    plt.xlabel('Dimension')
    plt.ylabel('Value')
    plt.title('Positional Encoding at Different Positions')
    plt.show()

# 正弦位置编码的性质
"""
性质1: 唯一性
每个位置都有唯一的编码

性质2: 相对位置关系
PE(pos+k) 可以表示为 PE(pos) 的线性函数
这意味着模型可以学习相对位置

性质3: 有界性
所有值都在[-1, 1]之间，数值稳定

性质4: 外推性
可以处理训练时未见过的更长序列
"""

class LearnedPositionalEncoding(nn.Module):
    """
    可学习位置编码
    """
    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 可学习的嵌入矩阵
        self.pos_embedding = nn.Embedding(max_len, d_model)

        # 初始化
        nn.init.normal_(self.pos_embedding.weight, std=0.02)

    def forward(self, x):
        """
        Args:
            x: [batch_size, seq_len, d_model]
        """
        seq_len = x.size(1)

        # 创建位置索引
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)

        # 添加位置嵌入
        x = x + self.pos_embedding(positions)

        return self.dropout(x)

class RotaryPositionalEmbedding(nn.Module):
    """
    旋转位置编码 (RoPE)
    通过旋转矩阵注入相对位置信息
    """
    def __init__(self, dim, max_seq_len=2048, base=10000):
        super().__init__()

        # 计算旋转角度
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

        # 预计算位置编码
        t = torch.arange(max_seq_len)
        freqs = torch.einsum('i,j->ij', t, inv_freq)  # [max_seq_len, dim/2]

        # 复数形式: cos + i*sin
        self.register_buffer('cos_cached', freqs.cos()[None, None, :, :])  # [1, 1, seq_len, dim/2]
        self.register_buffer('sin_cached', freqs.sin()[None, None, :, :])  # [1, 1, seq_len, dim/2]

    def forward(self, x, seq_len=None):
        """
        Args:
            x: [batch, heads, seq_len, head_dim]
        """
        if seq_len is None:
            seq_len = x.shape[2]

        cos = self.cos_cached[:, :, :seq_len, :]
        sin = self.sin_cached[:, :, :seq_len, :]

        return self.apply_rotary_pos_emb(x, cos, sin)

    def apply_rotary_pos_emb(self, x, cos, sin):
        """
        应用旋转位置编码

        Args:
            x: [batch, heads, seq_len, head_dim]
            cos: [1, 1, seq_len, head_dim/2]
            sin: [1, 1, seq_len, head_dim/2]

        Returns:
            [batch, heads, seq_len, head_dim]
        """
        # 将x分成两部分（偶数维和奇数维）
        x1, x2 = x[..., ::2], x[..., 1::2]

        # 旋转: [x1, x2] @ [[cos, -sin], [sin, cos]]
        # 即: x1' = x1 * cos - x2 * sin
        #     x2' = x1 * sin + x2 * cos
        rotated = torch.stack([
            x1 * cos - x2 * sin,
            x1 * sin + x2 * cos
        ], dim=-1)

        # 将最后两维展平，恢复原始形状
        return rotated.flatten(-2)

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

def scaled_dot_product_attention(
    query: torch.Tensor,   # [batch, seq_q, d_k]
    key: torch.Tensor,     # [batch, seq_k, d_k]
    value: torch.Tensor,   # [batch, seq_k, d_v]
    mask: torch.Tensor = None,
    dropout: nn.Dropout = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    缩放点积注意力

    关键点：
    1. 为什么除以 sqrt(d_k)?（完整推导）

       假设 q 和 k 的每个分量独立同分布，均值为0、方差为1：
         q = [q_1, q_2, ..., q_{d_k}],  q_i ~ (0, 1)
         k = [k_1, k_2, ..., k_{d_k}],  k_i ~ (0, 1)

       点积 q·k = Σ_{i=1}^{d_k} q_i * k_i

       每一项 q_i * k_i 的方差：
         Var(q_i * k_i) = E[q_i²] * E[k_i²] - (E[q_i] * E[k_i])²
                        = 1 * 1 - 0 = 1

       因为各项独立，所以：
         Var(q·k) = Σ Var(q_i * k_i) = d_k

       当 d_k=64 时，点积标准差 ≈ 8；d_k=128 时 ≈ 11.3
       这些大值会将 softmax 推入饱和区（梯度接近 0）

       除以 √d_k 后：
         Var(q·k / √d_k) = Var(q·k) / d_k = d_k / d_k = 1

       方差归一化为 1，softmax 的输入保持在合理范围，梯度可以正常流动。

    2. mask 的两种用途:
       - padding mask: 屏蔽填充位置
       - causal mask: 屏蔽未来位置（解码器）
    """
    d_k = query.size(-1)

    # Step 1: QK^T / sqrt(d_k) → [batch, seq_q, seq_k]
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # Step 2: 应用 mask（将需要屏蔽的位置设为 -inf）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 3: Softmax → 权重矩阵
    attn_weights = F.softmax(scores, dim=-1)

    # Step 4: Dropout（训练时）
    if dropout is not None:
        attn_weights = dropout(attn_weights)

    # Step 5: 加权求和 → [batch, seq_q, d_v]
    output = torch.matmul(attn_weights, value)

    return output, attn_weights

# ---------- 验证 ----------
batch, seq_len, d_k = 2, 5, 64
Q = torch.randn(batch, seq_len, d_k)
K = torch.randn(batch, seq_len, d_k)
V = torch.randn(batch, seq_len, d_k)

# 无 mask
out, weights = scaled_dot_product_attention(Q, K, V)
print(f"输出形状: {out.shape}")           # [2, 5, 64]
print(f"权重形状: {weights.shape}")        # [2, 5, 5]
print(f"权重行和: {weights.sum(-1)}")      # 每行和为 1.0

# 因果 mask（解码器用）
causal_mask = torch.tril(torch.ones(seq_len, seq_len))  # 下三角矩阵
out_masked, weights_masked = scaled_dot_product_attention(Q, K, V, mask=causal_mask)
print(f"\n因果mask权重（第0行只看自己，第4行看所有位置）:")
print(weights_masked[0].detach())

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制

    参数量分析（d_model=1024, h=16, d_k=64）:
    W_Q: 1024×1024 = 1M
    W_K: 1024×1024 = 1M
    W_V: 1024×1024 = 1M
    W_O: 1024×1024 = 1M
    总计: 4M 参数
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须被 n_heads 整除"  # assert断言：条件False时抛出AssertionError

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 每个头的维度

        # 四个投影矩阵
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        query: torch.Tensor,   # [batch, seq_q, d_model]
        key: torch.Tensor,     # [batch, seq_k, d_model]
        value: torch.Tensor,   # [batch, seq_k, d_model]
        mask: torch.Tensor = None,
    ) -> torch.Tensor:
        batch_size = query.size(0)

        # 1. 线性投影 → [batch, seq, d_model]
        Q = self.W_Q(query)
        K = self.W_K(key)
        V = self.W_V(value)

        # 2. 分头: [batch, seq, d_model] → [batch, n_heads, seq, d_k]
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # view重塑张量形状
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 缩放点积注意力
        d_k = self.d_k
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(0).unsqueeze(0)
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 4. 加权求和 → [batch, n_heads, seq_q, d_k]
        context = torch.matmul(attn_weights, V)

        # 5. 合并多头: [batch, seq_q, d_model]
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 6. 输出投影
        output = self.W_O(context)

        return output

# 验证
mha = MultiHeadAttention(d_model=512, n_heads=8)
x = torch.randn(2, 10, 512)

# Self-attention: Q=K=V=x
out = mha(x, x, x)
print(f"Self-attention 输出: {out.shape}")  # [2, 10, 512]

# Cross-attention: Q来自decoder, K/V来自encoder
enc_out = torch.randn(2, 20, 512)
dec_in = torch.randn(2, 10, 512)
cross_out = mha(dec_in, enc_out, enc_out)
print(f"Cross-attention 输出: {cross_out.shape}")  # [2, 10, 512]

class FeedForward(nn.Module):
    """标准FFN，4倍扩展"""
    def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        d_ff = d_ff or 4 * d_model  # 默认4倍扩展

        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # [batch, seq, d_model] → [batch, seq, d_ff] → [batch, seq, d_model]
        return self.w2(self.dropout(F.relu(self.w1(x))))

# GeLU (GPT-2/BERT 使用)
# GELU(x) = x · Φ(x)，其中 Φ 是标准正态分布的CDF
# 比 ReLU 更平滑，不会在 x=0 处产生不可导点

class GeLUFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        d_ff = d_ff or 4 * d_model
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w2(self.dropout(F.gelu(self.w1(x))))

# SwiGLU (LLaMA/Qwen/Mistral 使用)
# SwiGLU(x, W, V, W2) = (Swish(xW) ⊙ xV) W2
# 引入门控机制，效果优于 GeLU，但参数量增加 50%

class SwiGLUFeedForward(nn.Module):
    """LLaMA-style SwiGLU FFN"""
    def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        # LLaMA 使用 8/3 * d_model 作为 d_ff 以保持参数量不变
        d_ff = d_ff or int(8 / 3 * d_model)
        # 取最接近的 256 的倍数（GPU对齐优化）
        d_ff = ((d_ff + 255) // 256) * 256

        self.w1 = nn.Linear(d_model, d_ff, bias=False)  # gate projection
        self.w3 = nn.Linear(d_model, d_ff, bias=False)  # up projection
        self.w2 = nn.Linear(d_ff, d_model, bias=False)  # down projection
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # Swish(xW1) ⊙ xW3 → 门控机制
        gate = F.silu(self.w1(x))  # SiLU = Swish(x) = x * sigmoid(x)
        up = self.w3(x)
        return self.w2(self.dropout(gate * up))

# 对比
d_model = 1024
ffn_relu = FeedForward(d_model)
ffn_gelu = GeLUFeedForward(d_model)
ffn_swiglu = SwiGLUFeedForward(d_model)

for name, module in [("ReLU FFN", ffn_relu), ("GeLU FFN", ffn_gelu), ("SwiGLU FFN", ffn_swiglu)]:
    params = sum(p.numel() for p in module.parameters())
    print(f"{name}: {params:,} 参数")
# ReLU FFN: 8,393,728 参数
# GeLU FFN: 8,393,728 参数
# SwiGLU FFN: 8,650,752 参数 (d_ff调整后参数量接近)

class LayerNorm(nn.Module):
    """
    标准层归一化 (BERT/GPT-2 使用)
    LayerNorm(x) = γ · (x - μ) / √(σ² + ε) + β
    """
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))   # 可学习缩放
        self.beta = nn.Parameter(torch.zeros(d_model))    # 可学习偏移
        self.eps = eps

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

class RMSNorm(nn.Module):
    """
    RMS层归一化 (LLaMA/Qwen/Mistral 使用)
    RMSNorm(x) = γ · x / RMS(x)
    其中 RMS(x) = √(1/n · Σx²)

    优势：去掉了均值计算和偏移参数 β
    → 计算更快（约10-15%），效果相当
    """
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return self.gamma * x / rms

# 对比
x = torch.randn(2, 10, 512)
ln = LayerNorm(512)
rmsn = RMSNorm(512)

out_ln = ln(x)
out_rmsn = rmsn(x)
print(f"LayerNorm 输出统计: mean={out_ln.mean():.4f}, std={out_ln.std():.4f}")
print(f"RMSNorm 输出统计:   mean={out_rmsn.mean():.4f}, std={out_rmsn.std():.4f}")

class TransformerBlock(nn.Module):
    """Pre-Norm Transformer Block (LLaMA-style)"""

    def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLUFeedForward(d_model, d_ff, dropout)

    def forward(self, x, mask=None):
        # Pre-Norm + Residual
        x_norm = self.norm1(x)
        x = x + self.attn(x_norm, x_norm, x_norm, mask)
        x = x + self.ffn(self.norm2(x))
        return x

def create_causal_mask(seq_len: int) -> torch.Tensor:
    """
    创建因果注意力掩码

    例如 seq_len=4:
    [[1, 0, 0, 0],   ← 位置0只看自己
     [1, 1, 0, 0],   ← 位置1看0和1
     [1, 1, 1, 0],   ← 位置2看0、1、2
     [1, 1, 1, 1]]   ← 位置3看所有
    """
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# 训练时：一次前向传播同时预测所有位置
# 因为有因果mask，每个位置只看到左侧，等价于自回归

class WarmupCosineScheduler:
    """
    Warmup + Cosine Decay 学习率调度
    几乎所有现代大模型都使用这种方式

    阶段1 (warmup): 学习率从0线性增长到max_lr
    阶段2 (cosine): 学习率从max_lr余弦衰减到min_lr
    """
    def __init__(
        self,
        optimizer,
        warmup_steps: int,
        total_steps: int,
        max_lr: float = 3e-4,
        min_lr: float = 1e-5,
    ):
        self.optimizer = optimizer
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.max_lr = max_lr
        self.min_lr = min_lr
        self.step_count = 0

    def step(self):
        self.step_count += 1
        lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

    def get_lr(self):
        if self.step_count < self.warmup_steps:
            # 线性 warmup
            return self.max_lr * self.step_count / self.warmup_steps
        else:
            # 余弦衰减
            progress = (self.step_count - self.warmup_steps) / (
                self.total_steps - self.warmup_steps
            )
            return self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (
                1 + math.cos(math.pi * progress)
            )

# 训练循环骨架
"""
model = DecoderOnlyTransformer(...)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
scheduler = WarmupCosineScheduler(optimizer, warmup_steps=2000, total_steps=100000)

for batch in dataloader:
    input_ids = batch['input_ids']                      # [B, T]
    targets = input_ids[:, 1:]                          # 右移一位
    logits = model(input_ids[:, :-1])                   # [B, T-1, vocab]
    loss = F.cross_entropy(logits.reshape(-1, vocab_size), targets.reshape(-1))

    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 梯度裁剪
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()
"""

"""KV Cache 原理示意"""

class CachedMultiHeadAttention(nn.Module):
    """带KV Cache的多头注意力（推理用）"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, kv_cache=None):
        """
        Args:
            x: [batch, 1, d_model] — 推理时只有当前新token
            kv_cache: (cached_K, cached_V) — 之前步骤的K/V
        Returns:
            output, new_kv_cache
        """
        B = x.size(0)

        # 只为当前 token 计算 Q/K/V
        Q = self.W_Q(x).view(B, 1, self.n_heads, self.d_k).transpose(1, 2)
        K_new = self.W_K(x).view(B, 1, self.n_heads, self.d_k).transpose(1, 2)
        V_new = self.W_V(x).view(B, 1, self.n_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            K_cached, V_cached = kv_cache
            # 拼接历史 KV → [batch, n_heads, seq_so_far+1, d_k]
            K = torch.cat([K_cached, K_new], dim=2)
            V = torch.cat([V_cached, V_new], dim=2)
        else:
            K, V = K_new, V_new

        # 只计算当前 Q 与所有 K 的注意力 → O(1×S) 而非 O(S×S)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)

        out = out.transpose(1, 2).contiguous().view(B, 1, self.d_model)
        output = self.W_O(out)

        return output, (K, V)  # 返回更新后的 cache

"""
无 KV Cache: 生成 T 个token → 总计算量 O(T² · d)
有 KV Cache: 生成 T 个token → 总计算量 O(T · d) + 缓存内存
  速度提升: ~T倍（序列越长，加速越明显）

代价: 需要额外内存存储 KV Cache
  LLaMA-2-7B, seq_len=4096:
  KV cache = 2(K+V) × 32(layers) × 32(heads) × 4096(seq) × 128(d_k) × 2(bytes/fp16)
           ≈ 2GB
"""

def generate(
    model, tokenizer, prompt: str,
    max_new_tokens: int = 100,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 0.9,
):
    """
    完整的文本生成函数，支持多种解码策略
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    kv_cache = None
    generated = []

    for _ in range(max_new_tokens):
        # 只输入最新 token（有 KV Cache 时）
        if kv_cache is not None:
            x = input_ids[:, -1:]
        else:
            x = input_ids

        logits, kv_cache = model(x, kv_cache=kv_cache)
        logits = logits[:, -1, :]  # 只取最后一个位置

        # Temperature scaling
        logits = logits / temperature

        # Top-K: 只保留概率最高的K个token
        if top_k > 0:
            topk_logits, topk_indices = torch.topk(logits, top_k)
            logits = torch.full_like(logits, float('-inf'))
            logits.scatter_(1, topk_indices, topk_logits)

        # Top-P (Nucleus Sampling)
        probs = F.softmax(logits, dim=-1)
        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

        # 移除累积概率超过 top_p 的token
        remove_mask = cumulative_probs - sorted_probs > top_p
        sorted_probs[remove_mask] = 0
        sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)

        # 采样
        next_token = torch.multinomial(sorted_probs, 1)
        next_token = sorted_indices.gather(1, next_token)

        if next_token.item() == tokenizer.eos_token_id:
            break

        generated.append(next_token.item())
        input_ids = torch.cat([input_ids, next_token], dim=-1)

    return tokenizer.decode(generated)