Enhancing LLM Efficiency: Targeted Pruning for Prefill-Decode Disaggregation in Inference

2025-08-29

AI Summary

To address the high computational and communication overhead caused by separating prefill and decoding stages in large language model (LLM) inference, this paper proposes a modular pruning framework tailored for Prefill-Decoding (PD) separation architectures. The method comprises two key components: (1) stage-aware block-level pruning and knowledge distillation, dynamically adapting sparsity patterns to the distinct computational characteristics of prefill and decoding; and (2) token-position-aware KV cache pruning, which selectively retains only the initial and final tokens’ KV states in critical layers, drastically reducing inter-stage data transfer volume. Experiments demonstrate consistent accuracy retention under both unified and separated deployment configurations. The approach achieves a 20.56% end-to-end inference speedup and reduces inter-stage bandwidth consumption by 4.95×. This work establishes a novel paradigm for efficient LLM inference at the edge and in distributed settings.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Large Language Models (LLMs) demonstrate exceptional capabilities across various tasks, but their deployment is constrained by high computational and memory costs. Model pruning provides an effective means to alleviate these demands. However, existing methods often ignore the characteristics of prefill-decode (PD) disaggregation in practice. In this paper, we propose a novel pruning method for PD disaggregation inference, enabling more precise and efficient block and KV Cache pruning. Our approach constructs pruning and distillation sets to perform iterative block removal independently for the prefill and decode stages, obtaining better pruning solutions. Moreover, we introduce a token-aware cache pruning mechanism that retains all KV Cache in the prefill stage but selectively reuses entries for the first and last token sequences in selected layers during decode, reducing communication costs with minimal overhead. Extensive experiments demonstrate that our approach consistently achieves strong performance in both PD disaggregation and PD unified settings without disaggregation. Under the default settings, our method achieves a 20.56% inference speedup and a 4.95 times reduction in data transmission bandwidth consumption.

Problem

• Reducing computational and memory costs in LLM deployment
• Addressing prefill-decode disaggregation characteristics in pruning
• Optimizing KV Cache pruning to minimize communication overhead

Innovation

• Iterative block pruning for prefill and decode stages
• Token-aware KV cache pruning mechanism
• Selective KV cache reuse for first and last tokens