docs/next-steps-draft-model-speculative.md

# Draft Model Speculative Decoding — 执行规格 v1

**目标模型**：Qwen3-0.6B（draft）+ Qwen3-235B-A22B（verifier，当前主模型）
**平台**：Ascend 910 × 16 NPU
**预期收益**：全 prompt 类型 +50~100% TG（目标 40-50 t/s 稳定）
**工程量**：1-2 周

---

## 0 · 为什么换 draft model

| 指标 | PLD (n-gram) | Draft Model (0.6B) |
|---|---|---|
| Accept rate 均值 | 0.3-1.5 (低) | 3-5 (预期) |
| 事实/代码 prompt | 需 guard 丢弃大量 draft → 几乎失效 | 语义理解 → 稳定 accept |
| 正反馈循环风险 | 高（n-gram 放大重复） | 低（小模型有自己的分布，不会完美复制主模型重复） |
| 内存占用 | ~0 | ~1.2GB BF16 |
| 计算成本 | ~0 | 每 draft step ~10-15ms |

n-gram PLD 已到方法论天花板；accept=0.3-1.5 意味着平均每次 verify 只赚 1-2 个 token，
且 guard 在事实/代码场景丢弃几乎所有 draft。

---

## 1 · 硬件 / 资源规划

### 1.1 NPU 分配决策

**选项 A**（推荐）：主模型 TP=16 不变，draft 模型**时分复用 device 0**
- 优点：无需改动主模型 sharding
- draft decode ~10ms，主模型 decode ~35ms → 串行总 45ms 仍然值得（vs pure decode 35ms）
  因为每次 draft 后的 batch verify 摊销到 3-5 token → 50/5 = 10ms/token

**选项 B**：主模型 TP=15，draft 独占 device 15
- 优点：draft / 主完全并行
- 缺点：主模型所有 sharding 常量需重算 `64/15` 不整除问题严重

**选项 C**：主模型 TP=8（仅用 device 0-7），draft 独占 device 8
- 优点：配置简洁
- 缺点：主模型从 TP=16 降到 TP=8，每 rank 权重加倍，可能 OOM（Qwen3-235B BF16 约 470GB / 8 = 58GB/rank，910 单卡 64GB 勉强够）

**→ 选 A**。draft decode 可重叠的时机有限，但工程风险最低。

### 1.2 权重下载 / 存储
- HF：`Qwen/Qwen3-0.6B`（约 1.2GB BF16）
- 位置：`/run/modelscope/Qwen3-0.6B-BF16/`（与主模型同目录结构）
- tokenizer **完全相同**（共享 Qwen3 vocab.bin）—— 这是选 0.6B 的关键优势

---

## 2 · 架构设计

### 2.1 模块

```
include/
├── draft_runner.h           ← 新增：TP=1 Runner 特化，device 独占
├── draft_engine.h           ← 新增：Qwen3 28 层 attn+MLP forward (非 MoE)
├── draft_model_config.h     ← 新增：Qwen3-0.6B hparams
src/
├── draft_runner.cpp         ← 简化 runner.cpp，去 TP/HCCL
├── main_cli.cpp             ← 新增 --draft-model-dir flag + spec decode loop
```

### 2.2 主循环流程（替换现有 PLD 分支）

```
while (!eos && decoded < n_predict) {
    // Phase 1: draft 模型连续 decode K 步
    draft_tokens = []
    draft_runner.set_cursor(main_runner.past_len + 1)  // 同步位置
    cur = next_id
    for k in 0..K-1:
        d_logits = draft_runner.decode(cur)            // ~10 ms (device 0 独占)
        cur = argmax(d_logits)
        draft_tokens.append(cur)
        if is_eos(cur): break
    
    // Phase 2: 主模型 batch verify (K+1 位置)
    batch_input = [next_id] + draft_tokens
    batch_logits = main_runner.decode_batch(batch_input)  // ~42 ms
    
    // Phase 3: 接受最长前缀 + 1 bonus
    accept = 0
    for i in 0..K-1:
        if argmax(batch_logits[i]) == draft_tokens[i]:
            accept++
        else:
            next_id = argmax(batch_logits[i]); break
    if accept == K:
        next_id = argmax(batch_logits[K])  // bonus
    
    // Phase 4: rewind 两个模型的 KV cache
    main_runner.rewind(K - accept)
    draft_runner.rewind(K - accept + (accept == K ? 0 : 1))  // draft 多走了一步
    
    // Phase 5: emit accepted + bonus
    for tok in draft_tokens[:accept] + [next_id]:
        emit(tok); hist.append(tok)
}
```

### 2.3 关键 invariant

| invariant | 保障 |
|---|---|
| draft 与 main 看同一 token 序列 | 每次 decode 前 `draft.position = main.past_len`，rewind 同步 |
| device 0 时序：draft decode 不冲突主模型 | 主模型 rank 0 在 HCCL 等待期间 device 0 空闲 → draft 用这段时间 |
| KV cache 独立 | draft 与 main 各有自己的 KV cache，大小分别按 0.6B / 235B |

### 2.4 device 0 争用分析

主模型每 decode step ~35ms，其中：
- rank 0 的 local compute ~10ms
- HCCL AllReduce 等待 ~22ms（device idle from rank 0 POV）
- final logits ~3ms

**理论上 draft decode 可在 HCCL 等待期间占用 device 0**。但实际：
- HCCL 等待期间 rank 0 的 stream 是被 HCCL 占用的（不是 idle）
- draft 需要独立的 stream 才能真正并行
- 若用 `aclrtCreateStream` 分出第二条 stream，draft 和主的 HCCL 共享 AI Core 资源 → 可能互相干扰

**保守估算**：串行执行 draft (K=4 步约 40ms) + verify (42ms) = 82ms per spec cycle，产出 3-5 token → amortized 16-27 ms/token → 37-60 t/s（对比当前 ~27 t/s 单 decode）。

---

## 3 · 实施步骤（按依赖排序）

### M1 · draft_model 独立推理（3 天）
- [ ] `draft_model_config.h`：Qwen3-0.6B hparams（28 层，16 Q head，8 KV head，hidden=1024, ...）
- [ ] `draft_engine.h`：simplified attention + MLP(SwiGLU)，无 MoE
- [ ] `draft_runner.cpp`：单 device，无 HCCL，复用 RoPE cache / workspace pool
- [ ] `test_draft_standalone.cpp`：与 HF Transformers 参考比对 rel ≤ 1e-3
- [ ] `test_draft_consistency.cpp`：同 prompt 多次 argmax 确定（greedy 应稳定）

**验证点**：draft 单模型 decode ≥ 80 t/s（TP=1，0.6B 应该很快）

### M2 · 协同控制（3 天）
- [ ] `main_cli.cpp` 新增 `--draft-model-dir` flag
- [ ] 主/draft cursor 同步 wrapper
- [ ] draft.rewind() 实现（沿用 main 的 rewind_cache 模式）
- [ ] 串行版 spec decode loop（不追求重叠）
- [ ] **正确性验证**：--draft 路径与 --pld 关闭路径对相同 prompt 应产生**完全相同**的 token 序列（greedy temp=0）
  - 这是黄金标准：任何偏差都是 bug
  - 用 `test_spec_correctness.sh` 自动回归

**验证点**：5 个 prompt（故事/代码/事实/数学/对话）输出与 baseline 逐 token 一致

### M3 · 性能测量与调优（3 天）
- [ ] K sweep：K ∈ {2, 4, 6, 8}，找 best accept × throughput
- [ ] 与 PLD+guard 对比：同 prompt 3 runs
- [ ] 文档更新：诚实对比表

**验证点**：
- 所有 prompt OK verdict（bench_pld_safe.sh 分类器不报 degraded）
- mean TG 比 base +40% 以上

### M4（可选）· 并行化 / 异步 draft（2-3 天）
- [ ] 第二条 aclrtStream 用于 draft
- [ ] 尝试 draft decode 与主 verify 的下一轮重叠
- [ ] 若测不出显著增益则砍掉，保持串行

---

## 4 · 正确性测试计划

**绝对底线**：greedy (temp=0) 下，启用 draft 的输出必须与不启用**逐 token bit-identical**。
理由：spec decoding 的正确性定义就是 "输出等价于 verifier 单独贪心解码"。

### 4.1 回归测试套件

```bash
# test_spec_correctness.sh
for prompt in "Paris" "Python function" "Once upon" "2+2" "长句测试":
    out_base = run --no-pld --temperature 0
    out_draft = run --draft-model-dir --temperature 0
    diff out_base out_draft  # must be identical
```

### 4.2 压力测试
- 长 prompt (>1K token) 稳定性
- 多轮对话下 cache rewind 正确
- EOS 在 draft 和 main 不一致时的处理

### 4.3 benchmark 诚实性
- 继续用 `bench_pld_safe.sh` + 新增 `--draft` 模式列
- OK/degraded 分组统计；degraded = fatal

---

## 5 · 风险与兜底

| 风险 | 概率 | 兜底 |
|---|---|---|
| device 0 独占争用让 draft 延迟 >15ms | 中 | 保持串行执行，不追求重叠 |
| draft accept rate 不如预期（<3） | 中 | K 调小至 2-3，仍有 +20-30% 收益 |
| rewind 导致 KV cache 污染 | 高 | 严格测试 rewind 正确性；失败就回滚到"不 rewind，允许小概率错" |
| 910 初代 0.6B 权重加载失败 | 低 | 已验证 Qwen3-235B BF16 可加载，0.6B 结构子集 |
| 与 PLD guard 机制冲突 | 低 | draft 替代 PLD，同时不开启两个路径 |

---

## 6 · 时间表（单人 1 周密集 / 2 周分散）

| 天 | 任务 | 产出 |
|---|---|---|
| D1-D3 | M1 draft_standalone | test_draft_standalone PASS |
| D4-D6 | M2 协同 + 正确性 | test_spec_correctness 5/5 PASS |
| D7-D8 | M3 性能调优 | K sweep + 对比表 |
| D9-D10 | M4 可选重叠 + 文档 | 最终报告 |

---

## 7 · 决策 checkpoint

**第 3 天末**：若 draft 单模型 decode < 60 t/s（远不如预期），暂停，重新评估 0.6B 是否太慢 / 切换到 Qwen3-1.7B / 或放弃路径。

**第 6 天末**：若 5/5 prompt 正确性测试过不了，停下查 bug，不急于测性能。

**第 8 天末**：若 mean TG 提升 < 20% vs 当前 PLD+guard，判 ROI 不足，回滚保留 PLD+guard 为主路径。

---

*此规格服务于下一轮会话，作为单人专职工作的路线图。*