TLDR: This research introduces a novel pruning method to enhance the efficiency of Large Language Models (LLMs), particularly in prefill-decode disaggregated inference. The approach involves an iterative, stage-aware block removal strategy and a token-aware KV Cache pruning mechanism that selectively reuses only critical parts of the cache during the decode stage. This leads to a 20.56% inference speedup and a 4.95x reduction in data transmission bandwidth, demonstrating a significant improvement in LLM deployment efficiency.
Large Language Models, or LLMs, have become incredibly powerful tools, excelling at a wide range of tasks. However, their immense size often makes them costly and challenging to deploy in real-world applications. To tackle this, researchers are constantly looking for ways to make these models more efficient without sacrificing their impressive capabilities.
One common technique to reduce the computational and memory demands of LLMs is called model pruning. This involves removing redundant or less important parts of the model. While effective, existing pruning methods often overlook a crucial aspect of how LLMs are used in practice: the prefill-decode (PD) disaggregation during inference.
PD disaggregation is a strategy where the LLM inference process is split into two distinct stages. The ‘prefill’ stage processes the initial input sequence, generating a comprehensive understanding and a memory of past computations called the KV Cache. The ‘decode’ stage then uses this information to generate new tokens one by one. By separating these stages, each can be optimized for its specific resource needs. However, this separation also introduces challenges for traditional pruning methods, particularly due to the distinct sensitivities of each stage to pruning and the significant data transfer required for the KV Cache between the prefill and decode nodes.
A new research paper, titled Enhancing LLM Efficiency: Targeted Pruning for Prefill-Decode Disaggregation in Inference, proposes a novel pruning method specifically designed to integrate deeply with PD disaggregation. Authored by Hao Zhang, Mengsi Lyu, Yulong Ao, and Yonghua Lin from the Beijing Academy of Artificial Intelligence, this approach aims for more precise and efficient pruning of both model blocks and the KV Cache.
Smart Block Removal for Each Stage
The paper introduces an iterative block removal strategy that is independently tailored for the prefill and decode stages. This is crucial because the prefill stage is often more sensitive to pruning, meaning removing parts here can have a bigger negative impact on performance. The method constructs ‘pruning and distillation sets’ to guide this iterative process, allowing it to find better combinations of blocks to remove or merge. By distinguishing between the stages, the system can identify which blocks are best removed in the prefill, decode, or both, leading to more optimized solutions compared to previous methods.
Token-Aware KV Cache Pruning
Beyond structural block pruning, the research also addresses the substantial bandwidth overhead caused by transferring the KV Cache between prefill and decode nodes. The proposed ‘token-aware cache pruning’ mechanism is quite clever. During the prefill stage, the entire KV Cache is generated and utilized. However, during the decode stage, for selected layers, only the KV Cache entries corresponding to the first and last token sequences from the prefill stage are selectively reused. This significantly reduces the amount of data that needs to be transmitted, with minimal impact on performance.
Also Read:
- Accelerating LLM Thinking: A New Benchmark for Speculative Decoding
- Shrinking AI for Healthcare: Quantization’s Role in Biomedical NLP
Impressive Results
Extensive experiments demonstrate the effectiveness of this new method. Under default settings, the approach achieved a remarkable 20.56% inference speedup. Furthermore, it led to a 4.95 times reduction in data transmission bandwidth consumption. These improvements were consistent across various LLMs and benchmarks, even performing well in unified (non-disaggregated) settings. The research also shows that the method is robust to different calibration sets and hyperparameter settings, and it can be effectively combined with other compression techniques like quantization.
This work represents a significant step forward in making large language models more practical and efficient for deployment, especially in distributed inference environments where prefill-decode disaggregation is increasingly common. By intelligently pruning both model structure and memory, it paves the way for faster and more cost-effective LLM applications.
