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

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

Speakers: Marco Federici, Davide Belli, Mart Van Baalen, Markus Nagel, Paul Whatmough

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This talk, presented by Davide Belli and his colleagues at Qualcomm AI Research, addresses a critical challenge in the deployment of large language models (LLMs) on edge devices: the rapidly diverging growth rates of LLM sizes versus available on-device memory. Titled "Efficient LLM Inference using Dynamic Input Pruning and Cache-Aware Masking," the work proposes novel techniques to enable the execution of gigabyte-scale LLMs on smartphones, which typically offer limited and shared DRAM resources. The core problem lies in the inability of modern LLMs to fit entirely into a device's DRAM, necessitating slow flash storage access.

The presentation highlights that while LLM sizes are growing exponentially, compute speed on processing units is increasing at a slower pace, and DRAM capacity in smartphones is growing almost linearly. This disparity creates a significant bottleneck, especially for auto-regressive decoding, where the entire model needs to be accessed for each token generated. The research specifically focuses on optimizing the memory and compute efficiency of Multi-Layer Perceptron (MLP) blocks, which constitute approximately 75% of an LLM's memory footprint. The proposed solutions, Dynamic Input Pruning (DIP) and its cache-aware variant DIPCA, aim to overcome these memory limitations without severely compromising model accuracy or inference throughput.

The significance of this work lies in its potential to democratize access to powerful LLMs by enabling their efficient execution directly on user devices. This "on-device" capability reduces reliance on cloud infrastructure, enhances privacy, lowers latency, and potentially decreases operational costs. By tackling the fundamental memory constraints, Qualcomm AI Research opens pathways for more sophisticated and personalized AI experiences directly on smartphones, paving the way for a new generation of intelligent edge applications.

Background

▶ Watch: Introduction: LLM growth and hardware memory limitations (0:00)

The rapid advancement and proliferation of large language models (LLMs) have introduced significant challenges for their deployment on resource-constrained edge devices, particularly smartphones. The fundamental problem stems from a stark mismatch in growth rates: LLM sizes are expanding exponentially, while on-device DRAM capacity is growing at a much slower, almost linear rate. This creates a scenario where even moderately sized LLMs, especially after quantization, often exceed the available DRAM on a typical smartphone. For instance, a 4-bit quantized Llama 3 medium model can still occupy over 7 GB, whereas most smartphones offer 4-8 GB of DRAM, a substantial portion of which is reserved for the operating system and other applications.

A naive approach to this memory constraint involves storing the full model in slower flash memory and sequentially loading parts into DRAM as needed. While this might be tolerable for initial prompt processing, it becomes prohibitively slow for auto-regressive decoding. Reading from flash memory is one to two orders of magnitude slower than accessing DRAM, meaning that loading model weights for each token generation step would lead to unacceptably high latencies. Given that MLP blocks account for roughly 75% of an LLM's memory footprint, optimizing these components becomes a critical area of focus for on-device efficiency.

Prior work has explored several avenues to address this problem:

• Quantization: Reducing the numerical precision of model weights (e.g., from 32-bit floating point to 4-bit integers) is a common technique to decrease memory footprint. However, aggressive quantization to 2 or 3 bits often leads to a significant drop in model performance, making it an insufficient standalone solution. Even 4-bit quantization, as seen with the Llama 3 medium example, may not shrink the model enough to fit comfortably within device DRAM.
• Static Pruning: This involves permanently removing a fixed set of weights or neurons from the model. Unstructured pruning, which can achieve high sparsity, introduces overheads in tracking indices for sparse matrix multiplications, which can negate performance gains. Moreover, static pruning, by fixing the pruned weights regardless of the input query, tends to degrade performance as sparsity levels increase.
• Dynamic Pruning (ReLU-based): This approach leverages the inherent sparsity in activations produced by ReLU (Rectified Linear Unit) non-linearities. ReLU activations frequently result in zero values, allowing for dynamic pruning where rows or columns in subsequent weight matrices corresponding to these zero activations are not loaded. This can be extended to predictive dynamic pruning, as exemplified by the Dejavu paper, where a small auxiliary MLP predicts the sparsity pattern based on input values, allowing for pre-emptive pruning of weights in earlier layers (e.g., up-projection matrix).
• Caching Strategies: To manage the larger model footprint, systems like LLM in a Flash store the full model in flash memory but cache frequently used parts of the weights in the faster, limited DRAM. If a required weight is not in the cache, it is loaded from flash, and older elements are evicted to make space.

However, a significant shift in recent LLM architectures has challenged the efficacy of these dynamic pruning methods. Modern LLMs increasingly replace ReLU with GLU (Gated Linear Unit) variants, such as SwiGLU. Unlike ReLU, SwiGLU activations exhibit hardly any natural sparsity. This lack of inherent zeros makes ReLU-based dynamic pruning ineffective. Furthermore, predicting the magnitude of SwiGLU activations, which are highly non-linear, is considerably more difficult than simply predicting the sign of ReLU activations. This fundamental architectural change renders existing predictive dynamic pruning approaches largely unsuitable for contemporary LLMs, necessitating novel solutions that can introduce and exploit sparsity in these newer, denser activation functions.

Key Findings

▶ Watch: Challenge: Existing dynamic pruning fails with modern SwiGLU MLPs (7:00)

The research presented by Qualcomm AI Research addresses the critical challenge of efficient LLM inference on resource-constrained devices by introducing novel pruning and caching strategies tailored for modern LLM architectures. The key findings and contributions are multifaceted:

Firstly, the talk conclusively demonstrates that predictive dynamic pruning methods, which were effective for older LLMs utilizing ReLU non-linearities, are largely ineffective for recent LLMs that employ SwiGLU units. This failure is attributed to the inherent lack of natural sparsity in SwiGLU activations and the significantly higher difficulty in predicting their non-linear magnitudes compared to the simpler sign prediction for ReLU. This finding underscores the necessity for new approaches to dynamic sparsity in contemporary LLMs.

To overcome this limitation, the researchers introduce Dynamic Input Pruning (DIP), a novel non-predictive dynamic pruning method. DIP applies top-K pruning directly to the input and activations of MLP blocks, retaining only the strongest activations and their corresponding weights. Compared to existing approaches, including various static pruning strategies (structured, semi-structured, unstructured) and other dynamic pruning methods like Dejavu and CAT, DIP consistently achieves a superior memory-performance tradeoff. Specifically, DIP can prune up to 40% of the MLP weights with a minimal impact on model performance—less than a 1% drop in accuracy on MLU benchmarks and only a 0.5 perplexity increase on WikiText.

A crucial finding is that DIP combines exceptionally well with model quantization. By applying DIP on top of quantized models (e.g., 2-bit or 3-bit), the researchers demonstrate that increasing dynamic sparsity with DIP yields significantly better performance tradeoffs than merely increasing the quantization aggressiveness alone. This synergy allows for achieving desired memory footprints with less severe quantization, thereby preserving more model accuracy.

Furthermore, the paper introduces a second significant contribution: DIP Cache-Aware (DIPCA). This variant of DIP enhances throughput by intelligently prioritizing neurons that are already present in the device's DRAM cache. DIPCA employs a re-weighting mechanism that considers both the input magnitude and the current cache state, thereby minimizing the number of slow flash memory reads. This strategic cache-awareness results in an approximately 10% improvement in throughput over the standard DIP method, under tight memory constraints, while maintaining high accuracy.

Finally, the research team has open-sourced the full code for their work. This includes a hardware simulator designed to estimate throughput and memory allocation across various hardware specifications, caching, and pruning strategies. This open-source contribution is invaluable for the community, enabling further research, experimentation, and adoption of these efficient inference techniques.

Technical Deep Dive

▶ Watch: DIP performance: Outperforming prior static and dynamic pruning methods (10:40)

The core of the proposed solution lies in addressing the challenges posed by modern LLM architectures, particularly the shift from ReLU to SwiGLU (Gated Linear Unit) non-linearities within Multi-Layer Perceptron (MLP) blocks. A typical MLP block in an LLM consists of an up-projection layer, a non-linearity, and a down-projection layer. While ReLU-based MLPs exhibited natural sparsity in their activations (many zero values), SwiGLU units lack this characteristic, making traditional dynamic pruning methods ineffective. The complexity of predicting the highly non-linear magnitude of SwiGLU activations, as opposed to the simpler sign prediction for ReLU, further exacerbates this issue for predictive pruning techniques.

To circumvent these limitations, the researchers introduce Dynamic Input Pruning (DIP), a non-predictive dynamic pruning approach. Instead of attempting to predict future sparsity, DIP directly leverages the input activations to dynamically prune the MLP. The fundamental mechanism involves applying top-K pruning to the activation magnitudes. This means that for a given input, only the 'K' strongest (highest magnitude) activations are retained, and all other weaker activations are effectively pruned. This technique is applied to both the input vector X and the activations after the up-projection layer. By selecting the strongest activations, DIP ensures that the most salient information is preserved, leading to minimal accuracy degradation.

Specifically, for an MLP block, DIP operates by taking the top-K highest input elements. These selected elements correspond to specific neurons. The method then loads only the corresponding weight rows from the flash memory into DRAM and subsequently to the compute core for matrix multiplication. This dynamic selection and loading prevent unnecessary data transfers for weights associated with less significant activations, thereby reducing the effective memory footprint during inference. The pruning is applied to the up and gate layers of the MLP based on the input, effectively pruning consecutive weights in the matrix multiplication operations.

Building upon DIP, the paper introduces a second, more advanced contribution: DIP Cache-Aware (DIPCA). DIPCA is designed to further optimize throughput by intelligently interacting with the device's memory hierarchy. Traditional DIP, while effective at reducing the overall memory footprint, might still trigger frequent transfers from slow flash memory if the selected top-K neurons are not already present in the faster DRAM cache. DIPCA addresses this by incorporating the current cache state as a prior into its neuron selection process.

The mechanism of DIPCA involves a re-weighting mechanism applied to the normalized input magnitudes. This re-weighting function is designed to prioritize the selection of neurons that satisfy two criteria:

1. They have high activation magnitudes, similar to standard DIP.
2. They are already available in the DRAM cache.

By giving preference to neurons already in cache, DIPCA significantly minimizes the number of memory reads from flash. This is crucial because flash reads are orders of magnitude slower than DRAM access. The re-weighting effectively biases the top-K selection towards cached elements without entirely sacrificing the importance of high-magnitude activations. The goal is to find an optimal balance: maintaining high accuracy by selecting strong neurons, while simultaneously maximizing throughput by reducing the costly flash-to-DRAM transfers.

The overall system design implies a flexible memory management strategy. The full, unpruned model weights are stored in flash. During inference, DIP (or DIPCA) dynamically identifies the necessary weight sub-sections based on the current input. If these sub-sections are not in DRAM, they are fetched from flash and placed into a cache, evicting older elements if necessary. The pruned weights are then loaded into the compute core. This dynamic, input-dependent pruning combined with cache awareness allows for efficient on-device execution of LLMs that would otherwise exceed DRAM capacity.

To facilitate the evaluation and deployment of these techniques, the researchers also developed and open-sourced a hardware simulator. This simulator is a crucial tool for estimating the throughput and memory allocation under various conditions, including different hardware specifications, caching policies, and pruning strategies. This allows developers and researchers to model the performance characteristics of their systems without needing immediate access to specific hardware, accelerating the design and optimization process for on-device LLM inference.

Experimental Setup & Results

▶ Watch: Combining Dynamic Input Pruning with model quantization effectively (12:00)

The experimental evaluation of Dynamic Input Pruning (DIP) and DIP Cache-Aware (DIPCA) focused on demonstrating their effectiveness in reducing memory footprint and improving throughput for LLM inference on resource-constrained devices, while maintaining high accuracy. The experiments were conducted using contemporary LLMs and benchmark datasets relevant to on-device performance.

Models and Quantization:

The primary model used for evaluation was a Llama 3 medium model. While "53 medium model" was mentioned, the context of a 7 GB footprint for 4-bit quantization strongly suggests a Llama 3 7B-parameter class model, which is commonly referred to as "medium" in the Llama family. The experiments also included variants of Llama and Mistral models, indicating the generalizability of the proposed methods.

The study extensively explored the interaction of DIP with model quantization. This included blockwise quantization and vector quantization at various precision levels: 4-bit, 3-bit, and 2-bit. This allowed for a comprehensive analysis of the memory-perplexity/accuracy tradeoff when combining pruning with different quantization regimes.

Baselines for Comparison:

The performance of DIP was rigorously compared against several existing pruning strategies:

• Static Pruning: This included unstructured, structured, and semi-structured variants, representing methods where pruning patterns are fixed prior to inference.
• Dynamic Pruning:
• Dejavu: A predictive dynamic pruning method that attempts to predict active neurons using a small auxiliary MLP.
• CAT: Another dynamic pruning method, serving as a non-predictive baseline for comparison.

Datasets and Metrics:

Model performance was evaluated on two standard benchmarks:

• WikiText: Used to measure perplexity, a common metric for language models, reflecting how well a probability model predicts a sample.
• MLU five shots: Used to measure accuracy, likely referring to a few-shot learning setup on a general language understanding benchmark.

The primary performance metrics were perplexity (lower is better) and accuracy (higher is better). Additionally, the research focused on memory footprint (in GB) and throughput (in tokens per second or similar units), which are critical for on-device deployment.

Headline Results:

1. DIP Performance vs. Baselines: DIP consistently demonstrated superior performance tradeoffs compared to all evaluated static and dynamic pruning baselines. For instance, DIP achieved up to 40% MLP sparsity with remarkably small performance degradation: less than a 1% drop in accuracy on MLU and only a 0.5 perplexity increase on WikiText. This highlights DIP's ability to significantly reduce memory footprint with minimal impact on model quality.
2. DIP and Quantization Synergy: The experiments revealed that DIP combines exceptionally well with model quantization. Applying DIP on top of quantized models (e.g., 2-bit or 3-bit) resulted in much better performance tradeoffs than simply increasing the aggressiveness of quantization alone. This means that a given memory budget could be met with less severe quantization if DIP was employed, thereby preserving higher model accuracy.
3. DIPCA Throughput Improvement: The cache-aware variant, DIPCA, further improved inference throughput. Under tight memory constraints, DIPCA delivered an approximately 10% higher throughput compared to the standard DIP method. This improvement stems from its ability to minimize costly flash memory reads by prioritizing neurons already in DRAM cache.
4. Memory Footprint: The initial problem statement highlighted that a 4-bit Llama 3 medium model could be over 7 GB, exceeding typical smartphone DRAM (4-8 GB, with OS overhead). The proposed methods directly address this by reducing the effective memory footprint during inference, allowing these models to operate within the constraints.
5. Hardware Simulator: A key part of the experimental setup included an open-sourced hardware simulator. This tool allowed for detailed estimation of throughput and memory allocation under various hardware specifications, caching, and pruning strategies, providing valuable insights into real-world performance without necessitating physical hardware deployment for every iteration.

Unanalyzed Aspects:

It is important to note that the work explicitly states that power consumption was not directly analyzed. The primary focus remained on memory footprint and throughput, which were considered the main constraints for enabling model execution on device in the first place. The overhead of the top-K computation itself was also considered negligible compared to the overhead of flash reads. While perplexity and error bars for DIPCA were not shown in the talk slides, it was confirmed that these are available in the full paper and poster, showing consistent performance increases.

Overall, the results strongly validate DIP and DIPCA as effective, practical solutions for deploying large language models on edge devices with limited memory, offering superior performance-memory tradeoffs compared to prior art.

Practical Implications

▶ Watch: DIPCA mechanism: Selecting active neurons based on cache state (14:00)

The research on Dynamic Input Pruning (DIP) and its cache-aware variant (DIPCA) carries significant practical implications for a wide range of stakeholders involved in the development and deployment of AI/ML systems, particularly in the context of edge computing.

For Practitioners and Model Builders:

• Enabling On-Device LLMs: The most direct implication is the ability to deploy larger and more capable LLMs directly onto resource-constrained edge devices like smartphones. This overcomes the critical bottleneck of limited DRAM capacity, which previously restricted the complexity of models that could run locally. Developers can now consider using models like Llama 3 medium (7B parameters) on consumer devices, whereas previously, this was often infeasible without heavy cloud reliance.
• Reduced Cloud Dependency & Costs: By enabling on-device inference, the need for continuous cloud API calls is significantly reduced or eliminated. This translates into lower operational costs for applications, as developers no longer pay for per-token inference. It also opens doors for new business models that are less reliant on cloud infrastructure.
• Enhanced Privacy and Security: Local inference means user data (prompts, generated text) remains on the device, never leaving for a remote server. This drastically improves data privacy and reduces security risks associated with data transmission and storage in the cloud, which is a major concern for many users and regulatory bodies.
• Lower Latency and Improved User Experience: Eliminating network round-trips to the cloud results in near-instantaneous inference, providing a much smoother and more responsive user experience. This is crucial for interactive applications, real-time assistants, and creative tools where latency directly impacts usability.
• Better Accuracy-Memory Tradeoffs: DIP provides a superior way to achieve memory reduction compared to aggressive quantization or static pruning. Practitioners can now achieve a specific memory footprint with less severe quantization, thereby preserving more of the model's original accuracy. This allows for higher-quality on-device AI experiences.

For Infrastructure Teams and Deployers:

• Optimized Resource Utilization: DIPCA's cache-aware mechanism helps infrastructure teams optimize the utilization of precious on-device DRAM and reduce reliance on much slower flash memory. This intelligent management of the memory hierarchy leads to higher throughput and more efficient use of hardware resources.
• Hardware-Software Co-Design Insights: The open-sourced hardware simulator is a powerful tool for infrastructure teams. It allows them to model and predict the performance impact of different hardware specifications, caching strategies, and pruning techniques without extensive physical prototyping. This facilitates informed decisions in hardware procurement and system architecture design for future edge AI devices.
• Simplified Deployment for Edge AI: By providing robust methods for fitting large models into limited memory, the complexity of deploying LLMs at the edge is reduced. This can streamline the development pipeline for edge AI applications, making it easier to bring powerful AI capabilities to a wider range of devices.

Tradeoffs and Limitations:

• Power Consumption (Not Analyzed): While the work significantly addresses memory and throughput, it explicitly did not analyze power consumption. On-device inference, even with optimized memory access, can still be power-intensive. Future work would need to consider the energy efficiency implications, especially for battery-powered devices.
• Flash Memory Dependency: The methods still rely on flash memory to store the full model. While DIPCA minimizes flash reads, they are not entirely eliminated. The speed and endurance of flash memory remain a factor, especially for very frequent model loading or updates.
• Overhead of Dynamic Pruning: Although the overhead of the top-K computation for DIP is deemed negligible compared to flash reads, any additional computation adds to the overall processing load. For extremely low-power or low-compute scenarios, this might still be a consideration.
• Complexity of Implementation: Implementing dynamic pruning and cache-aware strategies can add complexity to the inference runtime compared to simply loading a dense, unpruned model. However, the benefits in memory efficiency and throughput often outweigh this added complexity for on-device deployment.

In summary, DIP and DIPCA offer a compelling solution to the escalating memory challenge in on-device LLM deployment. They empower developers to build more capable, private, and responsive edge AI applications, while providing infrastructure teams with tools for more efficient resource management.

Key Takeaways

• Modern LLMs with SwiGLU units pose a new challenge for dynamic pruning, as they lack the natural sparsity of older ReLU-based architectures, rendering traditional predictive pruning methods ineffective.
• Dynamic Input Pruning (DIP) is a novel, non-predictive method that effectively prunes MLP blocks by retaining only the top-K strongest input activations, achieving significant memory savings (up to 40% MLP sparsity) with minimal accuracy loss (<1% on MLU, 0.5 perplexity on WikiText).
• DIP synergizes powerfully with model quantization, enabling better memory-performance tradeoffs than aggressive quantization alone, thus preserving higher model quality for a given memory budget.
• DIP Cache-Aware (DIPCA) further boosts inference throughput by approximately 10% by intelligently leveraging the DRAM cache state to prioritize neurons already in memory, thereby minimizing slow flash memory reads.
• The proposed techniques are crucial for enabling the deployment of large language models on resource-constrained edge devices like smartphones, addressing the widening gap between LLM size and on-device DRAM capacity.
• The open-sourced code and hardware simulator provide valuable tools for researchers and developers to further explore, optimize, and deploy these efficient LLM inference strategies.

About the Speaker(s)

The talk was presented by Davide Belli, who, along with Marco Federici and colleagues Mart Van Baalen, Markus Nagel, and Paul Whatmough, are affiliated with Qualcomm AI Research. Their work focuses on advancing artificial intelligence technologies, particularly in the realm of efficient machine learning inference for edge devices. Their research contributes to making powerful AI models more accessible and deployable on a wide range of hardware, including smartphones.