• 大模型

你是否也好奇：3.8B 参数的模型在端侧设备上为什么回答还是慢？本文以 Phi-3-mini-128k-instruct_rai_1.8.0_medusa_1.0_npu 模型为例，讲透Medusa 推理加速为何成为端侧 AI 的关键技术，并详解其内建的**投机解码（Speculative Decoding）**原理。

1️⃣ 端侧 AI 的瓶颈：为什么"一个字一个字蹦"很痛

大模型生成文本是自回归的：每生成 1 个 token，都必须等上一个 token 算完，才能开始下一个。这意味着：

• 每一步都只"喂"给芯片 1 个 token，算力严重吃不饱；
• 推理时间几乎全花在反复搬运模型权重上，属于典型的"访存瓶颈"；
• 端侧设备（手机、开发板、NPU 加速卡）的内存带宽本来就有限，瓶颈被进一步放大。

👉 结论：端侧 LLM 体验差，往往不是模型"笨"，而是解码太慢。这正是 Medusa 推理加速要解决的问题。

2️⃣ 投机解码核心思想：先起草，再批量验证

投机解码的灵感来自"起草 + 校对"：

1. 起草：由廉价的预测机制一次性提出 K 个候选后续 token；
2. 验证：主模型用一次前向计算同时校验这整串候选；
3. 采纳：从开头逐位比对，匹配到哪一位就采纳哪一位，在首个分歧点回退，重新正常生成。

妙处在于：用生成 1 个字的开销，验证了一串字。只要预测命中率足够高，端到端速度就能翻倍，且数学上可以保证输出分布与原始模型完全一致（贪心解码下输出逐字相同）。

3️⃣ Medusa：把"草稿模型"换成轻量"解码头"

传统投机解码要额外训练一个小型草稿模型，工程上偏重。Medusa 换了一个更巧妙的思路：

• 在原模型之上加装多个轻量解码头（Decoding Heads），每个头负责往前多看一位——第 1 个头预测"下一个词"，第 2 个头预测"下下个词"，以此类推；
• 各头的预测结果按前缀关系组织成一棵候选树（通常 K=64 条候选序列）；
• 只需对解码头做轻量训练，不需要单独维护一个草稿模型，增量权重只有几十 MB，非常适合随主模型一起下发到端侧。

📌 本模型名中的medusa_1.0表示采用了Medusa v1 配置：2 个解码头、候选树大小 K=64。

4️⃣ 树注意力（Tree Attention）：一次前向验证整棵候选树

"一次前向验证一串字"靠的是Tree Attention：

• 构造特殊的注意力掩码：候选树上每个节点只能看到自己的前缀路径，兄弟分支之间互不干扰；
• 主模型只需执行一次并行前向，就把 K=64 条候选序列全部校验完毕；
• 之后取最长匹配分支继续生成，分歧点之后回退到常规解码。

效果对比一目了然：

5️⃣ 30 秒读懂模型名：每个后缀都是"卖点"

128K 长上下文 × 端侧 NPU × 2× 解码加速，正是"端侧 AI 关键"三者缺一不可的原因。

6️⃣ 快速上手：下载模型与启用 Medusa 解码

6.1 通过 ModelScope SDK 下载

模型卡与下载说明见仓库内的 README.md，SDK 下载方式如下（完整示例见 README.md）：

pip install modelscope

from modelscope import snapshot_download model_dir = snapshot_download('amd/Phi-3-mini-128k-instruct_rai_1.8.0_medusa_1.0_npu')

仓库中的.safetensors、.pt等权重文件均按 .gitattributes 的规则以 Git LFS 托管，克隆时会自动拉取真实权重。

6.2 用 transformers 原生运行投机解码

transformers 自 v4.44.0 起原生支持 Medusa 投机解码，核心参数就一个——num_beams=64（对应候选树大小 K）：

from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "amd/Phi-3-mini-128k-instruct_rai_1.8.0_medusa_1.0_npu" model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cpu") tokenizer = AutoTokenizer.from_pretrained(model_id) prompt = "If you could have any superpower, what would you choose?" inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate( **inputs, max_new_tokens=512, num_beams=64, # 对应 Medusa 候选树 K=64，投机解码生效的关键 assistant_token_id="assistant", ) print(tokenizer.decode(outputs[0], skip_special_tokens=False))

⚠️ 注意：num_beams在这里不是常规 beam search，而是驱动候选树验证的参数，保持 64 即可。

7️⃣ 实战提醒：加速从哪来，代价是什么

✅收益

• 输出越长，摊薄效果越明显——长文档总结、代码生成等任务收益最大；
• 不牺牲质量：验证机制保证输出与原模型分布一致；
• 权重增量极小，端侧 OTA 分发成本低。

⚠️成本

• 候选树并行验证会略微增加显存/内存占用，内存紧张的端侧设备需实测；
• 首 token 延迟不会下降（第一串仍需等主模型验证）；
• npu后缀版本需确认所用 NPU 运行时的导出兼容性，建议先跑通最小用例再上生产。

8️⃣ 总结：端侧 AI 落地前请对照这份清单

1. 延迟敏感？是 → 优先评估 Medusa 推理加速（本模型即开箱即用）；
2. 框架够新？确认 transformers ≥ 4.44，或 NPU 工具链支持树注意力；
3. 内存够吗？预留候选树验证的额外开销；
4. 先长后短：用长输出任务验证 2× 加速后再推广到短回复场景。

📎本文引用的仓库文件

• 模型卡说明：README.md
• 下载方式（SDK / Git）：README.md
• 权重文件 LFS 管理规则：.gitattributes

• 大模型