自回归推理与 KV-Cache 显存优化：打造毫秒级打字机

===================================================================================
                       朴素自回归推理的重复计算噩梦
===================================================================================
生成第 1 个字 ("公"):
  打字员朗读上文: "今天天气真好，我们一起去"                  ──> 算了 11 个字的特征

生成第 2 个字 ("园"):
  打字员从头重读: "今天天气真好，我们一起去公"                ──> 算了 12 个字的特征

生成第 3 个字 ("散"):
  打字员再次从头重读: "今天天气真好，我们一起去公园"          ──> 算了 13 个字的特征
  ...
生成第 100 个字:
  打字员又从头重读: [ 前面全部 110 个字全部重读一遍！ ]       ──> 算了 111 个字的特征
===================================================================================

if repetition_penalty != 1.0 and generated_ids:
    exempt_ids = {0, 1, 2, 3, 14, 262, 263, 273, 274, 275, 534, 1227, 1916}
    for tid in set(generated_ids):
        if tid in exempt_ids:
            continue
        if logits[tid] < 0:
            logits[tid] *= repetition_penalty
        else:
            logits[tid] /= repetition_penalty

logits = logits / temperature

sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

# 下面这两行到底在干嘛？！
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0

原始累加是否超过 0.8:  [ False,   True,   True,   True ]
                                   ↘       ↘       ↘
向右平移一位后结果:    [   ?,   False,   True,   True ]
强制将第一名设为 False: [ False,  False,   True,   True ]
                         (北京)   (上海)   (广州)  (火星)
                          保留     保留     剔除    剔除！

sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() # 向右平移
sorted_indices_to_remove[..., 0] = 0 # 无论如何，绝对不能把第一名剔除掉

import torch
import torch.nn.functional as F
from typing import List, Optional

def sample_next_token(
    logits: torch.Tensor,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 0.9,
    repetition_penalty: float = 1.1,
    generated_ids: Optional[List[int]] = None
) -> int:
    """
    自回归核心采样函数：将模型打分转化为最终选出的单字 ID
    """
    # 保证操作不影响原张量，且展平为一维 [Vocab_Size]
    logits = logits.clone().squeeze()

    # 1. 重复惩罚：削减历史出现过文字的分数，拒绝复读机
    if repetition_penalty != 1.0 and generated_ids:
        for tid in set(generated_ids):
            if logits[tid] < 0:
                logits[tid] *= repetition_penalty
            else:
                logits[tid] /= repetition_penalty

    # 2. 极端贪心搜索保护：如果温度无限接近 0，直接取第一名
    if temperature <= 1e-4:
        return torch.argmax(logits, dim=-1).item()

    # 3. 温度缩放：平滑或陡峭化概率曲线
    logits = logits / temperature

    # 4. Top-K 截断：只保留前 K 个最高分的强力候选
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        # 找出第 K 名的分数，凡是比它小的全部赋为负无穷大 (-inf)
        k_th_val = torch.topk(logits, top_k)[0][..., -1, None]
        indices_to_remove = logits < k_th_val
        logits[indices_to_remove] = -float("Inf")

    # 5. Top-P 核采样：动态累加概率，切除长尾荒谬候选
    if top_p < 1.0:
        # 从大到小排序
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        # 计算累加概率和
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # 找到超出阈值的截断点
        sorted_indices_to_remove = cumulative_probs > top_p
        # 巧妙右移，保留临界词
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # 将被剔除的索引还原映射回原 logits，并置为负无穷大
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = -float("Inf")

    # 6. 转为概率分布，并进行轮盘赌抽奖 (Multinomial 依概率采样)
    probs = F.softmax(logits, dim=-1)
    next_token_id = torch.multinomial(probs, num_samples=1).item()
    return next_token_id

class TextGenerator:
    def __init__(self, model, tokenizer, device: str = "cpu"):
        self.model = model.to(device).eval() # 切换到评估模式，关闭 Dropout
        self.tokenizer = tokenizer
        self.device = device

    @torch.no_grad() # 推理阶段绝不计算梯度，节约海量显存与计算时间
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 128,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.9,
        repetition_penalty: float = 1.1,
        stream: bool = True
    ) -> str:
        # 第一步：把人类输入的提示词编码成数字序列
        prompt_ids = self.tokenizer.encode(prompt, add_bos=True)
        generated = list(prompt_ids)

        if stream:
            print(prompt, end="", flush=True)

        past_key_values = None
        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=self.device)

        # 第二步【首字填充 Prefill 阶段】：
        # 一次性将完整的 Prompt 送入模型，产生初始的 KV-Cache 便签条，并得到第 1 个新字
        logits, _, past_key_values = self.model(input_ids, past_key_values=None, start_pos=0)
        next_token_logits = logits[0, -1, :] # 提取 Prompt 最后一个位置的预测得分

        # 采样出第 1 个生成的词
        next_id = sample_next_token(
            next_token_logits,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            generated_ids=generated
        )

        generated.append(next_id)
        if stream:
            # 实时将这个词翻译回字符，并打在终端屏幕上！
            print(self.tokenizer.decode([next_id], skip_special_tokens=False), end="", flush=True)

        start_pos = len(prompt_ids)

        # 第三步【增量自回归 Decode 阶段】：
        # 依靠 KV-Cache，每次只喂入最新生成的 1 个字！
        for _ in range(max_new_tokens - 1):
            # 如果生成了句子结束符 </s>，说明文章写完了，优雅退出
            if next_id == self.tokenizer.eos_token_id:
                break
            # 如果长度超出了模型最大上限，停止生成
            if start_pos >= self.model.config.max_seq_len:
                break

            # 核心提速秘诀：单字输入！形状仅仅为 [1, 1]
            step_input = torch.tensor([[next_id]], dtype=torch.long, device=self.device)
            logits, _, past_key_values = self.model(
                step_input,
                past_key_values=past_key_values, # 传入并持续追加便利贴
                start_pos=start_pos
            )
            next_token_logits = logits[0, -1, :]

            # 采样下一个字
            next_id = sample_next_token(
                next_token_logits,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                generated_ids=generated
            )

            generated.append(next_id)
            start_pos += 1

            if stream:
                print(self.tokenizer.decode([next_id], skip_special_tokens=False), end="", flush=True)

        if stream:
            print() # 换行收尾

        # 返回全部新生成的文字内容
        return self.tokenizer.decode(generated[len(prompt_ids):])