PDTrim: Targeted Pruning for Prefill-Decode Disaggregation in Inference
This work addresses deployment efficiency for LLM inference, offering incremental improvements in pruning techniques for disaggregated settings.
The paper tackles the computational and memory costs of deploying Large Language Models by proposing a targeted pruning method for prefill-decode disaggregation inference, resulting in improved performance, faster inference, and a 4.95× reduction in data transmission bandwidth consumption.
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 same (default) settings, our method achieves improved performance and faster inference, along with a 4.95$\times$ reduction in data transmission bandwidth consumption.