Dual-Flow Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation

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

Abstract: As large language models serve more requests, cumulative inference cost is becoming increasingly important relative to one-time training cost. The two inference phases stress hardware differently: prompt prefill is parallel and typically compute-bound, whereas autoregressive decode is sequential and often memory-bandwidth-bound. Conventional width or depth scaling increases both costs together because every added layer is evaluated in both phases. We ask whether additional learned computation can instead be allocated to continuation prediction while preserving the prompt-wide primary computation and a single persistent key-value (KV) cache.

摘要： 随着大型语言模型处理的请求日益增多，累积推理成本相对于一次性训练成本变得愈发重要。推理的两个阶段对硬件的压力各不相同：提示词预填充（Prompt Prefill）是并行的，通常受限于计算能力；而自回归解码（Autoregressive Decode）是串行的，通常受限于内存带宽。传统的宽度或深度扩展会同时增加这两个阶段的成本，因为每一层新增的层都需要在两个阶段中进行评估。我们探讨是否可以将额外的学习计算分配给续写预测，同时保留提示词处理的主要计算以及单一的持久化键值（KV）缓存。

We introduce the Dual-Flow Transformer. Its primary flow is a complete causal language model that processes the prompt and writes the KV cache. The auxiliary flow is omitted during prompt processing and activated only from the final prompt position onward, adding continuation-prediction computation without writing persistent state or influencing the primary flow. The two flows share major attention, MLP, and output matrices, while using separate token embeddings and lightweight coupling. Sharing weights and the primary cache also creates opportunities to reuse loaded weights and cached keys and values during grouped execution.

我们引入了“双流 Transformer”（Dual-Flow Transformer）。其主流（Primary flow）是一个完整的因果语言模型，负责处理提示词并写入 KV 缓存。辅助流（Auxiliary flow）在提示词处理阶段被忽略，仅从提示词的最后一个位置开始激活，从而在不写入持久化状态或影响主流的情况下，增加续写预测的计算量。这两个流共享主要的注意力机制、MLP 和输出矩阵，同时使用独立的 Token 嵌入和轻量级耦合。共享权重和主缓存也为在分组执行期间重用已加载的权重以及缓存的键值对创造了机会。

Across matched-token comparisons, Dual-Flow achieves lower validation loss across architectures and data configurations. In MoE models, the separation makes primary and auxiliary expert fan-outs independent controls over prompt cost, continuation cost, and predictive quality. We study two regimes: increasing decode computation at fixed prefill expert computation, and reallocating a fixed decode expert budget between the two flows. These experiments expose a prefill-decode-quality trade-off and demonstrate the potential of phase-specific expert allocation.

在匹配 Token 的对比测试中，双流架构在各种模型架构和数据配置下均实现了更低的验证损失。在混合专家模型（MoE）中，这种分离使得主专家和辅助专家的分发（Fan-out）能够独立控制提示词成本、续写成本和预测质量。我们研究了两种模式：在固定预填充专家计算量的情况下增加解码计算量，以及在两个流之间重新分配固定的解码专家预算。这些实验揭示了预填充-解码-质量之间的权衡，并展示了针对特定阶段进行专家分配的潜力。