从研究到生产:EAGLE-3在Vertex AI加速开源LLM推理2-3倍

推测解码(Speculative Decoding)能显著提升LLM推理速度,但传统方法需额外训练低效的草稿模型。Vertex AI采用EAGLE-3创新方案,仅在目标模型内部层添加轻量级草稿头(仅占2-5%模型大小),简化训练并实现2x-3x解码加速。本文详述从数据清洗、嵌入生成、训练到使用SGLang在Vertex AI大规模部署的完整管道,分享关键挑战与经验教训,包括合成数据管道、离线训练、聊天模板应用、掩码处理及Zero-Overhead Overlap Scheduler优化。基准测试显示,在Llama 4 Scout 17B上,EAGLE-3显著降低TPOT并提升吞吐量。(128字)

EAGLE-3 Vertex封面图

对于从事LLM开发的从业者来说,一次只生成一个token的瓶颈是常见痛点。标准的自回归生成本质上是顺序过程,这导致内存带宽受限的计算模式:GPU核心闲置,而时间主要耗费在为每一步从内存读取海量模型权重。

解决方案便是speculative decoding(推测解码)。这项优化技术通过引入草稿机制,加速大型LLM(目标模型)逐token生成的缓慢过程。草稿机制快速提出多个后续token,目标模型则并行批量验证这些提议,接受最长匹配前缀并从中继续生成。

但并非所有草稿机制都高效。经典的草稿-目标方法需单独部署小型LLM作为drafter,增加托管资源和成本。

传统草稿模型架构

这时,EAGLE-3(Extrapolative Attention Guided LEarning)登场。它无需完整独立模型,而是直接在目标模型内部层附加极轻量级draft head(仅2-5%目标模型大小)。该头在特征和token层面运作,利用隐藏状态推断未来token树状结构。

结果?保留推测解码所有优势,同时消除训练和运行第二模型的开销。相较复杂多亿参数草稿模型,EAGLE-3仅训练轻量头,针对Llama 70B等模型实现2x-3x解码加速(视多轮对话、代码、长上下文等工作负载而定)。

EAGLE-3目标模型架构

将EAGLE-3从论文推向大规模生产云服务,仍需克服诸多工程挑战。本文分享技术管道、关键问题及宝贵经验。

挑战1:数据准备

EAGLE-3头需训练,但公共数据集常存问题:

  • 严格使用条款:生成模型禁止用于竞争性开发。
  • PII污染:包含姓名、位置、金融标识等。
  • 质量无保障:仅适配demo,非客户生产负载。

经验1:构建合成数据生成管道

解决方案:根据客户场景选优质数据集,提取用户prompt,经DLP和PII过滤,应用chat template,tokenize后输入目标模型(如Llama 3.3 70B)生成响应。此法产出合规、干净且匹配模型分布的数据,完美适配draft head训练。

数据生成管道

挑战2:工程化训练管道

数据输入方式分online training(实时生成嵌入)和offline training(预计算嵌入)。我们选offline,因硬件需求低:预计算特征/嵌入存GCS,用于训练。鉴于EAGLE-3头小巧,初始训练仅需单主机一天;数据集扩展后增至数天。

训练管道

经验2:聊天模板不可或缺

指令调优模型训练中,若未用目标模型chat template(如Llama 3),嵌入错误,头将学习偏差分布。

经验3:注意掩码

训练输入含prompt和response,但EAGLE-3仅预测response。需在损失函数手动掩码prompt,否则头浪费容量预测已知prompt,性能下降。

掩码处理示意

挑战3:服务与扩展

经验4:服务框架至关重要

与SGLang团队合作,将EAGLE-3高效部署生产。SGLang的tree attention kernel专为EAGLE-3的draft tree(分支路径)并行验证设计,避免性能损失。

经验5:勿让CPU拖累GPU

EAGLE-3加速后,CPU开销(如kernel启动、元数据管理)成新瓶颈。同步调度器中,GPU Draft后闲置待CPU Verify准备。

普通调度器

SGLang的Zero-Overhead Overlap Scheduler解决此问题:利用FutureMap,CPU并行准备下一Draft/Extend,而GPU执行Verify,消除闲置,实现额外10%-20%加速

重叠调度器

基准结果

成果显著:在SGLang上,用Llama 4 Scout 17B Instruct基准,EAGLE-3实现2x-3x解码延迟降低及吞吐提升

指标1:中位输出Token时间(TPOT)

TPOT基准图

绿色EAGLE-3线在所有并发度下,TPOT均低于蓝色基线,延迟更优。

指标2:输出吞吐量

吞吐量基准图

绿色线在并发度下持续大幅超越基线,彰显EAGLE-3优势。可查完整notebook自行基准。