Rimuru boosts memory use and speed in large language model serving
Rimuru: Unified Memory Management for Modern LLM Serving via Elastic Expert Offloading
Operating SystemsMachine Learning
Summary
Large language models (LLMs) need a lot of memory for their big parts and the information they use while talking, which can slow things down especially when memory is limited. The authors show that splitting memory into fixed regions for different tasks wastes space and limits speed when demand changes. Their system Rimuru shares memory flexibly among all parts by letting different model pieces use the same pages differently. Rimuru also adjusts memory use based on demand and can directly use scattered model parts to avoid copying data. This leads to much faster processing and reduces waiting time when running LLMs.
What this means in practice
- •For llm system engineers: Serve large language models on devices with limited memory by dynamically reallocating memory for model parts and attention states to improve throughput and reduce response delay.
- •For cloud infrastructure teams: Optimize memory use in cloud servers running LLM workloads by sharing memory across different tensor types to increase concurrency under heavy loads.
Authors
Zhiyuan Fang, Dingyang Li, Jianfeng Zhong, Bin Wang, Weixi Zhu, Cunshu Ni, Zicong Hong, Wuhui Chen, Zibin Zheng
Abstract
Modern LLMs strain device memory through large model weights and attention states, particularly in resource-constrained scenarios. Existing systems separately provision memory for model weights and attention states, imposing fixed fixed boundaries that cannot adapt to the changing memory demands of workloads. Under heavy workloads, insufficient state capacity limits request concurrency and reduces throughput; under light workloads, idle memory in one region cannot be reallocated to another, wasting device memory. Meanwhile, modern LLM architectures introduce heterogeneous tensors with diverse layouts and lifetimes, further complicating memory management. These problems motivate unified memory management for heterogeneous tensors. However, a shared pool alone does not make memory directly reusable across tensors with different allocation granularities and execution layouts. This paper presents Rimuru, a modern LLM serving system that unifies memory management through elastic expert offloading. Rimuru establishes a shared page pool to enables elastic memory sharing across heterogeneous tensors. To accommodate different allocation granularities, Rimuru introduces memory views that allow each tensor type to interpret the shared memory space using its own page size and page table. Furthermore, Rimuru combines huge-page alignment with a multi-level page allocator that splits and merges pages, enabling cross-type capacity reuse. A page-table-aware fused MoE kernel directly accesses non-contiguous expert weights, avoiding copy overhead. Building on these mechanisms, Rimuru employs a workload-aware adjuster that dynamically reallocates pages across tensor types according to runtime demands. Implemented on top of vLLM, Rimuru remains compatible with modern LLMs. Evaluations show that Rimuru achieves up to 19.03x baseline throughput and reduces TTFT by 98.1% on average.