SlimInfer: Accelerating Long-Context LLM Inference via Dynamic Token Pruning

TL;DR

Long-context inference with LLMs is bottlenecked by quadratic attention and heavy KV-cache memory. SlimInfer introduces dynamic, layer-wise hidden-state pruning using fine-grained TokenUnits guided by an information-diffusion principle, paired with predictor-free, asynchronous KV-cache prefetching and CPU offloading to hide IO. Across LLaMA-3.1-8B-Instruct and Qwen2.5-7B-Instruct, SlimInfer achieves up to 2.53x Time-to-First-Token and 1.88x End-to-End latency reductions, while maintaining near-full accuracy on LongBench and reducing prompt KV-cache memory by up to 56%. The approach enables practical long-context deployment on GPUs and edge devices, with open-source code for replication and extension.

Abstract

Long-context inference for Large Language Models (LLMs) is heavily limited by high computational demands. While several existing methods optimize attention computation, they still process the full set of hidden states at each layer, limiting overall efficiency. In this work, we propose SlimInfer, an innovative framework that aims to accelerate inference by directly pruning less critical prompt tokens during the forward pass. Our key insight is an information diffusion phenomenon: As information from critical tokens propagates through layers, it becomes distributed across the entire sequence. This diffusion process suggests that LLMs can maintain their semantic integrity when excessive tokens, even including these critical ones, are pruned in hidden states. Motivated by this, SlimInfer introduces a dynamic fine-grained pruning mechanism that accurately removes redundant tokens of hidden state at intermediate layers. This layer-wise pruning naturally enables an asynchronous KV cache manager that prefetches required token blocks without complex predictors, reducing both memory usage and I/O costs. Extensive experiments show that SlimInfer can achieve up to $\mathbf{2.53\times}$ time-to-first-token (TTFT) speedup and $\mathbf{1.88\times}$ end-to-end latency reduction for LLaMA3.1-8B-Instruct on a single RTX 4090, without sacrificing performance on LongBench. Our code is available at https://github.com/Longxmas/SlimInfer.

Figures (11)

• Figure 1: Accuracy vs. inference efficiency across different acceleration approaches on LongBench bai2024longbenchbilingualmultitaskbenchmark for LLaMA-3.1-8B-Instruct grattafiori2024llama3herdmodels.
• Figure 2: (Left) Illustration of a probing experiment on LLaMA3.1-8B-Instruct grattafiori2024llama3herdmodels. Pruning the hidden state of the critical prompt token "278" (which indicates the correct answer: "$278") in a later layer (right) results in the correct output, whereas pruning prompt tokens in an early layer (middle) leads to an incorrect output. (Right) Visualization of layer-wise attention weights from the decoding token (i.e., response token) "278" to prompt tokens.
• Figure 3: SlimInfer reduces the latency of a layer (i.e., QKV Generation+ Attention+FFN) by prefetching KV cache that is offloaded to CPU, overlapping KV cache fetching with computation. "Sel." means selecting tokens in KV cache (b) or hidden state (c) to prune.
• Figure 4: Overview of the proposed SlimInfer. (i) During inference, early Preserve Layers retain all prompt blocks to support information diffusion (Section \ref{['sec:inf']}), while later Slim Layers prune less relevant blocks to reduce computation (Section \ref{['sec:overview']}). (ii) Each block is divided into fine-grained Token Units for accurate importance scoring (Section \ref{['sec:pruning']}). (iii) When $\texttt{overlap}<\gamma$ (Algorithm \ref{['alg:kv_swap']}), SlimInfer triggers asynchronous KV cache swapping, which naturally overlaps data transfer (prefetching$+$offloading) with computation (Section \ref{['sec:free']}). "T" denotes the current inference step.
• Figure 5: Inference efficiency comparison for LLaMA3.1-8B-Instruct grattafiori2024llama3herdmodels. (Upper) TTFT and (Lower) E2E latency acceleration ratio vs. context length. SlimInfer far outperforms the baselines at long context lengths ($\geq$24k) and remains on par with them in other cases.