一篇教程讲清楚三件事：怎么把一个带多 AI 词典：Token">Token 预测（MTP，Multi-Token Prediction）头的模型冻结并导出成端侧可跑的格式，怎么用校准集把量化带来的精度损失控制到可接受范围，以及怎么在 Pixel 上用统一标准测出真实的延迟与吞吐。
这篇能做出什么
跑完这篇，你会得到一条可复现的端侧流水线：
1. 把训练态模型里的 MTP 头单独切出来、冻结参数、以独立子图导出；
2. 对主干和 MTP 头分别做量化与校准，并量化评估"草稿接受率"这个真正影响体验的指标；
3. 用 adb 把产物推到 Pixel，跑出 TTFT（首 Token 延迟）、TPOT（每 Token 延迟）、有效吞吐和接受率四组数；
4. 用同一批 prompt 和 Gemini Nano 基线对齐比较，得出"MTP 头在这台设备上到底值不值得"的结论。
必须先说清一条边界：Gemini Nano 由设备上的 AICore 提供，权重不对外分发，无法直接"导出"。所以本教程的做法是——冻结与导出这一步，在你可控的同构小模型上完成（结构与 Gemini Nano 这类端侧小模型同类：十亿级参数、分组查询注意力、KV Cache 显式管理）；Pixel 上的 Gemini Nano 则作为对照基线，通过 ML Kit GenAI 的端侧接口调用。这样得到的结论——MTP 头该不该冻结、校准怎么做、加速比大概落在什么区间——可以直接迁移到真实项目。
前置条件清单
动手前逐项确认：
- [ ] 一台支持 AICore 的 Pixel 设备，系统为较新的 Android 版本，已开启开发者选项与 USB 调试；
- [ ] 宿主机（macOS 或 Linux），Python 3.10 以上，已装
adb； - [ ] 一个带 MTP 头的可控小模型权重，或你自己训练的 draft head；
- [ ] 校准数据集：200～1000 条贴近真实业务的短 prompt，覆盖你的目标语言和长度分布；
- [ ] 一台能长时间连着的设备（测量过程要反复跑，别用无线调试）；
- [ ] 时间预算：导出加校准通常要跑几轮，别指望一次成功。
涉及工具版本、模型下载地址、镜像 tag 的地方，一律以官方文档当前版本为准，本文只给通用命令。
步骤一：先把"冻结"的含义钉死
"冻结 MTP 头"在工程上有三层含义，少做一层，后面的测量就不可比：
参数冻结：requires_grad=False，并切到 eval()。这一步最简单，也最容易被跳过——训练态的 Dropout 和 LayerNorm 统计会让 draft 分布偏乐观，测出来的接受率虚高。
图冻结：训练图里 MTP 头往往和主干共享 KV Cache、依赖 teacher forcing 的输入排布。导出前要把它改成"只吃主干最后一层 hidden state、输出 γ 个位置的 logits"的纯推理子图。γ 就是草稿长度，必须写死成常量。
策略冻结：采样温度、top-k、draft 长度在整轮测量里保持不变。端侧 speculative decoding 的接受率对温度极其敏感，温度一变，前后数据就没有可比性。
把这三层写成一个函数，后面所有实验都走它：
```python
import torch
def freeze_mtp(head: torch.nn.Module, gamma: int, temperature: float) -> torch.nn.Module:
head.eval()
for p in head.parameters():
p.requires_grad_(False)
head.gamma = gamma # 草稿长度写死，导出时作为常量
head.temperature = temperature # 记录到实验元数据里
return head
```
步骤二：搭环境
```bash
python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install torch ai-edge-torch ai-edge-quantizer
```
ai-edge-torch 负责 PyTorch 到端侧格式的转换，ai-edge-quantizer 负责量化与校准。两者对 PyTorch 版本有配套要求，装不上就按报错信息调整，以官方文档当前版本为准。
装完先确认设备通：
```bash
adb devices
adb shell getprop ro.product.model
adb shell dumpsys package com.google.android.aicore | head -n 20
```
最后一条用来记录设备上 AICore 的版本号。这个号码要写进你的实验记录——基线对比时，AICore 版本不同，结果不可比。
步骤三：把 MTP 头切出来并导出
MTP 头的结构通常很轻：一个归一化、一个残差投影、一个共享或独立的输出投影。参数量控制在主干的百分之几以内，否则端侧 decode 阶段会被内存带宽拖死。
```python
import torch
class MtpHead(torch.nn.Module):
"""只接收主干最后一层 hidden state，输出 gamma 个位置的 logits。"""
def __init__(self, hidden_size, vocab_size, gamma):
super().__init__()
self.norm = torch.nn.RMSNorm(hidden_size) # 按你的实现替换
self.proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
self.lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)
self.gamma = gamma
def forward(self, hidden): # hidden: [1, 1, H]
h = self.norm(hidden)
outs = []
for _ in range(self.gamma):
h = h + self.proj(h)
outs.append(self.lm_head(h))
return outs
```
导出时要注意：KV Cache 必须作为显式输入输出，不能藏在模块内部。藏起来的话，每次调用都会重新分配缓存，TTFT 会被这一项吃掉一大截，而且你很难从数据里看出问题出在哪。
```python
import torch
import ai_edge_torch
hidden_size, vocab_size, gamma = 2048, 32000, 3
head = freeze_mtp(MtpHead(hidden_size, vocab_size, gamma), gamma=gamma, temperature=0.0)
sample_hidden = torch.randn(1, 1, hidden_size, dtype=torch.float16)
edge_model = ai_edge_torch.convert(head, (sample_hidden,))
edge_model.export("mtp_head.tflite")
```
导出后第一件事是数值对齐：拿同一批 hidden state，对比 PyTorch 输出和端侧输出的 top-1 一致率。低于预期就别急着量化，先查归一化层和 dtype。
步骤四：量化与校准
端侧量化的核心矛盾在这里：主干可以承受 int4/int8 的权重压缩，但 MTP 头不行。draft 头的输出分布本来就比主干更"平"，量化噪声很容易把 top-1 打偏，而接受率对 top-1 是否一致是零容忍的——差一个 token，整段草稿就作废。
推荐策略是混合量化：主干走低比特，MTP 头保留 int8 权重，lm_head 这类输出投影不量化或走更高比特。
```python
from ai_edge_quantizer import quantizer
from ai_edge_quantizer import recipe
qt = quantizer.Quantizer("mtp_head.tflite")
权重-only 动态量化：不需要校准数据，先用它跑通链路
qt.load_quantization_recipe(recipe.int8_dynamic_activation_int8_weight())
qt.quantize().export_model("mtp_head_dyn_int8.tflite")
```
需要静态激活量化时，才引入校准集。校准集的构造方式直接决定精度：
- 从真实业务日志采样，不要用通用语料凑数；
- 去重，控制长度分布与线上一致（长 prompt 和短 prompt 的激活范围差很多）；
- 对话场景必须包含完整的 system prompt 形态，因为实际推理时它一直在上下文里；
- 分批喂进去，形状与导出时的 sample 完全一致。
```python
def representative_dataset(calib_batches):
for batch in calib_batches: # 每项形状与 sample_hidden 相同
yield (batch,)
qt_static = quantizer.Quantizer("mtp_head.tflite")
qt_static.load_quantization_recipe(recipe.int8_static_activation_int8_weight())
qt_static.calibrate(lambda: representative_dataset(calib_batches))
qt_static.quantize().export_model("mtp_head_static_int8.tflite")
```
量化器 API 在不同版本间会调整，函数名以官方文档当前版本为准。
校准之后必须回归测接受率，而不是只看 perplexity。 这是最容易踩的坑：perplexity 掉了 2%，你可能觉得没事，但接受率从 0.7 掉到 0.5，端到端反而比不加 MTP 更慢。评估方法很简单——用同一批 prompt，让主干部署走正常 decode，同时让 MTP 头出草稿，统计 top-1 命中比例：
```python
def acceptance_rate(main_logits, draft_logits, gamma):
"""main_logits: [T, V]，draft_logits: [T, gamma, V]，均为贪心解码下的结果。"""
hit = 0
total = 0
for t in range(len(main_logits) - gamma):
for k in range(gamma):
total += 1
if draft_logits[t][k].argmax() == main_logits[t + k + 1].argmax():
hit += 1
return hit / max(total, 1)
```
步骤五：推到 Pixel 上跑起来
```bash
adb push mtp_head_static_int8.tflite /data/local/tmp/
adb push main_model.tflite /data/local/tmp/
adb shell "cd /data/local/tmp && ./benchmark_model \
--graph=mtp_head_static_int8.tflite \
--num_threads=4 \
--warmup_runs=10 \
--num_runs=50 \
--use_gpu=true"
```
benchmark_model 是 LiteRT / TFLite 官方发布里的通用跑分工具，从官方发布渠道获取，不要用第三方打包版本。
三个参数不能省：--warmup_runs 至少 10 次（首次调用包含 delegate 初始化和图编译，慢一个数量级）；--num_threads 要与你的实际部署配置一致；--use_gpu 开着跑一遍、关掉再跑一遍，确认 delegate 真的生效——GPU delegate 静默回退到 CPU 是很常见的情况，日志里能看到回退原因。
想测端到端而不是单算子，就写一个最小 Android app，或者用命令行驱动你的推理进程，把每次请求的时间戳打出来。
步骤六：统一测量标准
要报四个数，少一个结论就站不住：
TTFT：从发起请求到第一个 token 出来的毫秒数。端侧交互体验主要由它决定。
TPOT：首 token 之后每个 token 的平均耗时。
有效吞吐：输出 token 数 ÷ 端到端墙钟时间。注意 token 数要用 tokenizer 数，不是数字数。
接受率 α 与草稿长度 γ：这两个决定了 MTP 到底能不能加速。给定 α 和 γ，平均每轮被接受的 token 数是：
```
E[accepted] = (1 - α^(γ+1)) / (1 - α)
```
α=0.7、γ=3 时约等于 2.53 个 token；α 掉到 0.5 就只剩 1.88 个，而每个验证轮次的开销没变——这就是为什么接受率比 perplexity 重要得多。
测量脚本至少要做这几件事：
```python
import statistics
def summarize(records):
return {
"ttft_ms": statistics.median(r["ttft_ms"] for r in records),
"tpot_ms": statistics.median(r["tpot_ms"] for r in records),
"tps": statistics.median(r["out_tokens"] / (r["total_ms"] / 1000) for r in records),
"alpha": statistics.mean(r["accepted"] / r["drafted"] for r in records),
"n": len(records),
}
```
每个配置至少跑 5 轮取中位数，不要用平均值（首轮编译和偶发降频会把均值拉歪）。跑之前把设备温度记下来：adb shell cat /sys/class/thermal/thermal_zone*/temp。连续跑几十次以后设备会降频，此时的数据和冷机数据不可比，中间要留冷却间隔。
步骤七：和 Gemini Nano 基线对比
在同一台 Pixel 上，用 ML Kit GenAI 的端侧接口调 Gemini Nano，跑同一批 prompt，记录 TTFT 与完成时间。调用方式和可用能力随 ML Kit 版本变化，具体接口以官方文档当前版本为准。
对比时守住两条：
1. 同一份 prompt 集、同一个输出长度上限，否则吞吐数字没有意义；
2. 记录 AICore 版本号。基线不是恒定的，系统组件更新后数值会漂，实验记录里必须带上版本。
对比的结论通常不是"谁更快"，而是"在什么场景下 MTP 头划算"：短输出、低温度、模板化 prompt 的场景接受率高，收益明显；开放式创作类任务接受率低，加上 draft 的额外开销后可能持平甚至更慢。
常见坑与排错
KV Cache 藏进模块内部。表现为 TTFT 异常高、且随输出变长越来越慢。改成显式输入输出。
动态 shape 触发重编译。GPU delegate 遇到新形状会重新编译，第一次特别慢。把常用形状固定下来，或确保 warmup 覆盖到实际形状。
量化后接受率断崖下跌。优先怀疑 lm_head 和 MTP 头被一起压到了低比特。把这两层单独提回 int8 或 fp16 再测。
校准集和线上分布不一致。用通用语料校准、跑业务 prompt，激活范围会偏。校准集必须来自真实日志。
只报 tokens/s 不报 TTFT。端侧用户感知的是"多久出第一个字"，吞吐高但 TTFT 长的方案体验更差。
发热降频没记录。跑分越跑越慢，误以为是软件问题。加冷却间隔并记录温度。
delegate 静默回退。日志里看实际使用的 delegate，别信配置项。
MTP 头参数过大。decode 阶段是内存带宽瓶颈，头每多一份参数，每步就多一次带宽开销。头占主干参数的比例要压得很低。
用字符数当 token 数。吞吐直接算错，中文场景误差尤其大。
拿单次结果下结论。端侧噪声很大，至少 5 轮取中位数。
下一步建议
跑通这条流水线后，有三个方向值得继续：
自适应草稿长度。固定 γ 不是最优解。按滑动窗口统计的接受率动态调 γ——接受率高就多猜几个，掉下来就缩回去。实现成本不高，收益在混合负载下比较明显。
把 MTP 头放到 NPU。主干跑 GPU、draft 头跑 NPU，理论上能并行；实际能不能吃到收益，取决于设备上两个 delegate 的调度行为，需要实测。
换更强的 draft 结构。MTP 头是相对朴素的方案，树状草稿、特征复用的 draft 模型在同样接受率下往往能拿到更高加速比，代价是实现复杂度上升。
最后一句提醒：所有对比数字都要带上设备型号、系统版本、AICore 版本、量化配置和 γ/温度这五项元数据。缺任何一项，三个月后你自己也复现不出来。