FFN Fusion: Rethinking Sequential Computation in Large Language Models

Akhiad Bercovich (NVIDIA) · Mohammed Dabbah (Computer Science Departmen, Technion-Israel Institute of Technology) · Omri Puny (Weizmann Institute of Science) · Ido Galil (Technion) · Amnon Geifman (Weizmann Institute) · Yonatan Geifman (Deci) · Izik Golan (NVIDIA) · Ehud Karpas (NVIDIA) · Itay Levy (NVIDIA) · Zach Moshe (NVIDIA) · Najeeb Nabwani (NVIDIA) · Tomer Ronen (NVIDIA) · Itamar Schen (NVIDIA) · Ido Shahaf (NVIDIA) · Oren Tropp (NVIDIA) · Ran Zilberstein (NVIDIA) · Ran El-Yaniv (Technion & Deci.AI)
253b modelarchitectural optimizationattention layersfeed-forward networkffn fusioninference latencyllama-3.1-405b-instructmodel behaviorneural architecture designparallel operationsparallelizationprincipled methodologysequential computationtransformer blocks

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.