ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

如何为SSD扩展新模型?Qwen3/Llama3实现与EAGLE3 Draft开发指南

如何为SSD扩展新模型?Qwen3/Llama3实现与EAGLE3 Draft开发指南

【免费下载链接】ssd

A lightweight inference engine supporting speculative speculative decoding (SSD).

项目地址:https://gitcode.com/gh_mirrors/ssd10/ssd
点击查看免费下载

SSD(Speculative Speculative Decoding)是一个轻量级 LLM 推理引擎,通过“推测+验证并行”的投机解码(Speculative Decoding)算法实现无损加速。本文以 SSD 中 Qwen3 与 Llama3 两个现成模型实现为蓝本,手把手讲解扩展新模型的完整流程,并深入剖析 EAGLE3 Draft 模型的架构设计与权重加载要点,帮助你在自己的模型上复现 4.0x 解码加速。

一、先看懂 SSD:一张图理解核心算法

SSD 投机推测解码架构与解码速度对比

传统投机解码(SD)中小模型起草(Draft)和大模型验证(Verify)在同一硬件上串行执行;而 SSD 让两者在不同硬件上并行:小模型提前预判大模型可能的验证结果,并对所有分支同时投机。猜对了,起草开销直接归零。图中右侧柱状图展示了 AR、SD、SSD 三种方案的解码速度对比,SSD 达到约 4.0x 加速。

引擎层面,SSD 已支持 Tensor Parallelism、PagedAttention、CUDA Graphs、torch 编译与前缀缓存,模型扩展只需关注"模型本体 + 权重加载 + 注册"三件事。

二、新模型落地只需改动的 4 个位置

环节文件职责
模型定义ssd/models/qwen3.py / ssd/models/llama3.py5 层结构的 torch Module 实现
模型注册ssd/engine/model_runner.py按model_type分发模型类
权重加载ssd/utils/loader.pysafetensors → 参数映射(含 packed 权重)
路径配置ssd/paths.py模型快照目录的环境变量

三、第一步:读懂 Qwen3 的 5 类结构

每个模型文件都遵循统一的分层套路(以 Qwen3ForCausalLM 为例):

  1. Qwen3Attention:QKV 投影(QKVParallelLinear)→ 逐 head 的 QK RMSNorm(RMSHeadNorm)→ RoPE → Attention
  2. Qwen3MLP:gate/up 合并投影 +SiluAndMul激活 + 下行投影
  3. Qwen3DecoderLayer:pre-norm 结构,residual 贯穿流式传递
  4. Qwen3Model:Embedding + N 个 DecoderLayer + 最终 Norm,forward(input_ids, positions) → hidden_states
  5. Qwen3ForCausalLM:加上lm_head与compute_logits,并声明packed_modules_mapping

⚠️ 注意:构造参数中必须透传投机解码相关开关——draft、speculate、spec_k、async_fan_out、draft_async以及张量并行参数tp_group/tp_size,这些是 Attention 层支持 SSD 树状注意力的前提(见 Qwen3Attention 构造)。

Qwen3 与 Llama3 的差异点(对照 LlamaAttention):

  • Qwen3 有 QK Norm,Llama3 没有
  • rope_theta:Qwen3 默认 1,000,000,Llama3 默认 500,000
  • Qwen3 保留rope_scaling,Llama3 主动置空
  • qkv_bias:Qwen3 支持(来自 config),Llama3 恒为 False

四、第二步:注册模型并配置路径

引擎通过 HuggingFace config 的model_type字段分发模型类,逻辑在 setup_and_warmup_model_and_cudagraphs:

  • model_type == 'llama'→LlamaForCausalLM
  • model_type == 'qwen3'→Qwen3ForCausalLM
  • 其他类型直接抛ValueError: Unsupported model type

因此新模型接入只需:导入你的模型类,在分发处加一个elif分支。同时按 paths.py 的样式,用环境变量(如SSD_TARGET_MODEL/SSD_DRAFT_MODEL)指向本地快照目录(需包含config.json),并在 bench_helpers.py 的size_to_model映射中登记模型名。

五、第三步:让权重"对号入座"

loader.py 的加载规则很简单:

  • 权重名若命中packed_modules_mapping(如q_proj → qkv_proj.q),走分片合并加载——这是把 HF 里 4 个独立矩阵合并成 2 个 packed 矩阵的关键
  • 否则按参数名直接copy_

所以新模型的两条铁律:参数命名与 HF 权重 key 一致(或提供映射表);packed_modules_mapping覆盖所有被合并的投影层。

六、EAGLE3 Draft 模型开发要点

Eagle3DraftForCausalLM 是 SSD 里最精巧的部分,与"完整小模型当 Draft"路线不同,EAGLE3 只用单层 Transformer做投机:

  1. 目标层激活采集:LlamaModel前向时在指定的eagle_layers(默认[2, L//2, L-3],见 config.py)截获残差流,拼成3 × d_model_target维的特征
  2. fc 投影:nn.Linear(3*d_model_target, hidden_size)把目标特征压到 Draft 宽度;后续步则复用 Draft 自身的输出做自条件
  3. 双路拼接输入:Eagle3Attention 的qkv_proj输入维度是2 * hidden_size——token 嵌入与条件特征各自过 RMSNorm 后拼接
  4. 小词表加速:lm_head只覆盖draft_vocab_size,通过d2t/t2d映射表把小词表 logits 散射回完整目标词表(compute_logits)

权重加载同样有专门通道:load_eagle_model 会把midlayer.*前缀映射到model.layer.*、norm.weight映射到final_norm.weight,并在 Draft 权重缺失 embedding 时自动从目标模型目录回补(hidden_size 需一致)。

七、验证你的新模型

引擎自检与基准测试都在 bench/ 下,--eagle会启用 EAGLE3 自动选 Draft(如 8B 用EAGLE3_YUHUILI_8B,70B 用 SpecForge 版本,见 bench_helpers.py):

cd bench # 同步投机解码:70B 目标 + 1B Draft python -O bench.py --llama --size 70 --gpus 4 --spec --k 6 --b 1 --temp 0 --numseqs 128 --output_len 512 --all # EAGLE3 + 异步 SSD(注意:当前限制 Llama 目标、temp=0、必须 --async) python -O bench.py --llama --size 70 --gpus 5 --eagle --async --k 7 --f 3 --b 1 --temp 0 --numseqs 128 --all

Qwen 目标则用--qwen --size 32。当前限制提示:EAGLE 在 bench 脚本中仅对 Llama 目标开放(bench.py 的 assert),而Qwen3ForCausalLM也显式禁止了use_eagle参数(qwen3.py)——若你要扩展 Qwen3 的 EAGLE Draft,这正是需要修改的入口。

八、常见坑清单 🛠️

  • model_type没注册:启动即报Unsupported model type,先补 model_runner 分发
  • rope_theta 不一致:EAGLE Draft 的 config 常带错误默认值,config.py 会自动用目标模型值覆盖,自定义加载时别漏掉
  • packed 权重漏映射:显存加载"成功"但输出乱码,优先检查packed_modules_mapping
  • use_eagle用错位置:它必须配 Draft 侧的Eagle3DraftForCausalLM,配到LlamaForCausalLM上会触发 assert
  • 始终用python -O运行基准,关闭 debug 开销

📌 扩展路径总结:仿写ssd/models/下的 5 类结构 → 在model_runner.py注册 → 对齐权重命名与 packed 映射 →paths.py配路径 → bench 验证。掌握这套范式,GPT-OSS、Kimi-K2.5 等新模型(README Roadmap 中的规划)接入 SSD 就是水到渠成的事。

【免费下载链接】ssd

A lightweight inference engine supporting speculative speculative decoding (SSD).

项目地址:https://gitcode.com/gh_mirrors/ssd10/ssd
点击查看免费下载

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进