我是怎么理解 Orthrus 这套并行生成实现的

最近在读 Orthrus 这个仓库，先说结论：这是一个想把大语言模型生成提速到大约 $4\times$ 到 $5\times +$，同时又尽量保持原模型分布不变的并行解码方案。 README 里给出的结果是，Qwen3-1.7B、4B、8B 版本平均加速大约分别到 $4.25\times$、$5.20\times$、$5.36\times$，而且它主打的不是“牺牲一点效果换吞吐”，而是 strictly lossless generation。

如果你平时看过 speculative decoding 或 dLLM 相关工作，就会知道这里真正有吸引力的不是“并行生成”四个字本身，而是这几个条件同时出现：

• 生成要快
• 输出分布不能乱
• 显存开销不要再多挂一套 draft model

很多并行解码方案只能同时满足其中两条。要么快，但是会偏离原模型分布；要么保真，但是要额外维护一个草稿模型；要么实现上能跑，但一到长上下文，缓存开销就开始难看。Orthrus 让我觉得值得细看的地方就在这里：它试图在一个模型内部，同时把速度、保真和缓存成本这三件事一起处理。

Orthrus 是什么？如果用一句尽量直白的话说，它是一个 把自回归语言模型改造成“双视角生成器” 的方案。它还是基于原来的 backbone，这个仓库里用的是 Qwen3，但推理时不再只有一条严格串行的 next-token 路径，而是多出一条并行 proposal 路径。于是同一个模型在生成时就有两条路：

• 一条是正常的自回归路径，负责严肃地按因果顺序做 next-token prediction
• 一条是 diffusion 风格的并行路径，负责一次性多猜几个 token

然后再由自回归路径去验证并行路径的 proposal。关键点在于，这两条路径不是两套模型，而是同一个模型内部的两种前向方式，并且共享同一份历史 KV cache。也正因为这样，它才有资格去挑战两个老问题：一是 speculative decoding 额外 draft model 的系统成本，二是很多并行生成方案在 fidelity 上的漂移。

所以这篇文章我不打算复述论文摘要，而是直接回答三个更实际的问题：

1. Orthrus 到底是什么，为什么它的结果值得看
2. 它为什么不需要再挂一个 draft model
3. 它怎么在代码里把“并行 proposal + 原模型校验 + 分布修正”串成一个生成循环

下面我按“问题 -> 直觉 -> 公式 -> 代码”的顺序拆开讲。

1. 它到底想优化什么

先看最普通的自回归生成。

给定前缀 $x_{<t}$，第 $t$ 个 token 的生成分布是：

$$p(x_t\mid x_{<t})$$

整段序列的分解是：

$$p(x_{1:T})=\prod _{t=1}^{T}p(x_t\mid x_{<t})$$

这件事的好处是严格、稳定、分布定义清楚；坏处也很明显：第 $t+1$ 个 token 必须等第 $t$ 个 token 出来以后才能算。所以推理天生是串行的。

如果一次 forward 只能确认 1 个 token，那么总步数大约就是：

$$N_{\text{steps, AR}}\approx T$$

这就是吞吐瓶颈的来源。

很多加速方案的思路是：先让一个便宜模型一次提议多个 token，再让大模型验证。这类方法可以把步数从 $T$ 降到接近：

$$N_{\text{steps}}\approx \frac{T}{L_{\text{accept}}}$$

其中 $L_{\text{accept}}$ 是每轮平均能接受多少个 token。

问题是，这类方法通常要额外引入一个 draft model。于是你虽然减少了步数，但又多了一套模型参数和一套 KV cache。在长上下文里，缓存的系统成本经常比“少跑几步”更敏感。

我理解 Orthrus 的第一原则就是：不要再引入第二个模型。

2. Orthrus 的直觉其实很简单

我先用一句不那么论文化的话概括它：

不要让小模型替大模型猜，而是让大模型自己在内部切出一个“并行猜测模式”。

这件事落到仓库里，就是 src/model.py 里的两条路径：

• is_diffusion_pass=False：正常自回归路径
• is_diffusion_pass=True：并行提案路径

也就是说，Orthrus 不是：

$$\text{target model}+\text{draft model}$$

而更像：

$$\text{one model}=\text{AR view}+\text{diffusion view}$$

这两个 view 共享同一套主体结构，但在 attention 里有各自的投影参数。这个差别非常重要，因为它直接决定了后面最关键的一件事：两条路径可以共享同一份高保真的 KV cache。

3. 代码里怎么体现“同一个模型，两种视角”

我读这个仓库时，第一个真正让我抓到重点的地方，是 OrthrusAttention 的定义，见 src/model.py。

它不是又封了一套完整模型出来，而是在 attention 里额外放了一组 diffusion 专用投影：

self.q_proj = ...
self.q_proj_diff = ...
self.k_proj = ...
self.k_proj_diff = ...
self.v_proj = ...
self.v_proj_diff = ...
self.o_proj = ...
self.o_proj_diff = ...

这意味着：

• AR 路径继续用原本的 q_proj / k_proj / v_proj / o_proj
• diffusion 路径切到 *_diff 这组参数
• embedding、MLP、层堆叠、RMSNorm 等主体结构仍然共享

如果写成一个很粗糙的形式，可以把它理解成：

$$h_{\text{AR}}^{(l+1)}={\text{Block}}_{\text{AR}}(h^{(l)})$$

$$h_{\text{Diff}}^{(l+1)}={\text{Block}}_{\text{Diff-Attn}}(h^{(l)})+\text{Shared-MLP}(h^{(l)})$$

这里真正被“分叉”的核心是 attention 视角，不是整套 Transformer。

所以我觉得 README 里“只微调部分参数就注入并行能力”这件事，不能只当成一句训练口号看。它在代码里的真实含义是：并行能力主要靠 attention 视角扩展注入，而不是复制一整套推理系统。

4. 它为什么能省内存：因为 KV cache 是共享的

这个仓库最有价值的实现细节，我觉得不是生成循环本身，而是 diffusion 路径怎么使用 cache。

看 src/model.py 这一段：

shared_cache = past_key_values.layers[self.layer_idx]
shared_key_states = shared_cache.keys
shared_value_states = shared_cache.values

key_states = torch.cat([shared_key_states, key_states], dim=2)
value_states = torch.cat([shared_value_states, value_states], dim=2)

这里的意思非常直接：

1. AR 路径已经把历史前缀的 K/V 算好并存进 past_key_values
2. diffusion 路径不会再重建一份“草稿缓存”
3. 它直接拿 AR 路径已经确认过的 K/V 来当上下文
4. 然后只把当前 draft block 新产生的 K/V 拼接到后面

如果把历史前缀长度记成 $n$，当前并行 block 长度记成 $b$，那么 diffusion 路径在注意力里看到的 key/value 序列长度就是：

$$n+b$$

但这里前 $n$ 个不是新算一遍，而是直接复用 AR cache。

所以从系统角度看，它节省的不是某一层的一个小常数，而是避免了这件事：

$${\text{KV}}_{\text{AR}}+{\text{KV}}_{\text{draft}}$$

变成：

$${\text{KV}}_{\text{shared}}+{\text{KV}}_{\text{current block}}$$

如果只讲一句人话，那就是：Orthrus 快不快先不说，至少它没有为了并行提案再供一套历史缓存。

这也是它和传统 speculative decoding 最大的工程差异之一。传统 speculative decoding 更像：

• 一个模型负责便宜地提案
• 另一个模型负责昂贵地确认

Orthrus 则更像：

• 同一个模型的并行视角先提案
• 同一个模型的 AR 视角再确认

两者的“验证”这件事表面类似，但系统形态差很多。

5. 整个生成流程其实可以画成一个很直白的循环

这个仓库最值得读的函数，是 src/model.py 的 OrthrusLM.generate()。

我把它翻译成最朴素的流程，大概就是下面这样。

第一步：先按正常 AR 路径把 prompt 跑完

给定输入前缀 $x_{1:n}$，先跑一次标准模型，得到：

$$p(\cdot \mid x_{1:n})$$

然后采样出下一个 token：

$$x_{n+1}\sim p(\cdot \mid x_{1:n})$$

这一步之后，历史前缀的 KV cache 已经建立好了。

第二步：构造一个 block 让 diffusion 路径并行猜

如果 block size 是 $b$，那 Orthrus 会构造一个输入：

$$[x_{n+1},m,m,\dots ,m]$$

这里 $m$ 是 mask_token_id。

在代码里对应的是：

diff_block_ids = torch.full((1, diff_len), mask_token_id, ...)
diff_block_ids[:, 0] = output_ids[:, start_idx]

也就是说，block 的第一个位置是真实 token，后面全是 mask。这样 diffusion 路径的任务就变成：在已知起点的情况下，一次把后面一串位置都提议出来。

第三步：diffusion 路径并行给 proposal

记 diffusion 路径给出的 proposal 分布为 $q$。那么对 block 内位置 $i$，它会给出：

$$q_i(x_{n+i}\mid x_{\le n+1},m,\dots ,m)$$

然后并行采样得到候选 token：

$${\hat{x}}_{n+2:n+b}$$

这里可以把 $\hat{x}$ 理解为“草稿答案”。

第四步：AR 路径对 proposal 做严格校验

接下来代码会把：

$$[x_{n+1},{\hat{x}}_{n+2},\dots ,{\hat{x}}_{n+b}]$$

再喂回 AR 路径。

AR 路径给出的是真实因果分布 $p$：

$$p_i(x_{n+i}\mid x_{<n+i})$$

于是系统现在同时有两套信息：

• diffusion proposal 分布 $q$
• AR 真正分布 $p$

后面的核心问题就变成：哪些 proposal 可以接受，怎样接受才不改变最终分布。

6. “无损”到底是什么意思

这个词在很多项目里都很容易写虚。Orthrus 这里我觉得相对扎实，因为它在 generate() 里真的把分布修正写出来了。

6.1 温度接近 0 时：退化成前缀精确匹配

当 temperature < 1e-5 时，代码走的是确定性分支，直接比较：

$${\hat{x}}_{n+2}\stackrel{?}{=}{x}_{n+2}^{\text{AR}}$$

$${\hat{x}}_{n+3}\stackrel{?}{=}{x}_{n+3}^{\text{AR}}$$

一直到第一个不相等的位置为止。

如果连续匹配了 $k$ 个 token，那么这轮就接受长度为 $k$ 的前缀，然后把第 $k+1$ 个位置的 AR token 当成下一步继续推进。

所以确定性场景里，平均步数大约可以写成：

$$N_{\text{steps}}\approx \frac{T}{k+1}$$

这里 $k$ 越大，提速越明显。

6.2 有温度时：关键不是“比不比得上”，而是“怎么校正回原分布”

真正有意思的是随机采样分支。

对 proposal token ${\hat{x}}_i$，代码比较的是：

$$q_i({\hat{x}}_i)\text{和}{p}_i({\hat{x}}_i)$$

接受概率写成：

$${\alpha }_i=\min \left(1,\frac{p_i({\hat{x}}_i)}{q_i({\hat{x}}_i)}\right)$$

这其实就是一个典型的 acceptance correction。它的含义很直白：

• 如果 proposal 在这个 token 上低估了真实概率，那么大概率接受
• 如果 proposal 在这个 token 上高估了真实概率，就不能无脑接受，得按比例降下来

一旦在某个位置拒绝，代码不会简单退回 AR 贪心，而是从残差分布里采样：

$$r_i(x)=\max (p_i(x)-q_i(x),0)$$

再归一化：

$${\tilde{r}}_i(x)=\frac{r_i(x)}{{\sum }_{x^{\prime }}{r}_i(x^{\prime })}$$

然后：

$$x_i\sim {\tilde{r}}_i$$

这一步非常关键。因为如果拒绝以后直接改成从 $p_i$ 重采样，那么前面已经发生过的“proposal 尝试”这件事就没有被分布上抵消干净。残差分布的作用，就是把这部分偏差扣回去。

所以我会把 Orthrus 的无损理解成下面这句话：

它不是要求 diffusion proposal 必须等于 AR 输出，而是要求“接受 + 拒绝后的修正”这个整体过程，最终仍然落在原始 AR 分布上。

7. 为什么它能并行猜，但又不至于信息泄漏

并行生成里还有一个经常被忽略的问题：你让一个位置并行预测时，它到底能看到哪些 token？

如果看得太少，proposal 不准；看得太多，又会破坏因果约束。

Orthrus 在训练侧用 generate_dual_pass_mask() 处理这个问题，见 src/model.py。

它定义了两类可见性：

对历史 AR 前缀

每个 diffusion query 只能看各自 causal_limit 允许范围内的前缀 token。也就是：

visible AR ( i ) = { j ∣ j ≤ c i } \text{visible}_{\text{AR}}(i) = \{j \mid j \le c_i\} visibleAR​(i)={j∣j≤ci​}

其中 $c_i$ 就是 causal_limit。

对当前 draft block

每个 query 只能看同一个 block 内的 draft token。代码里的判断是：

$$
\left\lfloor \frac{q_idx}{b} \right\rfloor

=

\left\lfloor \frac{kv_idx - ar_len}{b} \right\rfloor
$$

这里 $b$ 是 block_size。

翻译成人话就是：可以并行，但只能在自己的 block 里并行。

于是总的可见集合可以写成：

$$\text{visible}(i)={\text{visible}}_{\text{AR}}(i)\cup {\text{visible}}_{\text{block}}(i)$$

这个 mask 设计我觉得很值得学，因为它说明了一件很实际的事情：并行生成不是简单把 causal mask 去掉，而是要重新设计一套“哪些历史能看、哪些未来能并行看”的规则。

8. generate() 里还有一个很容易漏掉的小细节：cache 裁剪

在 AR 验证时，代码会把整段 proposal block 都跑一遍，所以 KV cache 会临时长到：

$$n+b$$

但这一轮真正接受的可能只有前缀 $k+1$ 个 token。于是有效长度其实应该回到：

$$n+k+1$$

这就是为什么代码里在更新 start_idx 后，还要做：

past_key_values.crop(start_idx)

如果没有这一步，后面轮次看到的 cache 就会混入“验证过但没有真的接受”的 token，整个状态会立刻变脏。

这不是一个实现小尾巴，而是并行验证方案成立的必要条件之一。很多时候一篇论文讲清了算法，但真正决定代码是不是能跑对的，往往就是这种状态回滚细节。

9. 这个仓库为什么值得细读

我自己读完最大的感受是：Orthrus 值得看的，不只是结果图，而是它把几个平时很容易散落在不同系统里的东西，塞进了一条很短的逻辑链里：

1. 用同一个模型提供 AR view 和 diffusion view
2. 让 diffusion view 直接复用 AR 的 KV cache
3. 用 AR 路径做严格验证
4. 用 acceptance correction + residual resampling 把最终分布校正回去

如果把这几件事连起来看，Orthrus 的核心就很清楚了。它不是想证明“diffusion 比 autoregressive 更好”，而是在尝试回答一个更工程的问题：

能不能不引入第二个模型，只在一个模型内部做并行提案，然后还把最终分布守住？

从这个仓库给出的实现看，它至少提供了一个相当漂亮的答案。

10. 最后一句总结

如果只让我留一句结论，我会这么写：

Orthrus 真正值得学的，不是“并行生成”这四个字，而是它把并行提案、共享 KV cache 和分布校正放进了同一个生成循环里。

对做 LLM 推理的人来说，这种设计比单纯的 benchmark 数字更有参考价值。因为速度曲线可以随着硬件和实现变化，但“怎么在系统上少引入一个模型、少维护一份缓存、同时不把分布搞坏”这件事，才是更难、也更通用的工程问题。