generate.py

"""推理：压缩 KV cache 增量解码 + 基于 MTP 头的自投机解码。

    python generate.py --ckpt checkpoints/best.pt --prompt "在深夜，" --spec

自投机解码（self-speculative decoding）的思路：
MTP 头本来就是训来预测「再下一个」token 的，那它天然就是一个 draft model，
而且和主干共享 KV cache，几乎不额外花钱。

每一轮：
  1. 主干一次前向同时吃 [上一个真 token, 上一轮的草稿 token]
  2. 第 1 个位置的输出告诉我们「真正的下一个 token 应该是什么」→ 用来验草稿
  3. 草稿被接受 → 这一轮白赚一个 token（第 2 个位置的输出直接就是再下一个）
     草稿被拒绝 → 回滚一格，损失仅一次多余的 KV 写入
验收用的是标准 speculative sampling 的接受-重采样规则，
输出分布和不开投机时**严格一致**，不是近似。
"""

import argparse
import sys
import time
from typing import Optional, Tuple

import torch
import torch.nn.functional as F

from config import ModelConfig
from dataio.tokenizer import load_tokenizer
from model import DeepSeekFlashMini


# --------------------------------------------------------------- 采样工具
def filter_logits(logits: torch.Tensor, temperature: float, top_k: int, top_p: float):
    """返回过滤+归一化后的概率分布 (1, V)。"""
    if temperature <= 0:                      # 贪心：退化成 one-hot
        probs = torch.zeros_like(logits)
        probs.scatter_(-1, logits.argmax(-1, keepdim=True), 1.0)
        return probs
    logits = logits / temperature
    if top_k and top_k > 0:
        k = min(top_k, logits.size(-1))
        thresh = logits.topk(k, dim=-1)[0][..., -1:]
        logits = logits.masked_fill(logits < thresh, float("-inf"))
    if top_p and 0 < top_p < 1.0:
        sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
        cum = sorted_logits.softmax(-1).cumsum(-1)
        remove = cum - sorted_logits.softmax(-1) > top_p
        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
        logits = torch.empty_like(logits).scatter_(-1, sorted_idx, sorted_logits)
    return logits.softmax(-1)


def sample_from(probs: torch.Tensor) -> int:
    return int(torch.multinomial(probs, num_samples=1).item())


# --------------------------------------------------------------- 生成器
class Generator:
    def __init__(self, model: DeepSeekFlashMini, tokenizer, device, max_seq_len=None,
                 attn_impl: str = "naive"):
        self.model = model.eval()
        self.tok = tokenizer
        self.device = device
        self.max_seq_len = max_seq_len or model.cfg.max_seq_len
        model.set_attn_impl(attn_impl)
        model.setup_cache(1, self.max_seq_len, device, torch.float32)

    def _t(self, ids):
        return torch.tensor([ids], dtype=torch.long, device=self.device)

    @torch.no_grad()
    def generate(self, prompt: str, max_new_tokens: int = 200, temperature: float = 0.8,
                 top_k: int = 50, top_p: float = 0.95, speculative: bool = False,
                 stream: bool = True, stop_on_eos: bool = True) -> Tuple[str, dict]:
        ids = self.tok.encode(prompt, bos=True)
        keep = max(1, self.max_seq_len - max_new_tokens - 2)
        ids = ids[-keep:] or [self.tok.bos_id]
        self.eos = self.tok.eos_id if stop_on_eos else -1
        gen = self._spec_loop if speculative else self._plain_loop
        t0 = time.time()
        out_ids, meta = gen(ids, max_new_tokens, temperature, top_k, top_p, stream)
        out_ids = out_ids[:max_new_tokens]      # 投机解码一轮可能吐 2 个，这里截齐
        meta["seconds"] = time.time() - t0
        meta["tokens"] = len(out_ids)
        meta["tok_per_s"] = len(out_ids) / max(meta["seconds"], 1e-9)
        meta["ids"] = out_ids
        meta["prompt_ids"] = ids
        text = self.tok.decode(ids + out_ids)
        return text, meta

    # -------------------------------------------------- 普通增量解码
    def _plain_loop(self, ids, max_new, temperature, top_k, top_p, stream):
        m = self.model
        pos = 0
        pending = list(ids)
        out = []
        printed = 0
        while len(out) < max_new:
            _, logits, _ = m.forward_trunk(self._t(pending), start_pos=pos)
            pos += len(pending)
            probs = filter_logits(logits[:, -1].float(), temperature, top_k, top_p)
            t = sample_from(probs)
            if t == self.eos:
                break
            out.append(t)
            pending = [t]
            printed = self._stream(out, printed, stream)
            if pos + 2 >= self.max_seq_len:
                break
        if stream:
            print()
        return out, {"mode": "plain", "accepted": 0, "rounds": len(out)}

    # -------------------------------------------------- MTP 自投机解码
    def _spec_loop(self, ids, max_new, temperature, top_k, top_p, stream):
        m = self.model
        assert m.cfg.n_mtp > 0, "该模型没有 MTP 头，无法投机解码"
        pos = 0
        pending = list(ids)
        out = []
        printed = 0
        draft: Optional[int] = None
        draft_q: Optional[torch.Tensor] = None
        rounds = accepted = 0

        while len(out) < max_new:
            T = len(pending)
            h, logits, _ = m.forward_trunk(self._t(pending), start_pos=pos)
            rounds += 1
            eos_hit = False

            if draft is None:
                # 首轮（prefill）：只出一个真 token
                p = filter_logits(logits[:, -1].float(), temperature, top_k, top_p)
                t_new = sample_from(p)
                pos += T
                if t_new == self.eos:
                    break
                out.append(t_new)
                next_tokens = pending[1:] + [t_new]
                h_chunk, chunk_start = h, pos - T
            else:
                # pending = [上一个真 token, 草稿]，用位置 T-2 的输出来验草稿
                p = filter_logits(logits[:, T - 2].float(), temperature, top_k, top_p)
                q = draft_q
                ratio = (p[0, draft] / q[0, draft].clamp_min(1e-10)).clamp(max=1.0)
                if torch.rand(1, device=ratio.device) < ratio:
                    # ---- 接受草稿：这一轮吐 2 个 token ----
                    accepted += 1
                    out.append(draft)
                    p2 = filter_logits(logits[:, T - 1].float(), temperature, top_k, top_p)
                    t_new = sample_from(p2)
                    pos += T
                    if draft == self.eos or t_new == self.eos:
                        eos_hit = True
                    else:
                        out.append(t_new)
                    next_tokens = pending[1:] + [t_new]
                    h_chunk, chunk_start = h, pos - T
                else:
                    # ---- 拒绝：从残差分布 max(0, p-q) 重采样，保证分布无偏 ----
                    resid = (p - q).clamp_min(0)
                    s = resid.sum()
                    resid = p if s < 1e-9 else resid / s     # p==q 的退化情形
                    t_new = sample_from(resid)
                    pos += T - 1                     # 回滚草稿占的那一格
                    if t_new == self.eos:
                        eos_hit = True
                    else:
                        out.append(t_new)
                    next_tokens = pending[1:T - 1] + [t_new]
                    h_chunk, chunk_start = h[:, :T - 1], pos - (T - 1)

            printed = self._stream(out, printed, stream)
            if eos_hit or pos + 3 >= self.max_seq_len or len(out) >= max_new:
                break

            # ---- 用 MTP 头造下一轮的草稿（顺带把 MTP 的 cache 补齐）----
            _, mtp_logits, _ = m.mtp_forward(
                h_chunk, self._t(next_tokens), start_pos=chunk_start, depth=0)
            draft_q = filter_logits(mtp_logits[:, -1].float(), temperature, top_k, top_p)
            draft = sample_from(draft_q)
            pending = [out[-1], draft]

        if stream:
            print()
        rate = accepted / max(rounds - 1, 1)
        return out, {"mode": "speculative", "accepted": accepted, "rounds": rounds,
                     "accept_rate": rate}

    def _stream(self, out, printed, stream):
        if not stream:
            return printed
        text = self.tok.decode(out)
        if len(text) > printed:
            sys.stdout.write(text[printed:])
            sys.stdout.flush()
            printed = len(text)
        return printed


# --------------------------------------------------------------- CLI
def load_model(ckpt_path, device):
    ck = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = ModelConfig.from_dict(ck["config"])
    model = DeepSeekFlashMini(cfg).to(device)
    model.load_state_dict(ck["model"])
    return model, cfg, ck


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="checkpoints/best.pt")
    ap.add_argument("--tokenizer", default="")
    ap.add_argument("--prompt", default="在深夜，")
    ap.add_argument("--max-new-tokens", type=int, default=160)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top-k", type=int, default=50)
    ap.add_argument("--top-p", type=float, default=0.95)
    ap.add_argument("--spec", action="store_true", help="开启 MTP 自投机解码")
    ap.add_argument("--attn", default="naive", choices=["naive", "absorb"])
    ap.add_argument("--device", default="auto")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = (torch.device("cuda") if torch.cuda.is_available() else
              torch.device("mps") if torch.backends.mps.is_available() else
              torch.device("cpu")) if args.device == "auto" else torch.device(args.device)

    model, cfg, ck = load_model(args.ckpt, device)
    tok_path = args.tokenizer or f"{ck.get('data_dir', 'data')}/tokenizer.json"
    tok = load_tokenizer(tok_path)

    gen = Generator(model, tok, device, attn_impl=args.attn)
    print(f"--- {'MTP 投机解码' if args.spec else '普通解码'} | {args.attn} 注意力 ---")
    print(args.prompt, end="")
    _, meta = gen.generate(args.prompt, args.max_new_tokens, args.temperature,
                           args.top_k, args.top_p, speculative=args.spec)
    line = (f"[{meta['tokens']} tokens / {meta['seconds']:.2f}s = "
            f"{meta['tok_per_s']:.1f} tok/s")
    if args.spec:
        line += f" | 草稿接受率 {meta['accept_rate']*100:.0f}%"
    print(line + "]")


if __name__ == "__main__":
    main()