Breaking Memory Wall for Fast Edge LLM Inference Using Contextual Sparsity

Abstract

Deploying Large Language Models (LLMs) on memory-constrained edge servers to serve requests from mobile devices is challenging due to their substantial resource demands. The Key-Value (KV) cache and Feed-Forward Network (FFN) parameters consume the majority of available memory. However, existing methods typically rely on static memory partitioning for these components. This rigidity leads to critical performance bottlenecks: (1) An insufficient region for neurons causes neuron swapping or direct computation on the CPU, introducing high latency. (2) As the KV cache grows until it exhausts the allocated region, the system must resort to high-latency fallbacks (e.g., KV cache offloading, recomputation, or neuron swapping). Contextual sparsity of ReLU-based or sparsity-friendly LLMs, where only a small subset of neurons is active during inference, is a promising solution. By leveraging the skewness in neuron activation frequency, we observe two opportunities. (1) We can utilize memory fragmentation to store hot neurons. (2) When memory is exhausted, we prioritize evicting cold neurons from the GPU. Both opportunities improve inference throughput. Based on these insights, we propose ElasticMem, a dynamic memory management framework. It splits KV cache and neuron parameters into blocks to enable flexible neuron placement and eviction. We further design several mapping tables to enable logical-to-physical mapping. Moreover, specialized FFN operators and a CPU-GPU hybrid scheduling pipeline support efficient execution of ElasticMem. Finally, experiments on real-world edge platforms show that, compared with PowerInfer, the state-of-the-art sparsity-aware CPU-GPU hybrid execution baseline, ElasticMem improves throughput by up to <inline-formula><tex-math notation="LaTeX">$5.97\times$</tex-math><alternatives><mml:math><mml:mrow><mml:mn>5</mml:mn><mml:mo>.</mml:mo><mml:mn>97</mml:mn><mml:mo>×</mml:mo></mml:mrow></mml:math><inline-graphic xlink:href="zhao-ieq1-3712257.gif"/></alternatives></inline-formula> across multiple LLMs and deployment scenarios. In addition, compared with dense llama.cpp under memory pressure, ElasticMem achieves up to <inline-formula><tex-math notation="LaTeX">$54.7\times$</tex-math><alternatives><mml:math><mml:mrow><mml:mn>54</mml:mn><mml:mo>.</mml:mo><mml:mn>7</mml:mn><mml:mo>×</mml:mo></mml:mrow></mml:math><inline-graphic xlink:href="zhao-ieq2-3712257.gif"/></alternatives></inline-formula> higher throughput.