# Jenga: Effective Memory Management for Serving LLM with Heterogeneity

arXiv: https://arxiv.org/abs/2503.18292 · HTML: https://mlsanity.com/paper/2503.18292/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 42 of its 63 references that are in the map.
- First posted on arXiv 24 Mar 2025.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [LLaVA-OneVision: Easy Visual Task Transfer](https://mlsanity.com/paper/2408.03326/) (2024)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [InternVL: Scaling up Vision Foundation Models and Aligning for Generic Visual-Linguistic Tasks](https://mlsanity.com/paper/2312.14238/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving](https://mlsanity.com/paper/2401.09670/) (2024)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [MMMU-Pro: A More Robust Multi-discipline Multimodal Understanding Benchmark](https://mlsanity.com/paper/2409.02813/) (2024)
- [PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling](https://mlsanity.com/paper/2406.02069/) (2024)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [Jamba: A Hybrid Transformer-Mamba Language Model](https://mlsanity.com/paper/2403.19887/) (2024)
- [PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU](https://mlsanity.com/paper/2312.12456/) (2023)
- [PaliGemma 2: A Family of Versatile VLMs for Transfer](https://mlsanity.com/paper/2412.03555/) (2024)
- [Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention](https://mlsanity.com/paper/2403.19708/) (2024)

## Built on by (most cited first, 25 of 25)

- [TokenCake: A KV-Cache-centric Serving Framework for LLM-based Multi-Agent Applications](https://mlsanity.com/paper/2510.18586/) (2025)
- [CacheWise: Understanding Workloads and Optimizing KVCache Management for Efficiently Serving LLM Coding Agents](https://mlsanity.com/paper/2606.16824/) (2026)
- [MindSpeed RL: Distributed Dataflow for Scalable and Efficient RL Training on Ascend NPU Cluster](https://mlsanity.com/paper/2507.19017/) (2025)
- [SPAD: Specialized Prefill and Decode Hardware for Disaggregated LLM Inference](https://mlsanity.com/paper/2510.08544/) (2025)
- [Gated Differentiable Working Memory for Long-Context Language Modeling](https://mlsanity.com/paper/2601.12906/) (2026)
- [Cornfigurator: Automated Planning for Any-to-Any Multimodal Model Serving](https://mlsanity.com/paper/2512.14098/) (2025)
- [Unifying Sparse Attention with Hierarchical Memory for Scalable Long-Context LLM Serving](https://mlsanity.com/paper/2604.26837/) (2026)
- [VibeServe: Can AI Agents Build Bespoke LLM Serving Systems?](https://mlsanity.com/paper/2605.06068/) (2026)
- [Effective Distillation to Hybrid xLSTM Architectures](https://mlsanity.com/paper/2603.15590/) (2026)
- [Not All Tokens Are Worth Caching: Learning Semantic-Aware Eviction for LLM Prefix Caches](https://mlsanity.com/paper/2605.18825/) (2026)
- [Characterization of Multi-Model Agentic AI Systems on General Tasks via Trace-Driven Simulation](https://mlsanity.com/paper/2606.01725/) (2026)
- [Tangram: Unlocking Non-Uniform KV Cache Compression for Efficient Multi-turn LLM Serving](https://mlsanity.com/paper/2606.06302/) (2026)
- [SwiftCache: Efficient LLM Serving for Multi-turn Conversations with Heterogeneous KV Cache Sharing](https://mlsanity.com/paper/2606.16135/) (2026)
- [SmartGen: Seamless Disaggregated LLM Inference with Selective KV Cache Transfer](https://mlsanity.com/paper/2607.28150/) (2026)
- [Joint Encoding of KV-Cache Blocks for Scalable LLM Serving](https://mlsanity.com/paper/2601.03067/) (2026)
- [Crossflow: Prefill-Decode Elasticity for Agentic LLM Serving](https://mlsanity.com/paper/2609.27085/) (2026)
- [Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference](https://mlsanity.com/paper/2605.06046/) (2026)
- [SPEAR: A System for Post-Quantization Error-Adaptive Recovery Enabling Efficient Low-Bit LLM Serving](https://mlsanity.com/paper/2606.11244/) (2026)
- [Characterizing Software Aging in GPU-Based LLM Serving Systems](https://mlsanity.com/paper/2606.11916/) (2026)
- [From Tensor Buffer to Distributed Memory Hierarchy: A Survey of KV Cache Management for LLM Serving](https://mlsanity.com/paper/2607.02574/) (2026)
- [TensorCast: The Missing Tensor Management Layer in Large Language Model Infrastructure](https://mlsanity.com/paper/2608.06007/) (2026)
- [BOOST: Concurrent Access to Host Memory and HBM to Accelerate LLM Inference](https://mlsanity.com/paper/2609.13592/) (2026)
- [TierKV: Long-Context On-Device LLMs via Predictive Multi-Tier KV Caching](https://mlsanity.com/paper/2609.21172/) (2026)
- [WavePP: High-Throughput Pipeline Parallel LLM Prefill under Prefix Reuse](https://mlsanity.com/paper/2609.35263/) (2026)
- [Scaling LLM Inference via Parallelism: Throughput-Latency Trade-off](https://mlsanity.com/paper/2610.19405/) (2026)

## Abstract (arXiv)

Large language models (LLMs) are widely used but expensive to run, especially as inference workloads grow. To lower costs, maximizing the request batch size by managing GPU memory efficiently is crucial. While PagedAttention has recently been proposed to improve the efficiency of memory management, we find that the growing heterogeneity in the embeddings dimensions, attention, and access patterns of modern LLM architectures introduces new challenges for memory allocation. In this paper, we present Jenga, a novel memory allocation framework for heterogeneous embeddings in LLMs. Jenga tackles two key challenges: (1) minimizing memory fragmentation when managing embeddings of different sizes, and (2) enabling flexible caching and eviction policies tailored to the specific token-dependency patterns of various layers. Jenga employs a two-level memory allocator, leveraging the least common multiple (LCM) of embedding sizes to optimize memory usage and providing APIs to express layer-specific caching logic to enhance memory reuse. We implemente Jenga on vLLM, a state-of-the-art LLM inference engine, and evaluate it with diverse LLMs, datasets, and GPU configurations. Evaluations show that Jenga improves GPU memory utilization by up to 79.6%, and increases serving throughput by up to 4.92x (1.80x on average).