FFN Fusion: Rethinking Sequential Computation in Large Language Models
Authors: name: Akhiad Bercovich, name: Mohammad Dabbah, name: Omri Puny, name: Ido Galil, name: Amnon Geifman, name: Yonatan Geifman, name: Izhak Golan, name: Ehud Karpas, name: Itay Levy, name: Zach Moshe, name: Najeeb Nabwani, name: Tomer Ronen, name: Itamar Schen, name: Elad Segal, name: Ido Shahaf, name: Oren Tropp, name: Ran Zilberstein, name: Ran El-Yaniv

We introduce 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 Llama-Nemotron-Ultra-253B-Base
(Ultra-253B-Base), an efficient and soon-to-be publicly available model that
achieves a 1.71X speedup in inference latency and 35X lower per-token cost
while maintaining strong performance across benchmarks. Through extensive
experiments on models from 49B to 253B parameters, we demonstrate that FFN
Fusion becomes increasingly effective at larger scales and can complement
existing optimization techniques like quantization and pruning. 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.