Efficient LLM Inference using Dynamic Input Pruning and Cache-Aware Masking

TL;DR

The paper tackles the memory bottleneck of on-device LLM inference by proposing Dynamic Input Pruning (DIP), a predictor-free sparsification method, and a cache-aware extension (DIP-CA) that leverages DRAM cache state to boost throughput with minimal accuracy loss. DIP eliminates the need for expensive retraining, while lightweight LoRA adapters help recover some performance loss; DIP-CA further improves cache hit rates and reduces DRAM loads via a re-weighting scheme. Across multiple SwiGLU-based models and hardware configurations, DIP outperforms static pruning and prior dynamic sparsity approaches in memory footprint and token throughput, with notable gains such as a 46% memory reduction and 40% throughput increase on Phi-3-Medium relative to the dense baseline. The work also demonstrates favorable compatibility with quantization and highlights cache-aware masking as a key factor for practical, mobile deployments, supported by open-source HW simulators and artifacts. Overall, the methods enable high-throughput, memory-efficient on-device LLM inference without reliance on costly predictor training.

Abstract

While mobile devices provide ever more compute power, improvements in DRAM bandwidth are much slower. This is unfortunate for large language model (LLM) token generation, which is heavily memory-bound. Previous work has proposed to leverage natural dynamic activation sparsity in ReLU-activated LLMs to reduce effective DRAM bandwidth per token. However, more recent LLMs use SwiGLU instead of ReLU, which results in little inherent sparsity. While SwiGLU activations can be pruned based on magnitude, the resulting sparsity patterns are difficult to predict, rendering previous approaches ineffective. To circumvent this issue, our work introduces Dynamic Input Pruning (DIP): a predictor-free dynamic sparsification approach, which preserves accuracy with minimal fine-tuning. DIP can further use lightweight LoRA adapters to regain some performance lost during sparsification. Lastly, we describe a novel cache-aware masking strategy, which considers the cache state and activation magnitude to further increase cache hit rate, improving LLM token rate on mobile devices. DIP outperforms other methods in terms of accuracy, memory and throughput trade-offs across simulated hardware settings. On Phi-3-Medium, DIP achieves a 46\% reduction in memory and 40\% increase in throughput with $<$ 0.1 loss in perplexity when compared to streaming the dense model from Flash. The open source code for HW simulator, methods, and experiments in this paper is available at https://github.com/Qualcomm-AI-research/dynamic-sparsity .

Figures (14)

• Figure 1: Overview of Cache-Aware Dynamic Input Pruning for a vector-matrix multiplication. The Cache-Aware masking component (DIP-CA) outputs a column sparsity mask based on current input activations and weights currently cached in DRAM. The subset of required weights are retrieved from DRAM cache, if available, or loaded from Flash in case of cache misses.
• Figure 2: Comparing the largest LLMs by year of release zhao2023survey against NPU and DRAM trends in recent iPhone devices. An exponential fit highlights the linear increase in DRAM size opposed to the exponential increases in NPU processing power and model sizes.
• Figure 3: GLU Activation Magnitude distribution for Layer 31 Mistral-7B (left), and the corresponding re-trained ReLU-fied Turbosparse Mistral (right).
• Figure 4: Layer activation density for different GLU thresholding strategies at 50% MLP target density on a Mistral-7B model. Darker shaded region indicate the standard deviation, while lighter shading covers minimum and maximum values. Setting a threshold for each layer or each token result in similar perplexity. The dense model perplexity is $5.249$.
• Figure 5: Diagrams describing the sparsity schemes for GLU pruning, Gate pruning, Predictive pruning and Input pruning. Activations with larger magnitude are visualized in darker colors, while dark blue indicates weights used for computation for the current input.