PD-Disaggregation Pruning for LLM Optimization

• PD-disaggregation pruning is a targeted model compression technique that decouples the prefill and decode phases to meet distinct computational constraints.
• It applies phase-aware block pruning and token-aware KV cache reduction, achieving up to 20.56% speedup and a 4.95× decrease in data transmission bandwidth.
• Using simulated annealing for block selection, the method optimally balances pruning operations to retain 94–96% accuracy despite 13–25% sparsity.

Prefill-Decode (PD) Disaggregation Pruning is a targeted model compression technique designed to optimize the inference efficiency of LLMs when their serving pipelines are architecturally split into distinct prefill and decode phases. This method addresses the heterogeneous computational and memory requirements present in modern LLM deployments—where prefill consumes high throughput and memory, generating the complete key-value (KV) cache from a long prompt, and decode iteratively generates output tokens, often on separate machines and under bandwidth limitations. Classic pruning approaches, which apply uniform sparsification, fail to account for these phase-specific performance and resource constraints. PD-disaggregation pruning instead combines phase-aware block pruning and token-aware KV cache pruning to simultaneously reduce compute latency and inter-machine data transfer, with minimal degradation in downstream accuracy (Zhang et al., 29 Aug 2025).

1. Motivation and PD Disaggregation Formulation

Modern Transformer-based LLMs implement a two-stage inference protocol: in the prefill stage, a sequence of $N$ prompt tokens is processed to build full hidden state representations and a layerwise KV cache $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$. In the decode stage, new tokens are generated autoregressively, appended to the cache for each step. Prefill and decode typically have different latency, compute, memory, and communication bottlenecks. Uniform pruning strategies—removing identical blocks or parameters for both—either spare needless compute in decode or erode overall accuracy. PD-disaggregation pruning explicitly seeks a solution decoupled along phase lines:

Let $\mathcal{B}$ be the set of blocks/layers, $k$ the quota of blocks to prune or merge, and $L$ the model depth. The objective is to minimize accuracy loss $(\Delta_{\text{acc}})$ under constraints

$$\text{sparsity} \geq k/L, \quad \text{bandwidth reduction} \geq B$$

contingent on accurate block sensitivity estimation and communication profiling, subject to the differing usage of each model subgraph in prefill and decode (Zhang et al., 29 Aug 2025).

2. Block-Level Pruning and Distillation Sets

PD-disaggregation pruning identifies candidate blocks for removal or pairwise merging using redundant representation criteria derived from a small calibration corpus. For each block $i$, redundancy is computed as cosine similarity between input and output activations:

$r_i = \cos(h_{i-1}, h_i)$

Blocks with large $r_i$ are considered weakly informative and seeded into an initial pruning set:

$$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$0

Distillation sets are constructed based on merged pair redundancy:

$$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$1

Pairs $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$2 are eligible for merging if $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$3, typically with $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$4, provided neither member is in $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$5. This step allows the framework to distill one block's capacity into a neighbor, trading parameter loss for preserved representational ability.

3. Iterative Block Removal via Simulated Annealing

Rather than removing or merging blocks greedily by individual redundancy measures, PD-disaggregation pruning uses a simulated-annealing search procedure to identify the near-globally optimal subset of $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$6 operations. The search alternates, at each step, between removing a current prune/merge candidate and adding a new one, scoring each candidate set $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$7 by calibration accuracy. Replacements are sampled with probability proportional to the redundancy score. The system retains and ultimately selects the configuration $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$8 that minimizes accuracy drop on the calibration batch.

After this search, phase-specific testing is performed for each candidate operation: if removing a block exclusively in the decode phase improves calibration accuracy by more than 3% relative to also removing it in prefill, the removal is committed for decode only. Otherwise, it is committed for both phases. Standard hyperparameters are $$K_{\text{pre}}^{(i)}, V_{\text{pre}}^{(i)} \in \mathbb{R}^{N \times d}$$9 (initial temperature), $\mathcal{B}$0 (decay), $\mathcal{B}$1, with $\mathcal{B}$2 typically set between 9% and 25% (Zhang et al., 29 Aug 2025).

4. Token-Aware KV Cache Pruning

Transmission of the full per-layer KV cache (size $\mathcal{B}$3) from prefill to decode nodes is bandwidth-intensive. PD-disaggregation pruning reduces this overhead by exploiting the observation that decode-stage attention scores concentrate on early and recent tokens. For each attention head $\mathcal{B}$4 in layer $\mathcal{B}$5, aggregate attended mass over the first and last $\mathcal{B}$6 tokens:

$\mathcal{B}$7

Layers where any head’s $\mathcal{B}$8 (e.g., $\mathcal{B}$9) are excluded. For others, layer scores are computed as

$k$0

with $k$1 the intra-layer mean and std, and the top $k$2 layers retained for partial cache transmission. Selected layers send only the first and last $k$3 entries; unselected layers send the full $k$4. Block-pruned layers do not transmit KV cache at all. The resulting bandwidth is

$k$5

yielding a proportional reduction of $k$6.

5. Experimental Evaluations

PD-disaggregation pruning has been evaluated across a range of LLMs (LLaMA3.1-8B, LLaMA2-13B, Qwen2.5-7B/14B/32B, OPT-6.7B) and benchmarks covering MMLU, PIQA, BoolQ, ARC, WNLI, SST-2, IFEval, and HumanEval within a PD-serving context. Default settings employ 256 calibration samples.

Key empirical results:

• Achieves a 20.56% average inference speedup on LLaMA3.1-8B (latency reduction from 287 ms to 228 ms per query).
• Reduces cross-phase data transmission bandwidth by 4.95× (4.0 GB to 0.8 GB).
• Maintains zero-shot accuracy drops at 2–5% absolute, preserving 94–96% of dense model performance at 13–25% block sparsity.
• Iterative search offers a 1–2% absolute accuracy advantage over greedy approaches; attention-filtered KV pruning outperforms random or uniform head selection by ∼1–2% (Zhang et al., 29 Aug 2025).

The technique is orthogonal to quantization: with 8-bit activation-aware quantization (AWQ), accuracy preservation remains similar.

6. Implementation Considerations and Trade-offs

The PD-disaggregation pruning algorithm is implemented in PyTorch using 2× H100 GPUs and HuggingFace Transformers. Calibration latency is modest: 4–9 s for pruning set construction, 4–10 s for cache calibration (versus 26–50 s for other methods), totaling around 44 s on an 8B model, compared to 183 s for SLEB. Pruning can slightly increase decode-stage workload if post-pruning block/KV alignment is suboptimal. Settings for $k$7 are tunable for practical SLO targets.

Potential extensions include memory management for dynamic subgraph topologies under PD, integration with Mixture-of-Experts pruning, and co-design with multi-tenant scheduling policies (Zhang et al., 29 Aug 2025).

7. Comparative Perspective and Future Directions

In direct comparison to baselines such as LLM-Pruner, FLAP, Shortened LLaMA, ShortGPT, SLEB, and SliceGPT, PD-disaggregation pruning consistently delivers the strongest speed and bandwidth reductions while retaining the majority of model accuracy. Ablation studies demonstrate that the decoupled phase-aware strategy is superior to uniform or greedy methods. Accuracy exhibits graceful degradation with increasing pruning ratios, and the orthogonality to quantization suggests a broad applicability to practical LLM deployments.

Future work may address memory management for dynamically pruned architectures, introduction of PD-disaggregation pruning to expert routing in Mixture-of-Experts settings, and integration with serving-side load balancing and multi-user resource scheduling policies. This suggests that phase-aware approaches may become increasingly critical as deployment architectures for LLMs and multimodal models diversify (Zhang et al., 29 Aug 2025).