Dual-Flow Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation
06:00 · August 15, 2026 · arXiv cs.AI RSS
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. 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. 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.
Summary
The Dual-Flow Transformer addresses the growing share of inference cost in large language model deployment by separating the hardware demands of the two serving phases. Prompt prefill processes tokens in parallel and is typically compute-bound, while autoregressive decode proceeds sequentially and is often limited by memory bandwidth. Conventional scaling increases both costs together because every added parameter participates in both phases. The new architecture instead supplies extra learned computation only during continuation prediction, while leaving the prompt-wide primary path and a single persistent key-value cache unchanged.
Its primary flow is a standard causal transformer that alone processes the entire prompt and writes the KV cache. The auxiliary flow remains inactive across earlier prompt positions and activates only from the final prompt token onward. It reads the primary state and cache but never writes persistent entries or feeds back into the primary computation. The two flows share the main attention, MLP, and output matrices, employ separate token embeddings, and connect through lightweight learned coupling vectors. A mixture objective over their next-token distributions supplies direct training signal to the auxiliary path. In grouped execution the shared weights and cache allow a single load of each weight tile and KV region to serve both flows.
When applied to mixture-of-experts models, the separation turns the primary and auxiliary expert fan-outs into independent controls. The primary fan-out sets prompt cost, while the sum of both fan-outs determines continuation cost. Experiments compare matched token budgets across dense and sparse configurations and record lower validation loss for Dual-Flow. Additional runs that either increase decode experts at fixed prefill budget or reallocate a fixed decode budget between flows expose concrete prefill–decode–quality trade-offs, showing how serving systems can tune expert allocation to workload characteristics without altering the shared parameter set or cache structure.
Why it matters
High technical depth and novelty make it directly actionable for Dutch AI researchers optimizing LLM serving costs; phase-specific allocation aligns with EU deployment needs for efficient, scalable models in SME and research settings.
