Gahow Wang
6485c87c5b
Contrast Small-Model (Phase 22-24, done), EAGLE3 (this phase's target), and Multi-Token Prediction (DeepSeek-V3 style, not applicable here). Includes the actual EAGLE3-Qwen3-8B weight tensor listing pulled from AngelSlim/Qwen3-8B_eagle3 on dash5: - 1 midlayer (attention + mlp) with hidden_size=4096 - fc.weight (4096, 12288) fusing 3 target hidden-state levels - q_proj (4096, 8192) taking concat(embed, fused_h) as input - lm_head only over draft_vocab_size=32000, mapped back with d2t table - ~750 MB total (vs 1.2 GB for Qwen3-0.6B), draft cost ~1/10 of target Also captures Qwen3-8B + EAGLE3 speedup benchmark on vLLM: ~1.97-2.02x across MT-bench/HumanEval/GSM8K/Alpaca. That's the number to beat. Next commits will implement Eagle3Head in xserv-model + hook target hidden states out of Qwen3::decode_core.
301 lines
14 KiB
Markdown
301 lines
14 KiB
Markdown
# Phase 25: 三种投机解码方法对比 — Small Model / EAGLE / MTP
|
||
|
||
> 目标：把 speculative decoding 三种主流范式（本项目已试过一种，另两种未实现）
|
||
> 讲清楚，并把 EAGLE3-Qwen3-8B 的实际权重结构展开来看。
|
||
|
||
## 1. 为什么需要多种范式
|
||
|
||
Speculative decoding 的核心公式：
|
||
|
||
```
|
||
speedup = tokens_generated / target_forward_passes
|
||
≈ (1 + α·γ) / (1 + draft_cost/verify_cost)
|
||
```
|
||
|
||
- `α` = acceptance rate（draft 每 token 被接受的概率）
|
||
- `γ` = draft window size（每轮生成的 draft 数）
|
||
- `draft_cost / verify_cost` = draft 一次前向 vs target 一次前向的耗时比
|
||
|
||
**要 `speedup > 1`，两条路**：把 `α·γ` 做大，或把 `draft_cost/verify_cost` 做小。
|
||
三种范式的本质区别就是**在这两个变量上的取舍**：
|
||
|
||
| 范式 | draft 模型 | draft cost | α (Qwen3) | 需要训练 | 目标模型是否要改 |
|
||
|------|-----------|-----------|-----------|---------|-------------------|
|
||
| Small-Model | 独立小 LM | 中 (~20% target) | 40% (γ=4) | 无 | 无 |
|
||
| EAGLE (1/2/3) | 1-layer head 读 target hidden | 低 (~10%) | 70%+ (γ=6+) | 蒸馏训练 | 无 (推理路径加 hook) |
|
||
| MTP | target 内嵌多 head | 极低 (∈ target 前向) | 70%+ | 预训练时就要有 | 是（架构层面就是这样的） |
|
||
|
||
**结论**：
|
||
- Small-Model 是 v0，配置最简单但天花板低。
|
||
- **EAGLE3 是当前性价比最高的落地方案**：draft cost 极低，α 高，需要一次蒸馏训练（约 100k tokens 数据），但对目标模型无侵入。
|
||
- MTP 是 DeepSeek-V3 / DeepSeek-R1 那种"模型天生就懂"的方案，加速比最高但**必须在预训练时就设计进去**，无法事后加装到 Qwen3。
|
||
|
||
---
|
||
|
||
## 2. Small-Model Speculative（本项目 Phase 22-24 已实现）
|
||
|
||
### 结构
|
||
|
||
- **Draft**: 独立的、小得多的同族 LM。要求：**tokenizer 完全一致**（vocab 也一致）。
|
||
- **Verify**: target 用 batched forward 一次算 γ 个位置的 logits，从左往右比较
|
||
`draft_tokens[i] == target_argmax[i]`，接受最长匹配前缀。
|
||
|
||
### 算法伪代码
|
||
|
||
```python
|
||
for _ in gen_tokens:
|
||
round_start = len(committed)
|
||
# 1. draft γ steps
|
||
draft_tokens = [draft.decode(prev) for _ in range(gamma)]
|
||
# 2. target verify all γ positions in one forward
|
||
verify_logits = target.forward(committed + draft_tokens[:γ])
|
||
# 3. accept longest matching prefix
|
||
accepted = 0
|
||
while accepted < γ and draft_tokens[accepted] == argmax(verify_logits[accepted-1]):
|
||
accepted += 1
|
||
# 4. correction: use target's answer as the next token
|
||
correction = argmax(verify_logits[accepted-1] if accepted>0 else prev_target_logits)
|
||
committed.extend(draft_tokens[:accepted] + [correction])
|
||
```
|
||
|
||
### 优点
|
||
|
||
- **零训练**。任何同 tokenizer 的两个 LM 组合都能跑。
|
||
- 语义正确性直接保证：只要 accept 逻辑严格，输出等价于纯 target greedy。
|
||
- 代码简单，是理解 speculative decoding 最好的教学入口。
|
||
|
||
### 缺点（本项目实测在 dash5 上）
|
||
|
||
Qwen3-0.6B / Qwen3-8B 组合：
|
||
|
||
| γ | acceptance | speedup_e2e |
|
||
|---|---|---|
|
||
| 1 | 66.5% | 0.57× |
|
||
| 4 | 40.3% | 0.49× |
|
||
| 8 | 25.1% | 0.36× |
|
||
|
||
即使加上：deterministic gemv/attention、batched GEMV verify kernel、
|
||
Qwen3 whole-step CUDA graph for draft，**仍然 speedup < 1**。
|
||
|
||
根本原因两点：
|
||
1. **Draft 太贵**：0.6B 一次 decode ~ 2.5 ms，target 8B ~ 12 ms → draft/verify ≈ 20%。
|
||
γ=4 时，draft 4×2.5=10 ms 单独就占了 verify (13 ms) 的 77%。
|
||
2. **Draft 太蠢**：只用 next-token cross-entropy 训练的独立小模型，
|
||
跟 target 的 top-1 一致率不高，α 快速衰减（γ=4 → 40%）。
|
||
|
||
理论上限（假设 verify 免费）：`speedup ≤ (1 + α·γ) ≈ 2.6×`。
|
||
实际上 verify 花掉了绝大部分预算，跑到 0.5× 就到头了。
|
||
|
||
### 什么时候能赢？
|
||
|
||
只有当 `draft_cost / verify_cost < acceptance_rate` 时才可能 >1×。
|
||
Qwen3-0.6B 的 draft_cost 太高，需要 draft 是 target 的 **~1/40** 才行
|
||
（8B target 需要 ~200M draft）。Qwen3 没有官方 200M 的成员。
|
||
|
||
---
|
||
|
||
## 3. EAGLE3（本 Phase 要做的方案）
|
||
|
||
### 3.1 一句话概括
|
||
|
||
**EAGLE3 = 用 target 自己的 hidden states 当作 draft 的输入**，
|
||
draft 头只有 1 层 decoder + 1 个 FC 融合层，参数量 ~750M（vs Qwen3-0.6B 的 1.2 GB），
|
||
且更重要的是：draft 前向**不需要重跑 embedding、不需要多层 attention 累积**，
|
||
成本大约是 target 一次 decode 的 **~1/10**。
|
||
|
||
### 3.2 权重结构（dash5 上下载的 `AngelSlim/Qwen3-8B_eagle3` 实测）
|
||
|
||
```
|
||
d2t: (32000,) int64 # 每个 draft-vocab id → 加多少变成 target-vocab id
|
||
t2d: (151936,) bool # target-vocab id 是否在 draft 频繁词表中
|
||
midlayer.self_attn.q_proj.weight: (4096, 8192) bf16
|
||
midlayer.self_attn.k_proj.weight: (1024, 8192) bf16
|
||
midlayer.self_attn.v_proj.weight: (1024, 8192) bf16
|
||
midlayer.self_attn.o_proj.weight: (4096, 4096) bf16
|
||
midlayer.mlp.gate_proj.weight: (12288, 4096) bf16
|
||
midlayer.mlp.up_proj.weight: (12288, 4096) bf16
|
||
midlayer.mlp.down_proj.weight: (4096, 12288) bf16
|
||
midlayer.hidden_norm.weight: (4096,) bf16 # 融合特征的 pre-attn norm
|
||
midlayer.input_layernorm.weight: (4096,) bf16 # draft 嵌入的 pre-attn norm
|
||
midlayer.post_attention_layernorm.weight: (4096,) bf16
|
||
norm.weight: (4096,) bf16
|
||
fc.weight: (4096, 12288) bf16 # 3×hidden → hidden fusion
|
||
lm_head.weight: (32000, 4096) bf16 # 输出 draft-vocab
|
||
```
|
||
|
||
**关键观察**：
|
||
- `fc.weight (4096, 12288)`：**输入是 target 三个不同层的 hidden state 拼起来**
|
||
（low + mid + high level），一次 FC 融合成 EAGLE 内部的 hidden dim。这是 EAGLE3
|
||
跟 EAGLE1/2 最大的区别（前两代只用 target 最后一层）。
|
||
- `q_proj.weight (4096, 8192)`：**8192 = 4096 × 2**。attention 输入是
|
||
`concat(embed(draft_token), fused_target_hidden)`，两个 4096 拼起来。
|
||
也就是每次预测下一个 token 时，"prompt" 是"上一个 draft token 的 embedding"
|
||
+"target 对上一个位置的隐状态"。
|
||
- `lm_head.weight (32000, 4096)`：**只输出 32000 个高频 token**（vs target 的 151936）。
|
||
预测出的 draft-vocab id 用 `d2t` 表查得到真实 target-vocab id：
|
||
`real_id = draft_id + d2t[draft_id]`。这一步把 lm_head 从 622 MB 压到 131 MB。
|
||
|
||
### 3.3 推理时的数据流
|
||
|
||
```
|
||
target 前向（正常执行）：
|
||
tokens t_0..t_n
|
||
→ embed → layer0 → layer1 → ... → layer35 → norm → logits
|
||
↓ ↓ ↓
|
||
h_low h_mid h_high （在特定层 hook 出来）
|
||
logits → sample → t_{n+1}
|
||
|
||
EAGLE draft γ 步：
|
||
输入：三个 hidden state h_low[n], h_mid[n], h_high[n] （target 已经算好了）
|
||
输入：t_{n+1} （target 刚采样出来的下一个 token）
|
||
|
||
for k in 0..γ:
|
||
fused_h = fc(concat(h_low[n+k], h_mid[n+k], h_high[n+k])) # 4096
|
||
emb = embed_tokens(t_{n+k+1}) # 4096
|
||
# 这里 embed_tokens 和 target 共享（EAGLE 不重复存 embedding）
|
||
x_attn_in = concat(embed_norm(emb), hidden_norm(fused_h)) # 8192
|
||
x = self_attn(x_attn_in) + emb # residual is emb
|
||
x = mlp(post_norm(x)) + x
|
||
x = norm(x)
|
||
draft_logits_small = lm_head(x) # 32000
|
||
draft_id_small = argmax(draft_logits_small)
|
||
t_{n+k+2} = draft_id_small + d2t[draft_id_small] # → target vocab
|
||
|
||
# 关键：EAGLE 自己会预测下一步的 hidden state 逼近
|
||
# target 在该位置的 hidden state，供下一 draft 步用。
|
||
h_low[n+k+1] = h_mid[n+k+1] = h_high[n+k+1] = x
|
||
```
|
||
|
||
**为什么快？**
|
||
1. 只有 1 层 decoder（vs Qwen3-0.6B 的 28 层）。
|
||
2. 每步计算量 = `attn(hidden=4096, kv_heads=8) + mlp(intermediate=12288) + lm_head(V=32000)`
|
||
≈ 1 层 Qwen3-8B decoder + 一个小 lm_head。整个 draft 步 ≈ target 单层 forward + 半个 lm_head，
|
||
远小于 target 完整 forward。
|
||
3. Draft 的 KV cache 也只有 1 层（vs 28 或 36）。
|
||
4. Embedding 表复用 target 的（不重复算）。
|
||
|
||
**Acceptance rate 高的原因**：draft 直接使用了 target 的隐状态，
|
||
不是"用另一个小模型独立猜"，α 通常 ≥70%。
|
||
|
||
### 3.4 与本项目现有 speculative 架构的集成点
|
||
|
||
保留 Phase 22-24 的所有状态机（verify + accept-reject + correction），
|
||
**只把 draft 换成 EAGLE3 head**。API 契约：
|
||
|
||
```rust
|
||
// 现在 (Qwen3 draft)
|
||
let draft_logits = draft_decoder.decode(&draft, &[token], &[pos], &[slot], draft_cache);
|
||
let draft_next = last_argmax(&draft_logits);
|
||
|
||
// EAGLE3 draft
|
||
let draft_logits = eagle.step(&target_hidden_low, &target_hidden_mid, &target_hidden_high, token, pos);
|
||
let draft_next_small = last_argmax(&draft_logits);
|
||
let draft_next = draft_next_small + eagle.d2t[draft_next_small as usize];
|
||
```
|
||
|
||
**新增到 xserv 的东西**：
|
||
1. Target 侧：改造 `Qwen3::decode_core` 让它在特定 3 层（比如 1/3、2/3、末层的
|
||
`post_attention_layernorm` 之后）把 hidden state export 出来。
|
||
2. 新模块 `eagle3.rs`：加载 `AngelSlim/Qwen3-8B_eagle3` 权重，暴露 `step()` 方法。
|
||
3. `bench-speculative` 增加 `--drafter eagle3` 分支，draft 改用 EAGLE head。
|
||
|
||
**不变的东西**：verify path、accept-reject 逻辑、near-tie fallback、CUDA graph
|
||
框架、matched=true 的正确性验证。
|
||
|
||
### 3.5 Acceptance 上限
|
||
|
||
按 EAGLE3 paper 的报告，Qwen3-8B 上 γ=6 acceptance ≈ 0.75，speedup 通常 2-3×。
|
||
本项目实测目标：`speedup_e2e > 1` 是保底，`> 2` 是 stretch goal。
|
||
|
||
---
|
||
|
||
## 4. Multi-Token Prediction (MTP)
|
||
|
||
### 4.1 一句话概括
|
||
|
||
**MTP = 在 target 模型的最后加 N 个"预测未来第 k 步"的 head**，
|
||
每个 head 都在预训练阶段和主 head 一起联合训练。推理时这些 head 天然可以并行
|
||
生成 γ 个 draft，然后主 head 一次前向验证。
|
||
|
||
代表实现：**DeepSeek-V3/R1、Meta MTP 论文（Gloeckle et al., 2024）**。
|
||
|
||
### 4.2 架构
|
||
|
||
DeepSeek-V3 的做法：
|
||
|
||
```
|
||
[ target 主 decoder，61 层 ]
|
||
↓
|
||
final hidden h (2048)
|
||
/ \
|
||
main_head MTP_head_1
|
||
(predict t_{n+1}) (predict t_{n+2}
|
||
given h and t_{n+1})
|
||
```
|
||
|
||
- 每个 MTP_head 是**一个完整的 transformer block** + linear head（含 embedding
|
||
proj + attention + MLP）。
|
||
- 训练时：MTP_head_k 的 target 是 `t_{n+k+1}`，loss 加权求和（DeepSeek-V3 训练时权重 0.3）。
|
||
- 推理时：main_head 得到 `t_{n+1}` 后，用 MTP_head_1 得到 `t_{n+2}`（作为 draft），
|
||
可以级联 MTP_head_2 得到 `t_{n+3}`……然后 target 主前向一次性验证。
|
||
|
||
**DeepSeek-V3 论文**（arxiv 2412.19437）报告：
|
||
- MTP module 1 层，depth=1，参数占总模型 ~2%。
|
||
- MTP accept rate ≈ 85-90%。
|
||
- 端到端 tps 提升 1.8×。
|
||
|
||
### 4.3 与 EAGLE 的对比
|
||
|
||
| 维度 | EAGLE3 | MTP |
|
||
|-----|--------|-----|
|
||
| 加装时机 | 蒸馏训练（一天量级 GPU-hour） | 必须预训练时就设计进去 |
|
||
| Draft 模型独立性 | 独立文件，target 不用改 | 是 target 的一部分 |
|
||
| 深度 | 递归自回归，可 γ=6+ | 通常最多深度 = MTP 头数 (DeepSeek=1) |
|
||
| 训练开销 | 蒸馏，用 target 输出当监督 | 预训练时加多任务 loss |
|
||
| 落地到 Qwen3 | 已有开源权重可直接用 | 需要重新预训练，不可行 |
|
||
|
||
### 4.4 为什么我们不做 MTP
|
||
|
||
- Qwen3-8B 没有预训练的 MTP head。要 MTP 就得**自己重新预训练 Qwen3**，不现实。
|
||
- 若要用现成 MTP，只能换到 DeepSeek-V3 这种自带 MTP 的模型；那对整个 xserv 目标
|
||
(Qwen3 + gpt-oss serving) 是绕道。
|
||
|
||
---
|
||
|
||
## 5. 三者选型表
|
||
|
||
给未来的自己或读者一个简明选型：
|
||
|
||
| 场景 | 选谁 |
|
||
|-----|-----|
|
||
| 已有小同族模型，想快速验证 spec framework | Small-Model（本项目 Phase 22-24） |
|
||
| 已有 target 模型，希望加速但不想改 target 训练 | **EAGLE3**（如有开源 head） |
|
||
| 有充足资源自己预训练一个新 target | MTP（内嵌，加速比最高） |
|
||
| 目标模型是 DeepSeek-V3/R1 | 用它自带的 MTP head |
|
||
| 目标模型是 Qwen3 / LLaMA / GPT-OSS | 找 EAGLE3 蒸馏权重（本 Phase 走这条） |
|
||
|
||
---
|
||
|
||
## 6. 本 Phase 的实施计划
|
||
|
||
1. **写这份文档**（正在做）。
|
||
2. **`xserv-model` 新增 `eagle3.rs`**：定义 `Eagle3Head` 结构，加载
|
||
`AngelSlim/Qwen3-8B_eagle3` 权重。
|
||
3. **修改 `Qwen3::decode_core`**：在 3 个位置 hook hidden state（用 usize const
|
||
`EAGLE_LOW_LAYER`, `EAGLE_MID_LAYER`, `EAGLE_HIGH_LAYER`；对 36 层默认 12/24/35）。
|
||
返回值改成 `(Tensor, Option<[Tensor; 3]>)`，第二个 tuple 只在开启 eagle 时填。
|
||
4. **新增 `Eagle3Head::step(hidden_states, token, pos) -> Tensor`**：一层 attention+
|
||
MLP + lm_head，输出 draft-vocab logits，caller 做 d2t 映射。EAGLE 自己也有
|
||
一个 1-层的 KV cache（每轮 spec 结束时清空）。
|
||
5. **`bench-speculative` 加 `--drafter [qwen3|eagle3]` 开关**。EAGLE 分支复用现有
|
||
verify+accept 逻辑，只替换 draft 环节。
|
||
6. **γ 扫**：预期 γ=6 时 acceptance > 0.7、speedup_e2e > 1.5×。
|
||
|
||
## Sources
|
||
|
||
- EAGLE-3 paper (arxiv 2503.01840): "Scaling up Inference Acceleration of Large Language Models via Training-time Test"
|
||
- SafeAILab/EAGLE GitHub: reference implementation
|
||
- AngelSlim/Qwen3-8B_eagle3 on ModelScope/HuggingFace: pre-trained head we're using
|
||
- DeepSeek-V3 Technical Report (arxiv 2412.19437): MTP architecture
|
||
- Gloeckle et al. 2024 "Better & Faster Large Language Models via Multi-token Prediction"
|