Fast Forward: Accelerating LLM Prefill with Predictive FFN Sparsity

TL;DR

FastForward targets the prefill bottleneck in long-context LLM inference by injecting structured FFN sparsity into the block-wise processing of prompts. It combines a lightweight Expert Neuron Predictor, an Error Compensation Network, and a Dynamic Layerwise Sparsity Schedule to selectively prune FFN computations per block while preserving accuracy, enabling substantial TTFT reductions on edge- and constrained-hardware scenarios. Across LLaMA and Qwen models up to 8B, it achieves up to 1.45× compute-bound speedup with around a 4–6% accuracy penalty at 50% sparsity on LongBench, and demonstrates robust ablations highlighting the value of per-layer budgeting, keeping sink blocks dense, and learned expert selection. This approach shifts sparsity to the prefill stage, providing practical efficiency gains for short-to-moderate context lengths and offering a path toward more bandwidth-efficient inference on resource-limited devices.

Abstract

The prefill stage of large language model (LLM) inference is a key computational bottleneck for long-context workloads. At short-to-moderate context lengths (1K--16K tokens), Feed-Forward Networks (FFNs) dominate this cost, accounting for most of the total FLOPs. Existing FFN sparsification methods, designed for autoregressive decoding, fail to exploit the prefill stage's parallelism and often degrade accuracy. To address this, we introduce FastForward, a predictive sparsity framework that accelerates LLM prefill through block-wise, context-aware FFN sparsity. FastForward combines (1) a lightweight expert predictor to select high-importance neurons per block, (2) an error compensation network to correct sparsity-induced errors, and (3) a layer-wise sparsity scheduler to allocate compute based on token-mixing importance. Across LLaMA and Qwen models up to 8B parameters, FastForward delivers up to 1.45$\times$ compute-bound speedup at 50% FFN sparsity with $<$ 6% accuracy loss compared to the dense baseline on LongBench, substantially reducing Time-to-First-Token (TTFT) for efficient, long-context LLM inference on constrained hardware.

Figures (7)

• Figure 1: Time to First Token (TTFT) for the LLaMA-3.1-8B-Instruct model with and without 50% sparsity across different context lengths, measured on a single NVIDIA A100 GPU.
• Figure 2: Time taken by different components of a transformer block to process tokens of varying context length for Llama-3.1-8B profiled in a single A100 GPU
• Figure 3: Block diagram showing the working of FastForward
• Figure 4: Blockwise attention score distributions across selected layers of the LLaMA-3.2-3B-Instruct model. Each histogram shows the sum of attention scores received by individual blocks (excluding the first block) within a block of length 128 during the prefill phase. The variation in score ranges across layers indicates differing degrees of token mixing, particularly among non-sink tokens.
• Figure 5: Mean attention scores per block during prefill across layers of the LLaMA-3.2-3B-Instruct model. Each value represents the average attention received by blocks of size 128, excluding the first block. These means serve as a proxy for the extent to which each layer contributes to mixing token embeddings, with higher values indicating stronger integration of non-sink tokens.