Source-linked AI summary
POP: Prefill-Only Pruning for Efficient Large Model Inference
Junhui He, Zhihui Fu, Jun Wang, Qingan Li
TL;DR
LLM and VLM inference is costly, and structured pruning often loses accuracy despite its hardware compatibility. POP makes pruning stage-aware by removing deep layers during prefill while preserving the full model for decode, using virtual-gate importance analysis and KV/cache boundary mechanisms. It achieves prefill speedups of up to 1.37× with accuracy comparable to the full model across evaluated models and modalities.
Problem
Existing structured pruning is hardware-compatible but often degrades accuracy because it applies pruning stage-agnostically despite different prefill and decode functions.
Method
POP uses virtual gates to identify layer importance, prunes deep layers during prefill, retains the full model during decode, and uses independent KV projections with boundary handling.
Results
POP achieves prefill speedup of up to 1.37× while maintaining accuracy comparable to the full model and demonstrating minimal performance loss across models and modalities.
Takeaways & Limitations
Stage-aware optimization can reduce prefill computation while preserving decode capacity, offering a hardware-compatible alternative to accuracy-costly structured pruning.
Takeaways & Limitations
POP still requires full model weights during decode, so it does not reduce peak VRAM usage and is best suited to compute-bound scenarios.
Abstract
from arXiv · showhide
Large Language Models (LLMs) and Vision-Language Models (VLMs) have demonstrated remarkable capabilities. However, their deployment is hindered by significant computational costs. Existing structured pruning methods, while hardware-efficient, often suffer from significant accuracy degradation. In this paper, we argue that this failure stems from a stage-agnostic pruning approach that overlooks the asymmetric roles between the prefill and decode stages. By introducing a virtual gate mechanism, our importance analysis reveals that deep layers are critical for next-token prediction (decode) but largely redundant for context encoding (prefill). Leveraging this insight, we propose Prefill-Only Pruning (POP), a stage-aware inference strategy that safely omits deep layers during the computationally intensive prefill stage while retaining the full model for the sensitive decode stage. To enable the transition between stages, we introduce independent Key-Value (KV) projections to maintain cache integrity, and a boundary handling strategy to ensure the accuracy of the first generated token. Extensive experiments on Llama-3.1, Qwen3-VL, and Gemma-3 across diverse modalities demonstrate that POP achieves up to 1.37$\times$ speedup in prefill latency with minimal performance loss, effectively overcoming the accuracy-efficiency trade-off limitations of existing structured pruning methods.
1 Introduction
LLM/VLM inference is computationally expensive, while existing pruning methods trade hardware compatibility against accuracy. POP addresses this by pruning deep layers only during prefill and retaining the full model during decode.
- LLM and VLM deployment is limited by substantial inference overhead from their massive parameter counts.
- Unstructured pruning can preserve accuracy but often needs specialized hardware, whereas structured pruning is hardware-compatible but can significantly degrade accuracy.
- Prefill encodes input history into the KV cache, while decode also models the next-token probability distribution, creating different pruning sensitivities.
- POP uses virtual gates to estimate layer importance and finds deep layers critical for generation but largely redundant for context encoding.
- POP prunes deep layers exclusively during prefill while retaining full model capacity during decode.
- POP uses mechanisms for missing KV states and stage-boundary handling to transition between pruned prefill and full decode stages.
2 Preliminary
Decoder-only Transformer inference maintains KV states across layers and tokens, while layer pruning bypasses selected computational blocks. Existing pruning applies the same skipped-layer set to prefill and decode despite their different functional goals.
- A decoder-only Transformer layer typically combines attention and feed-forward blocks with residual connections and layer normalization.
- During inference, attention computes queries, keys, and values for the current token, with positional information typically added through RoPE.
- New key-value pairs are appended to the cache so autoregressive generation avoids recomputing token history.
- Layer pruning removes entire attention and FFN blocks while preserving residual connections, reducing selected layers to identity mappings.
- Existing methods apply the skipped-layer set identically across prefill and decode, although prefill focuses on context encoding and decode also predicts the next token.
3 Method
POP estimates stage-specific layer importance and prunes deep layers only during prefill, retaining the full model for decode. Independent KV projections and boundary handling preserve cache availability and first-token accuracy while reducing prefill computation.
- Virtual-gate importance estimation: Layer importance is defined as the increase in calibration loss when that layer is pruned, with other parameters unchanged.
- Virtual-gate importance estimation: Virtual gates suppress a layer’s residual Attention and FFN updates when its scalar gate is set to zero, approximating physical layer removal.This avoids separately removing and evaluating every layer.
- Stage-aware importance analysis: Stage-aware gates separately estimate layer sensitivity during prefill and decode, whose distinct roles motivate different pruning strategies.Prefill encodes prompt information into the KV cache, whereas decode also predicts the next-token distribution.
- Stage-aware importance analysis: Experiments across Llama-3.1 text and Qwen3-VL multimodal data find asymmetric importance, with deep layers essential for decode but largely redundant for prefill.
- Prefill-only pruning: POP removes the deep layers during prefill while retaining the full model during decode; the default strategy prunes the last 1/3 of layers.The pruning ratio is empirically selected to balance efficiency and accuracy.
- Prefill-only pruning: Independent KV projections generate cache entries for skipped layers while bypassing their heavy Attention and FFN computation during prefill.Projection cost is reported as less than 5% for Llama-3, Qwen-3, and Gemma-3, preserving the speedup benefit while supplying complete decode-time KV caches.
- Boundary handling: The boundary treats the last input token as the first decode step so its prediction uses the full model capacity.Pruning deep layers for that token would degrade the first generated token and accumulate generation errors.
4 Experiments
Experiments across text-only and vision-language models show that POP preserves accuracy while accelerating prefill, especially for long contexts, with design ablations confirming its key components.
- Experimental Setup: Experiments cover Llama-3.1, Qwen3-VL, and Gemma-3 across text, multimodal, reasoning, long-context, and vision-understanding tasks.The evaluation uses multiple libraries, calibration datasets, and NVIDIA A100 GPUs.
- Accuracy on Downstream Tasks: POP preserves accuracy across benchmarks despite pruning approximately 33% of the model, outperforming structured baselines that often collapse on open-ended generation.On Llama-3.1, POP retains 97.00% of full-model GSM8K performance and 95.64% of HumanEval performance.
- Accuracy on Downstream Tasks: POP maintains near-lossless multimodal accuracy, including 50.67% vs 51.33% on MMMU and 86.40% vs 87.03% on ScreenSpot for Qwen3-VL.Its accuracy is comparable to Wanda while retaining structural hardware compatibility.
- Effectiveness of Design Choices: Removing deep-layer targeting, independent KV projections, or boundary handling causes substantial accuracy degradation in ablations on Qwen3-VL.Shallow pruning reaches 0.15% on GSM8K and 0% on HotpotQA, while removing independent KV projections causes catastrophic collapse.
- Sensitivity to Pruning Ratio: A 33% pruning ratio delivers 1.37× acceleration with negligible accuracy loss, whereas pruning beyond 50% sharply reduces performance, especially on HotpotQA.Lower ratios can preserve or slightly exceed full-model accuracy but provide limited speedup.
5 Related Work
Prior pruning methods trade accuracy against hardware efficiency, while token pruning reduces sequence length through perplexity- or attention-based selection.
- Model Pruning: Unstructured pruning removes individual weights using magnitude and activation information but often requires specialized kernels for wall-clock speedups.Examples include SparseGPT and Wanda.
- Model Pruning: Structured pruning removes coarse-grained components such as layers or channels to improve hardware compatibility, but can degrade accuracy.The passage contrasts parameter pruning with token pruning, which reduces sequence length instead.
- Token Pruning: Token pruning accelerates inference by selecting informative text or visual tokens using perplexity-based or attention-based criteria.Examples include LLMLingua, PyramidInfer, DAC, FastV, and DART.
6 Conclusion
POP exploits asymmetric pruning sensitivity between prefill and decode by pruning deep layers only during prefill and retaining the full model for decode. It achieves up to 1.37× prefill speedup with accuracy comparable to the full model.
- Conclusion: POP identifies deep layers as indispensable for decode but minimally useful for prefill, motivating stage-aware rather than stage-agnostic pruning.The method decouples context processing from token generation.
- Conclusion: POP achieves up to 1.37× prefill speedup while maintaining accuracy comparable to the full model and outperforming existing structured pruning methods.The conclusion presents stage-aware optimization as a promising direction beyond pruning, including quantization and architecture.
7 Limitations
POP improves efficiency and accuracy trade-offs but has limitations in memory usage and deployment architecture. It requires full-model weights during decode and currently relies on a monolithic inference pipeline.
- POP requires the full model weights during decode, so it does not reduce peak VRAM usage and suits compute-bound rather than capacity-bound scenarios.
- The current implementation uses a monolithic inference pipeline modified from the Transformers library.
- Adapting POP to disaggregated prefill-decode systems would involve non-trivial engineering efforts left for future research.
A Robustness to Representation Mismatch
Layer skipping creates representation drift, but POP's attention outputs remain stable across deep layers. The virtual-gate analysis and empirical measurements explain why this mismatch does not necessarily corrupt decoding.
- The virtual-gate importance score accounts for representation mismatch by assigning high importance to layers whose removal would cause large loss-gradient variance.
- Skipping layers causes hidden-state similarity to fall from 1.0 to 0.71, while KV-cache value similarity drops as low as 0.46.
- Attention outputs retain cosine similarity above 0.96 across layers 25-36 despite drifted keys and values.
- The attention mechanism stabilizes drifted representations through weighted aggregation over the context window, preserving semantic information for subsequent layers.
B Integration with Orthogonal Methods
POP is complementary to sparse attention and token pruning because it removes transformer depth rather than reducing sequence length or sparsifying attention computation. Their combination can improve the accuracy-efficiency frontier.
- POP reduces computation along the depth dimension by removing transformer layers, whereas token pruning reduces sequence length and sparse attention reduces attention-block computation.
- Table 5 reports long-context task accuracies and TTFT speedup ratios under greedy decoding, with TTFT measured at batch size 1 and input length 32K.
- Combining conservative FlexPrefill or LLMLingua-2 settings with POP achieves approximately 1.56× speedup while preserving higher HotpotQA accuracies.
- Aggressive FlexPrefill sparsity reaches 1.22× speedup but lowers HotpotQA accuracy from 57.83 to 47.79.