FFN Fusion: Rethinking Sequential Computation in Large Language Models

TL;DR

This work tackles the latency bottleneck in large language models by introducing FFN Fusion, an architectural optimization that parallelizes sequences of FFN layers after selectively removing attention via the Puzzle framework. The authors develop a principled methodology to identify independent FFN sequences and fuse them into wider, parallelizable units, reducing synchronization overhead and improving hardware utilization. They demonstrate the approach on Llama-3.1-405B-Instruct to produce Ultra-253B-Base, achieving a 1.71× latency speedup and a 35× reduction in per-token cost while preserving or improving benchmark performance, with results strengthening as model size increases. The work also explores the limits of full-block parallelization and provides empirical evidence that FFN Fusion complements quantization and pruning, suggesting promising future directions in neural architecture design and efficient inference for very large models.

Abstract

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.

Figures (12)

• Figure 1: An overview of our FFN Fusion approach. Step $1$: We apply Puzzle to partially remove FFN layers and remove entire attention layers. Step $2$: We fuse consecutive FFN layers into a single wide FFN layer.
• Figure 2: Block-wise dependency heatmap for Ultra-PreFusion (log-scale). Each coordinate $(i, j)$ encodes how much block $j$ depends on block $i$, measured by the cosine distance between $h^j({\bm{X}})$ and $\tilde{h}^j_i({\bm{X}})$. Darker blue hues indicate weaker dependencies. The attention-removed region (dotted box) shows consistently lower values, suggesting greater potential for parallelization. Darker red hues indicate stronger dependencies. Further analysis of this Figure can be found in Appendix \ref{['app:pair_block_analysis']}.
• Figure 3: Comparison of Ultra-253B-Base before and after applying an additional longer continual pretraining.
• Figure 4: Accuracy vs. latency performance of Ultra-253B-Base. Latency is measured on a single NVIDIA H100 node with tensor parallel (TP) 8, running in FP8. The red line represents the efficient frontier, highlighting models with the best accuracy-to-throughput tradeoff. Accuracy = $(\text{MT-Bench} \times 10 + \text{MMLU} + \text{MMLU-Pro} + \text{Arena Hard} + \text{HumanEval}) / 5$.
• Figure 5: Accuracy vs. Latency for FFN Removal vs. Fusion.

Theorems & Definitions (2)

• Theorem 3.1
• proof