免草稿模型!Orthrus 用扩散注意力模块实现 LLM 并行生成,提速不减质

你正在跑本地大模型,每次生成一个 token 都要等,想提速就得用“投机解码”(Speculative Decoding)——但代价是额外加载一个小模型、多占一份 KV 缓存,还要担心草稿模型猜不准分布导致接受率下降。Orthrus 是另一种路线:不挂草稿模型,在冻结的自回归 Transformer 每一层里塞一个可训练的扩散注意力模块,让它一次性并行吐出 32 个 token。关键点:输出分布和目标模型逐 token 采样一模一样,数学上保证。

对开发者来说最吸引人的数据:一次前向传播生成 32 个 token,而模型输出分布和原始自回归采样完全一致。如果工程实现能验证这个数学性质,你就能获得并行生成的速度,同时省掉投机解码里“草稿模型和主模型是否一致”的反复校验。

目前还是早期研究,不是 pip install 就能用的库。但它的架构思路值得理解——它指向了自托管推理的另一个设计空间:加速来自模型内部,而不是来自旁边挂个草稿模型。

Orthrus 怎么同时生成多个 token

基座 Transformer 保持冻结。Orthrus 在每一层插入一个扩散注意力模块,这个模块操作一组“占位位置”——在已发布配置中是 32 个未来 token。推理时,扩散模块通过少量去噪步骤逐步把这些占位符变成具体 token,每一步都复用当前层的激活值。

“输出分布完全保留”是它最不寻常的地方。投机解码通过拒绝采样实现分布保留:草稿模型提议、目标模型验证、不匹配的回滚。Orthrus 用不同的机制达到同样保证:扩散模块以冻结模型的隐藏状态为条件,并把这些隐藏状态作为收敛信号,因此产生的输出等同于自回归模型以相同温度逐 token 采样会得到的结果。代价从“草稿有时错,少接受几个 token”变成了“去噪可能需要更多步才能收敛”。

共享 KV 缓存是让它适合自托管部署的关键。投机解码的实现(比如 Medusa、Eagle)通常需要单独的草稿缓存,或者在主缓存上额外加上草稿条目。Orthrus 直接复用冻结模型的 KV 缓存,让内存占用更接近单个模型而不是模型加草稿对。

Orthrus 目前在研究论文中有描述,还未作为生产库发布。下面的性能数据来自架构设计和报告配置,不是独立基准测试的结果。请当作待验证的假设,在计划基础设施之前先用你自己的负载做测试。

和投机解码相比怎么样

投机解码已在生产中使用一段时间了。vLLM、TensorRT-LLM、llama.cpp 都支持某种形式。工作机制:你加载一个小草稿模型(有时是调优的 Medusa 头,有时是独立的 1B 模型),草稿模型提议 K 个 token,目标模型一次前向传播验证全部 K 个 token,运行时接受最长的匹配前缀。

Orthrus 改了这些部分:

  • 草稿费用:无需加载、训练或维护单独的模型。扩散模块作为基座模型的一部分嵌入。
  • KV 内存:和基座模型共享,不会被一个侧边草稿模型翻倍。
  • 接受行为:输出从结构上就和基座自回归采样分布一致,而不是通过拒绝采样概率性一致。
  • 训练成本:扩散注意力模块需要为每个基座模型训练一次。这不是免费的,但会被该检查点的每次部署分摊掉。

在没有针对知名基座模型(如 Llama、Qwen、Mistral)的公开实现和可复现的硬件基准测试之前,很难量化比 Eagle-2 或 Medusa-2 快多少。架构论证很有力,实证对比还在待办中。

如果你在自托管,这意味着什么

如果你在运行一个本地大模型用于开发者工具,关心的延迟是「首 token 时间」+「解码侧每秒 token 数」。投机解码主要加速解码侧。Orthrus 瞄准同样的指标,但成本曲线不同。

几个值得跟踪的实际问题:

  • 量化:大多数自托管环境在 4-bit 或 8-bit 权重的量化模型上运行。训练好的扩散模块能否扛住激进量化还是个开放问题——用 fp16 训练的模块不一定能在 GPTQ 或 AWQ 后干净地回退。
  • 批大小交互:投机解码的加速会随着批大小增大而缩小,因为验证器前向传播已经饱和了算力。Orthrus 的并行块生成和批处理之间的交互依赖于去噪步骤的调度方式,现有材料还没有多批次的对比。
  • 长上下文解码:32 个 token 的块对短回答是容易的。几千 token 的输出需要 100 多个块连续生成,每块的收敛成本在那个场景下比峰值并行度更重要。

如果你在用 AI 编程工具,并且它跑在本地推理服务器上,这类技术带来的端到端延迟改善正是让本地模型在编辑延迟上能与云 API 竞争的关键。

注意事项和缺失部分

围绕这个架构的讨论清晰地列出了尚未回答的问题:
– 没有针对流行基座模型(Llama、Qwen、Mistral)发布的检查点,开发者无法直接装入现有推理运行时。
– 没有和 Eagle-2 或 Medusa-2 在相同硬件和提示分布上的正面基准对比。
– 没有在工具使用或函数调用输出上的文档化表现——这类提示往往是投机解码表现最差的场景,因为下一个 token 的分布受到结构约束。

这些都不是对研究的否定——这是早期架构的常见差距。但如果你正在规划未来两个季度的自托管 LLM 基础设施,投机解码依然是默认选择。Orthrus 值得跟踪,但还不到押注的阶段。


原文发布于 pickuma.com

类似文章