FFN Fusion: Rethinking Sequential Computation in Large Language Models

We introduce \textit{FFN Fusion}, an architectural optimization technique that reduces sequential computation in large language models by identifying and exploiting natural opportunities for parallelization. Our key insight is that sequences of Feed-Forward Network (FFN) layers, particularly those remaining after the removal of specific attention layers, can often be parallelized with minimal accuracy impact. We develop a principled methodology for identifying and fusing such sequences, transforming them into parallel operations that significantly reduce inference latency while preserving model behavior. Applying these techniques to Llama-3. 1-405B-Instruct, we create a 253B model (253B-Base), an efficient and soon-to-be publicly available model that achieves a 1. 71$\times$ speedup in inference latency and 35$\times$ lower per-token cost while maintaining strong performance across benchmarks. Most intriguingly, we find that even full transformer blocks containing both attention and FFN layers can sometimes be parallelized, suggesting new directions for neural architecture design.

Authors

• Akhiad Bercovich
• Mohammed Dabbah
• Omri Puny
• Ido Galil
• Amnon Geifman
• Yonatan Geifman
• Izik Golan
• Ehud Karpas
• Itay Levy
• Zach Moshe
• Najeeb Nabwani
• Tomer Ronen
• Itamar Schen
• Ido Shahaf
• Oren Tropp
• Ran Zilberstein
• Ran El-Yaniv

Keywords

No keywords are indexed for this paper.