# Decode-Branch Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation

arXiv: https://arxiv.org/abs/2608.12385 · HTML: https://mlsanity.com/paper/2608.12385/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 19 of its 25 references that are in the map.
- First posted on arXiv 31 Jul 2026.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [Think before you speak: Training Language Models With Pause Tokens](https://mlsanity.com/paper/2310.02226/) (2023)
- [Hyper-Connections](https://mlsanity.com/paper/2409.19606/) (2024)
- [Parallel Scaling Law for Language Models](https://mlsanity.com/paper/2505.10475/) (2025)
- [Parallel Loop Transformer for Efficient Test-Time Computation Scaling](https://mlsanity.com/paper/2510.24824/) (2025)
- [Learning to Insert [PAUSE] Tokens for Better Reasoning](https://mlsanity.com/paper/2506.03616/) (2025)
- [Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models](https://mlsanity.com/paper/2607.08186/) (2026)

## Built on by (most cited first, 1 of 1)

- [Almost Free State Prediction Separation](https://mlsanity.com/paper/2609.03807/) (2026)

## Abstract (arXiv)

As large language models serve ever more requests, cumulative inference cost is growing relative to the one-time cost of training. In typical serving, prompt prefill runs in parallel and is compute-bound, whereas autoregressive decode is sequential and memory-traffic-bound. Conventional width or depth scaling raises both costs together, since every added layer is evaluated in both phases and enlarges the weights read at each decode step. We instead ask whether additional learned computation can be allocated to continuation prediction while preserving prompt-wide primary computation and a single KV cache. We realize this with the Decode-Branch Transformer. Its primary path alone processes the prompt and writes the KV cache; the decode branch is omitted during prefill and activated only from the final prompt position onward, adding continuation computation without writing state or affecting the primary path. The paths share attention, MLP, and output matrices, using separate token embeddings with lightweight coupling. Grouped decode reuses loaded weight tiles and the primary KV cache across both paths, so the added arithmetic does not proportionally increase dominant memory traffic or decode latency. Across matched-token comparisons, Decode-Branch achieves lower validation loss across architectures and data settings. In MoE models, the primary and branch expert fan-outs become independent knobs for trading prompt cost, decode cost, and predictive quality. We study two expert-allocation regimes, holding prefill or decode computation fixed, and expose a prefill-decode-quality trade-off enabled by phase-specific expert allocation.