[Paper Review] LLM as a System Service on Mobile Devices

Wangsong Yin, Mengwei Xu+2 more|arXiv (Cornell University)|Mar 18, 2024

Business Process Modeling and Analysis5 citations

TL;DR

The paper proposes LLMS, a system-service design that manages on-device LLM contexts through chunk-wise KV cache compression and swapping to enable low-latency, persistent LLM services across apps. It demonstrates up to two orders of magnitude switching-latency reduction over baselines and significant improvements without accuracy loss.

ABSTRACT

Being more powerful and intrusive into user-device interactions, LLMs are eager for on-device execution to better preserve user privacy. In this work, we propose a new paradigm of mobile AI: LLM as a system service on mobile devices (LLMaaS). Unlike traditional DNNs that execute in a stateless manner, such a system service is stateful: LLMs execution often needs to maintain persistent states (mainly KV cache) across multiple invocations. To minimize the LLM context switching overhead under tight device memory budget, this work presents LLMS, which decouples the memory management of app and LLM contexts with a key idea of fine-grained, chunk-wise, globally-optimized KV cache compression and swapping. By fully leveraging KV cache's unique characteristics, it proposes three novel techniques: (1) Tolerance-Aware Compression: it compresses chunks based on their measured accuracy tolerance to compression. (2) IO-Recompute Pipelined Loading: it introduces recompute to swapping-in for acceleration. (3) Chunk Lifecycle Management: it optimizes the memory activities of chunks with an ahead-of-time swapping-out and an LCTRU (Least Compression-Tolerable and Recently-Used) queue based eviction. In evaluations conducted on well-established traces and various edge devices, \sys reduces context switching latency by up to 2 orders of magnitude when compared to competitive baseline solutions.

Motivation & Objective

• Motivate LLM as a system service (LLMaaS) for mobile devices to improve privacy, availability, and efficiency.
• Highlight the memory bottleneck of persistent LLM contexts and the inadequacy of traditional app-level memory management.
• Introduce LLMS as a memory-management framework that decouples app memory from LLM context memory.
• Develop novel techniques to maximize context-switching performance under memory constraints.

Proposed method

• Divide KV cache into fixed-size chunks (e.g., 16 tokens) for independent compression and swapping.
• Apply Tolerance-Aware Compression to allocate higher compression to chunks with lower information density.
• Introduce a Swapping-Recompute Pipeline to overlap recomputation with I/O during context loading.
• Use a Chunk Lifecycle Management strategy with an LCTRU eviction queue and ahead-of-time swapping to hide reclaim latency.
• Provide an Android- Service-like API design where apps hold LLMCtx contexts and interact via callLLM.
• Empirically evaluate on COTS devices with Llama2-7B/OPT-7B across 72-hour traces, comparing to LMK, disk swapping, and vLLM baselines.

(a) The architecture of an LLM. We use Llama2 as an illustrative example. Embedding layer is omitted for simplicity.

Experimental results

Research questions

• RQ1 How can LLM context memory be managed efficiently on mobile devices when LLMs are exposed as a system service?
• RQ2 What chunking, compression, and scheduling strategies minimize context-switching latency while preserving accuracy?
• RQ3 Can a Swapping-Recompute pipeline and an LCTRU-based eviction policy outperform traditional app-level memory management for LLM contexts?
• RQ4 What are the performance gains of LLMS compared with baselines across real device traces and common LLMs?

Key findings

• LLMS reduces LLM context switching latency by up to 100x (two orders of magnitude) compared to conventional app-level memory management and disk swapping baselines.
• Compared to state-of-the-art vLLM with static 8-bit quantization, LLMS achieves up to 20x switching-latency reduction and on average 9.7x.
• 11 No accuracy loss is observed across 6 datasets when using LLMS with its chunk-wise tolerance-aware compression.
• LLMS demonstrates practical on-device LLMaaS viability on Jetson Orin NX, Jetson TX2, and MI14 smartphone with Llama2-7B and OPT-7B.
• The approach leverages the unique properties of KV cache (chunk-wise structure, compressibility, and recomputability) to optimize context management.

(b) The autoregressive inference with KV cache.