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 回归测试套件

# 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 为主路径。

---

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