双流Transformer：将主预填充路径与额外解码计算解耦

# 双流 Transformer：将主预填充路径与附加解码计算解耦

来源：https://arxiv.org/html/2608.12385

Liming Liu  
佐治亚理工学院（Georgia Institute of Technology）

&  
Mingze Wang  
北京大学（Peking University）

&  
Tuo Zhao  
佐治亚理工学院（Georgia Institute of Technology）

###### 摘要

随着大语言模型服务越来越多的请求，累积推理成本相对于一次性训练成本不断增长。在常见的服务模式下，两个推理阶段对硬件提出了不同的需求：提示词预填充（prompt prefill）以并行方式运行，通常受计算能力限制；而自回归解码是串行的，通常受内存带宽限制。传统的宽度或深度扩展会同时提高这两类成本，因为每一层新增的模型结构都会在两个阶段中都被计算。我们转而提出一个问题：是否可以将额外的学得计算分配给续写预测（continuation prediction），同时保留覆盖整个提示词的主计算和单一持久化键值（KV）缓存。我们通过 **双流 Transformer（Dual-Flow Transformer）** 实现了这一分离。其主流（primary flow）是一个完整的因果语言模型，它单独处理提示词并写入 KV 缓存；因此，辅助流（auxiliary flow）可以在提示词阶段被省略，只从其最后位置开始激活，在续写预测中添加额外计算，而无需写入持久化状态，也不会影响主流。两个流共享所有主要的注意力、MLP 和输出矩阵，并使用独立的 token 嵌入和轻量级耦合；由于它们共享权重和主缓存，分组执行创造了一个机会，可以在不同流之间复用已加载的权重以及缓存的键和值。在匹配 token 数量的比较中，双流模型在多种架构和数据配置下均取得了更低的验证损失。在 MoE 模型中，这种结构上的分离使得主专家和辅助专家的扇出（fan-out）成为独立可调的旋钮，用于权衡提示词成本、解码成本和预测质量。我们研究了两种分配机制：固定预填充专家计算而增加解码计算，以及固定解码专家计算而在两个流之间重新分配专家预算。这些实验共同揭示了预填充–解码–质量之间的权衡关系，并为按阶段分配专家提供了结构性机会。

## 1 引言

大语言模型服务的成本由两个截然不同的阶段决定。在**预填充（prefill）**阶段，所有提示词 token 并行处理。模型权重为批处理序列加载，提示词上的注意力通过高度并行的矩阵运算进行评估。随着上下文长度和批大小增大，预填充主要受算术吞吐量的限制。相比之下，在自回归**解码（decoding）**阶段，token 必须逐个串行生成。每生成一个新 token，都需要再次遍历模型权重，并再次读取不断增长的键值（KV）缓存。由于每个字节的移动只对应少量计算，解码通常受内存带宽而非峰值算术吞吐量的限制。这些不同的瓶颈对长上下文推理和交互式智能体尤其重要，因为长提示词之后往往是多次串行的模型调用以及新追加的观察内容。这种不匹配促使我们设计一种架构，使其能够改变用于续写预测的计算量，而不会成比例地改变提示词全局计算或持久化状态。

为什么需要将额外资源分配给解码？先前的研究表明，语言模型能力可以通过更多测试时计算来提升。思维链（chain-of-thought）、自洽性（self-consistency）和重复采样会在中间轨迹或替代轨迹上花费额外计算（Wei et al., 2022; Wang et al., 2023; Brown et al., 2024），而自适应方法则根据提示词难度来分配这些计算（Snell et al., 2025）。暂停 token 方法提供了更直接的证据，表明预测前的额外隐藏计算可以改善性能（Goyal et al., 2024; Kim et al., 2025）。这些研究共同支持将生成时计算视为一种能力扩展资源。预填充与解码之间的系统不对称性引出了一个具体的架构问题。现有的测试时方法通常通过生成更长的推理轨迹、插入延迟 token 或对多个候选进行采样，来在外部花费额外预算。我们转而问：单个模型能否同时满足四个性质：(i) 一条完整的主因果路径，它单独处理提示词；(ii) 没有持久化的辅助状态，因此 KV 缓存保持不变；(iii) 在续写预测步骤上应用额外的学得计算；(iv) 主计算与额外计算之间大量共享层权重和主缓存状态。这些性质共同定义了一个**相位解耦（phase-decoupled）**的架构目标：使预填充和解码计算可独立配置，同时保留固定的主提示词路径和单一持久化状态。

我们用双流 Transformer 来实现这一目标。主流是一条完整的、标准的因果 Transformer 路径，不依赖辅助状态；辅助流读取主流状态和缓存，但从不写入持久化状态，也从不反馈到主流。两个流共享所有主要的注意力、MLP 和输出矩阵，但使用独立的 token 嵌入和轻量级的学得耦合向量。由于主路径是自包含的，辅助流可以在较早的提示词位置被省略，只在续写预测步骤中被计算：对于第一个续写 token，在提示词最后一个位置计算；之后在每个生成位置计算。因此，提示词全局的主计算和持久化 KV 缓存与传统模型一致，仅在提示词最后一个位置多了一次辅助流计算。共享权重和主缓存还带来了内存流量上的机会：在分组执行中，一次加载的权重块和 KV 区域可以同时服务于两个流，从而在不扩大所读取权重和缓存条目集合的情况下提高解码算术量。两个流的下一个 token 分布上的混合似然为辅助轨迹提供了直接的训练信号。对于稀疏混合专家（MoE）模型，**路由器重放（router replay）**让辅助流复用主流所选专家，但使用自己的混合权重：这保持了所引用专家权重集合不变，不过所选专家会被应用于两个流的状态。

双流模型对预填充和解码计算的分离在 MoE 模型中尤其有用，因为主辅助两个流可以激活不同数量的专家。主扇出决定提示词全局的专家算术量，而两个扇出之和决定续写时的专家算术量。因此，扇出对定义了一个包含预填充成本、解码成本和预测质量的三方权衡。服务系统可以根据自身工作负载在这个分配空间中选取一个点：它可以保持提示词计算不变，将可用的解码计算用于更多辅助专家；也可以保持续写时专家算术量不变，同时为预填充密集型工作负载减少主扇出。这在不需要改变共享参数池或单一主缓存结构的前提下，提供了部署特定的灵活性。

我们的实验通过受控的 NanoGPT 数据扩展、稠密 LLaMA 风格模型规模扩展以及稀疏 MoE 基础配置来评估双流模型。在匹配 token 预算下，双流模型持续降低验证损失，而组件消融实验则分离了跨流耦合和混合读出的贡献。我们还加入了接近的训练计算量比较。对于 MoE，我们在两个互补机制中研究相位特定的专家分配：在保持预填充计算不变的情况下增加解码计算，以及在保持解码计算不变的情况下在两个阶段之间重新分配计算。这些实验揭示了由此产生的预填充–解码–质量权衡。我们的贡献如下：

- **架构**。我们提出 Dual-Flow，一种非对称共享 KV 设计，在结构上将提示词全局的主计算与额外的续写预测计算解耦。主流单独处理提示词并决定持久的 KV 缓存，而辅助残差轨迹只读取该缓存而从不写入，也从不影响主流。该构造结合了非对称共享 KV 注意力、共享稠密权重、独立嵌入、轻量级跨流耦合和混合目标；路由器重放将其扩展到 MoE，同时保持主流所选专家集在流之间一致。
- **经验质量**。在受控数据扩展、稠密模型规模扩展和稀疏基础配置下，Dual-Flow 在匹配 token 预算下一致降低验证损失。组件消融实验分离了跨流耦合和混合读出的贡献，并补充了训练计算量的比较，进一步支持了在交互流之间分配计算的好处。
- **相位特定的 MoE 分配**。与先前让每个流使用相同参数化和计算分配的并行流设计不同，我们使主辅助专家扇出可以独立配置。我们研究了互补的机制，分别保持预填充或解码计算固定，揭示了预填充成本、解码成本和预测质量之间的权衡，并支持按工作负载分配专家。

## 2 相关工作

##### 预测前的额外计算。

模型和数据扩展会增加训练和服务成本（Kaplan et al., 2020; Hoffmann et al., 2022），而测试时方法表明，计算可以转而花在对预测做出承诺之前。思维链在生成的 token 中暴露中间推理过程（Wei et al., 2022）；自洽性和重复采样评估多个候选轨迹（Wang et al., 2023; Brown et al., 2024）；自适应方法根据提示词难度分配测试时计算（Snell et al., 2025）；暂停 token 方法插入学得的虚拟位置，使模型在产生答案前获得额外隐藏计算（Goyal et al., 2024; Kim et al., 2025）。这些结果共同支持了一个基本前提：下一个 token 计算是一种有用的扩展资源。这些方法通常通过额外的生成位置、样本或串行潜在步骤来花费该资源，而不是改变普通自回归解码步骤的内部计算。

##### 无相位解耦的并行潜在计算。

并行流方法将这种额外计算移入模型内部。并行扩展（Parallel Scaling）对输入应用多种变换，用共享主干处理由此产生的流，并聚合其输出（Chen et al., 2025）。隐藏解码（Hidden Decoding）使用独立嵌入的流，并保留流特定的 KV 作为上下文（Liu et al., 2026）。状态预测分离（State-Prediction Separation）将输入 token 和预测 token 交错放置，并在局部窗口内保留预测 token 的 KV（Monea et al., 2026）。这些方法表明并行潜在轨迹可以改善语言建模，但它们的额外流也会参与提示词处理或保留流特定的历史。因此，提示词计算、持久化状态或两者都会随扩展计算而增长。它们的核心目标是通用的模型内部计算扩展，而不是在增加续写侧计算的同时保持提示词全局路径固定。

##### 相位解耦的潜在计算。

PHD-Transformer 和并行循环 Transformer（PLT）是分离提示词全局持久化状态与额外解码时计算的最接近先例。PHD 将 token 重复为原始副本和隐藏解码副本，只保留原始 token 的 KV 用于长距离注意力，并在预填充期间只计算原始路径（Wu et al., 2025b）。其训练序列在副本上使用结构化注意力布局，内核设计通过重排使稀疏模式对设备友好。PLT 将循环状态在不同连续 token 间错开，使解码时不同逻辑循环一起执行；它共享第一循环的全局 KV，并用门控滑动窗口状态增强后续循环（Wu et al., 2025a）。PLT 在 token 维度上是并行的，但在训练时循环维度上是串行的，因为每个循环都会消费前一个循环产生的移位状态。Dual-Flow 则保持两个同位残差轨迹，可以堆叠起来并用标准因果 FlashAttention 并行评估（Dao et al., 2022）。主轨迹从不读取辅助轨迹，而辅助轨迹通过学得的逐层耦合读取同层主中间状态。PHD 和 PLT 从最终副本或循环中读取预测；Dual-Flow 训练一个混合，其中两个分布都直接贡献于预测的下一个 token。因此，三者都将持久化提示词状态与额外解码计算解耦，但在训练图、跨轨迹交互和预测机制上实现方式不同。PHD 在架构形式上与 Dual-Flow 最接近；因此我们在 NanoGPT 组件消融中包含了一个受控的 PHD-2 实现。

##### 其他预测和状态共享机制。

多 token 预测为未来 token 附加预测头（Gloeckle et al., 2024），推测解码使用草稿过程降低串行验证成本（Leviathan et al., 2023）。Dual-Flow 通过交互残差轨迹预测同一个下一 token，因此既不是一组未来 token 头，也不是草稿-验证方法。多查询和分组查询注意力在查询头之间共享 KV 状态（Shazeer, 2019; Ainslie et al., 2023）；Dual-Flow 将类似的共享原则应用于完整残差状态流之间。路由器重放进一步使所选专家集在流之间对齐（Shazeer et al., 2017; Fedus et al., 2022），在保留所引用专家权重的同时，仍然将那些专家应用于两个隐藏状态。

## 3 背景

### 3.1 自回归语言建模

给定序列 $x_{1:T}$，自回归（AR）语言模型将其概率分解为 $p(x_{1:T}) = \prod_{t=1}^{T} p(x_t \mid x_{<t})$，其中每个条件分布通过一个以 $h_t$ 为输入的 softmax 头部进行参数化，$h_t$ 是第 $t$ 个位置最终 Transformer 层的隐藏状态。在带 KV 缓存的推理中，预填充阶段并行计算提示词的隐藏状态；解码阶段对每个新位置评估同一组层，读取之前位置的缓存键和值。

### 3.2 Transformer 块

一个标准 Transformer 块接收输入状态 $x$ 并计算 $x + \operatorname{Attn}(x) + \operatorname{MLP}(x + \operatorname{Attn}(x))$，其中注意力使用当前输入作为查询，并与缓存的键和值交互。自注意力的核心是 softmax 注意力 $o_i = \sum_j \operatorname{softmax}_j(q_i^\top k_j / \sqrt{d}) v_j$，其中 $q_i$、$k_j$、$v_j$ 分别通过对输入状态应用投影矩阵 $W_Q$、$W_K$、$W_V$ 得到。MLP 对每个位置独立应用：$\operatorname{MLP}(x) = f(x W_{\text{up}}) W_{\text{down}}$，其中 $f$ 是一个非线性激活函数，通常为 SiLU（或 GELU）。现代实现通常使用门控线性单元变体。

### 3.3 稀疏混合专家

MoE 层只激活每个 token 的一小组参数。对于每个输入状态 $x$，路由器计算 $r = \operatorname{softmax}(x W_r)$，其中 $W_r \in \mathbb{R}^{d \times E}$；top-$k$ 选择一组索引 $\mathcal{E}$，输出为所选专家输出的 $r$ 加权和。MoE 层将专家 MLP 的参数量与每 token 计算量分离开来。

## 4 双流 Transformer

Dual-Flow 使用两个残差轨迹：主流写持久化 KV，而辅助流增加额外的续写计算，但不产生持久化状态。我们在这两个轨迹之间共享层权重，使用独立的 token 嵌入，并引入逐层耦合。两个流共同预测下一个 token。

### 4.1 共享权重下的双注意力

对于位置 $t$ 和层 $\ell$，设 $x_\ell$ 为主流输入状态，$y_\ell$ 为辅助流输入状态。在每一层，主流计算：

$$
x_{\ell+1} = x_\ell + \operatorname{Attn}^{sg}(x_\ell, x_\ell, \mathrm{KV}_{x}) + \operatorname{MLP}(x_\ell + \operatorname{Attn}^{sg}(x_\ell, x_\ell, \mathrm{KV}_{x})),
$$

其中 $\operatorname{Attn}^{sg}$ 表示标准自注意力，查询来自第一个参数，键和值来自第二个参数，并读取持久化 KV 缓存 $\mathrm{KV}_{x}$。该公式等价于没有辅助状态的标准因果 Transformer：主流是一个完整模型，单独处理提示词。

辅助流计算：

$$
y_{\ell+1} = y_\ell + \operatorname{Attn}^{sg}(y_\ell, x_\ell, \mathrm{KV}_{x}) + \operatorname{MLP}(y_\ell + \operatorname{Attn}^{sg}(y_\ell, x_\ell, \mathrm{KV}_{x})),
$$

其中辅助状态 $y_\ell$ 使用与主流相同的头部和键/值投影矩阵。查询来自辅助流状态，而键和值来自主流的缓存状态；我们称这种结构为非对称共享 KV 注意力。辅助流只从主流读取，因此它可以被省略，直到需要额外计算时才激活。

在 MoE 层中，我们实现**路由器重放**：路由器权重是共享的，主流路由决定所选专家集合 $\mathcal{E}_x$；辅助流通过独立的（或共享的，取决于配置）路由器权重计算自己的混合权重，但使用相同的专家集合 $\mathcal{E}_x$。这保持了引用的专家权重集合与纯主流路径完全相同，同时仍然对两个流的状态应用这些专家。

### 4.2 耦合与预测

除了注意力耦合外，我们还引入逐层耦合向量。设 $u_\ell$ 和 $v_\ell$ 是第 $\ell$ 层可学习的维度为 $d$ 的向量。在注意力之前，我们计算：

$$
\tilde{x}_\ell = x_\ell + u_\ell \odot y_\ell, \qquad
\tilde{y}_\ell = y_\ell + v_\ell \odot x_\ell,
$$

其中 $\odot$ 表示逐元素乘法。主流的 $\tilde{x}_\ell$ 与 $x_\ell$ 之间的差异随 $y_\ell$ 变化，因此严格来说主流依赖于辅助流。为了保持主路径的完全自包含性，我们在主流中去除这种耦合。设 $\operatorname{sg}(\cdot)$ 表示停止梯度。主输入为：

$$
\tilde{x}_\ell^{\text{main}} = x_\ell + \operatorname{sg}(u_\ell \odot y_\ell),
$$

而辅助输入为：

$$
\tilde{y}_\ell^{\text{aux}} = y_\ell + v_\ell \odot \operatorname{sg}(x_\ell).
$$

在正向传播中，主流输入包含一个依赖于 $y_\ell$ 的项，因此辅助流确实参与主路径的前向计算；但反向传播中，$\partial \tilde{x}_\ell^{\text{main}} / \partial y_\ell = 0$，因此辅助参数不会从主流梯度接收信号。这保留了梯度意义上的自包含主路径。

最终预测是两个流分布的混合。令 $P_x$ 和 $P_y$ 为两个流在最终层隐藏状态上得到的 softmax 分布，$\lambda \in [0, 1]$ 为温度参数或可学习参数。预测分布为：

$$
P = \lambda P_x + (1 - \lambda) P_y.
$$

在训练中，该混合分布直接用于交叉熵损失。DiTTo 风格可学习的逐层耦合仅应用于辅助输入，因此主流前向路径在结构上不受辅助状态的影响。

### 4.3 预填充与解码的行为

设 $L$ 为层数，提示词长度为 $P$，新 token 数量为 $N$。一个标准 Transformer 在预填充时计算 $L$ 层乘以 $P$ 个 token，解码时计算 $L$ 层乘以 $N$ 个 token。Dual-Flow 的主路径在预填充时执行完全相同的 $LP$ 次层计算，并写入相同的 KV 缓存。辅助流在提示词内部位置被跳过，只在提示词的最终位置被求值；在解码时，每个生成位置都求值两个流。

因此，Dual-Flow 的预填充成本为 $LP$ 次主计算加上一次辅助计算，解码成本为 $LP$ 次主计算加上 $LN$ 次辅助计算。在典型解码设置中，$N$ 远大于 1，因此额外成本主要体现在解码侧。这种效果正是我们想要的：额外的学得计算集中于续写预测步骤，持久化 KV 缓存仅限于主流。

### 4.4 与之前并行流模型的比较

我们可以与 Parallel Scaling、Hidden Decoding 和 PHD 进行对比。Parallel Scaling 使用共享主干，但为每个流分配不同的输入变换。Hidden Decoding 使用独立嵌入，并保留流特定的 KV 上下文。在 PHD 中，原始的隐藏解码副本在预填充时也会被计算，尽管只保留原始 KV。这些方法在提示词和持久化状态上都会产生额外成本。Dual-Flow 则保持提示词路径不变，只通过辅助流增加续写时计算，并保留单一主 KV 缓存。因此，Dual-Flow 是唯一一种在不增加提示词计算和持久化状态的情况下增加续写计算的方法。

## 5 实验

### 5.1 设置

**数据。** 我们在 OpenWebText（OWT）上训练 NanoGPT 模型，该数据集是 OpenAI 的 WebText 数据集的开源重建版本。它包含约 90 亿个 token，从 Reddit 上获得至少 3 个赞的网页外链中提取。对于密集 LLaMA 风格实验，我们使用相同的 OWT 数据集以保持一致性。

**架构。** 我们的 NanoGPT 基线是一个 12 层、12 头、嵌入维度 768 的模型，总参数量约为 1.24 亿。Dual-Flow 变体共享相同的层数和头数，但每层增加一对耦合向量，以及两层额外的嵌入表：主嵌入表（与基线共享）和辅助嵌入表。总参数量增加约 1.5%。

**训练。** 我们使用 AdamW 优化器，权重衰减为 0.1，峰值学习率为 $6 \times 10^{-4}$，并采用余弦退火。批大小为 480 个序列，每个序列 1024 个 token，总 token 数为 60 亿。所有实验均在 8 块 A100 GPU 上运行。对于 LLaMA 风格实验，我们使用相同的优化器设置，但在最多 120 亿 token 上训练。

**评估。** 我们在 OWT 的固定验证集上报告验证困惑度（perplexity）。对于 MoE 实验，我们报告验证损失以及预填充和解码的专家算术量的消融。

### 5.2 组件消融

我们首先在 NanoGPT 规模下对 Dual-Flow 的主要设计选择进行消融：

1. **无耦合（No Coupling）**：移除逐层耦合向量，使辅助流仅通过共享权重与主流交互。
2. **无停止梯度（No Stop-Gradient）**：移除主路径上的停止梯度，使辅助参数也能从主流梯度中学习。
3. **无混合（No Mixture）**：直接从辅助流分布 $P_y$ 预测，而不是使用混合 $P$。

表 1 显示，完整模型、无耦合和无混合在验证困惑度上均优于基线，但完整模型是最好的。移除耦合会将困惑度提高约 0.4，移除停止梯度则再提高约 0.2。这证实了逐层耦合和停止梯度正则化的价值，以及混合读出相对于单流预测的优势。

### 5.3 数据扩展

我们评估在相同 token 预算下，Dual-Flow 能否持续优于基线。图 1（左）显示了数据扩展曲线：在 22 亿到 60 亿 token 之间，Dual-Flow 始终将验证损失降低约 0.03–0.05。这表明增益不会随训练数据增加而消失，而是在不同的数据规模下都保持稳定。

### 5.4 模型规模扩展

接下来我们测试增益是否可以推广到更大的稠密模型。我们训练了 LLaMA 风格的 3.5 亿和 7 亿参数模型，分别有 24 层和 36 层。表 2 显示，Dual-Flow 在两种规模下都降低了验证困惑度，7 亿参数模型的绝对改善更大。这表明双流结构在更大模型上也具有扩展性，增益不仅限于 NanoGPT 规模。

### 5.5 稀疏 MoE 中的相位特定专家分配

我们研究 MoE 模型中额外的辅助专家是否会影响质量。我们在 NanoGPT 规模上训练一个 MoE 基线，其中 12 层中的 8 层是 MoE 层，每个 MoE 层有 8 个专家，top-2 路由。这导致预填充和解码的专家算术量为每 token 16 个专家激活。我们考虑一个 Dual-Flow MoE 变体，其中辅助流使用 top-1 路由，从而总激活专家数为 3 个：2 个主专家加 1 个辅助专家。表 3 显示，给定相同的 token 预算，这种额外的辅助专家将验证损失降低了约 0.02，而预填充专家算术量保持不变（仍然是 2），解码专家算术量则从 2 增加到 3。这表明，在不增加提示词计算的情况下，可以将额外的解码计算转化为质量提升。

我们接着研究相反的机制：固定解码计算，同时重新分配预填充和解码之间的专家预算。我们保持总专家激活数为 2 个：主辅助流分别是 top-1 和 top-1，或者 top-2 和 top-0（基线）。如表 4 所示，在固定解码成本下，从主流中移除一个专家并为辅助流添加一个专家会略微降低验证损失，同时将预填充成本减半。这是一个有利的权衡：在预填充密集型工作负载中，可以牺牲少量质量来大幅降低提示词成本。

### 5.6 训练计算量的比较

最后，我们进行一个控制训练计算量的比较。我们使用相同总 FLOPs 训练一个更大的基线，并与使用相同 FLOPs 的 Dual-Flow 进行比较。表 5 显示，在相同训练计算量下，Dual-Flow 的验证损失仍然更低。这表明额外容量本身并不是主要因素，双流结构才是收益的来源。

## 6 结论

我们提出了 Dual-Flow Transformer，这是一种将持久化预填充状态与额外解码计算解耦的架构。其主流是标准因果 Transformer，单独处理提示词并写入单一 KV 缓存；辅助流从主缓存读取，但不写持久化状态，也不向主流提供梯度。共享权重、路由器重放和混合目标使该架构高效且易于训练。在数据扩展、模型规模扩展和稀疏 MoE 配置下，Dual-Flow 在匹配 token 预算下始终降低验证损失。对 MoE 的消融揭示了预填充成本、解码成本和预测质量之间的三元权衡。未来的工作可以将双流结构扩展到其他模态，并进一步研究跨流分配计算的自适应策略。