RAP: Runtime Adaptive Pruning for LLM Inference

TL;DR

RAP is a reinforcement learning-based framework that adaptively prunes large language models during inference to optimize memory usage and performance in real-time scenarios.

Contribution

It introduces a dynamic, runtime-aware pruning method that jointly considers model weights and KV-cache, outperforming fixed heuristic approaches.

Findings

RAP outperforms state-of-the-art baselines in experiments.

It effectively adapts to runtime memory variations and workload demands.

First method to jointly consider model weights and KV-cache dynamically.

Abstract

Large language models (LLMs) excel at language understanding and generation, but their enormous computational and memory requirements hinder deployment. Compression offers a potential solution to mitigate these constraints. However, most existing methods rely on fixed heuristics and thus fail to adapt to runtime memory variations or heterogeneous KV-cache demands arising from diverse user requests. To address these limitations, we propose RAP, an elastic pruning framework driven by reinforcement learning (RL) that dynamically adjusts compression strategies in a runtime-aware manner. Specifically, RAP dynamically tracks the evolving ratio between model parameters and KV-cache across practical execution. Recognizing that FFNs house most parameters, whereas parameter -light attention layers dominate KV-cache formation, the RL agent retains only those components that maximize utility within the current memory budget, conditioned on instantaneous workload and device state. Extensive experiments results demonstrate that RAP outperforms state-of-the-art baselines, marking the first time to jointly consider model weights and KV-cache on the fly.