Jenga: Two-Level Memory Allocation for Heterogeneous LLM Serving

This episode examines "Jenga," a paper from the team behind vLLM's PagedAttention, which asks what that memory manager costs on modern models that break its two core assumptions: that every layer stores identically sized per-token state, and that every layer needs the full prefix. It covers sliding-window layers, Mamba-style recurrent states, and cross-attention vision models. In a Llama 3.2 Vision example, text layers end up reserving slots for image tokens they never read. The discussion also covers why prefix caching and a single global LRU eviction policy no longer work when each layer type needs a different part of the prefix. It then walks through the proposed fix, a slab-style two-level allocator that sizes large pages at the least common multiple of the layer types' entry sizes (a worked toy of 256 and 384 bytes gives 768-byte pages) and lets each layer type split them into its own small pages. Listeners get a clear explanation of how this recovers memory and throughput while working with existing attention kernels.

Sources:

1. Jenga: Effective Memory Management for Serving LLM with Heterogeneity — Chen Zhang, Kuntai Du, Shu Liu, Woosuk Kwon, Xiangxi Mo, Yufeng Wang, Xiaoxuan Liu, Kaichao You, Zhuohan Li, Mingsheng Long, Jidong Zhai, Joseph Gonzalez, Ion Stoica, 2025

http://arxiv.org/abs/2503.18292

2. Efficient Memory Management for Large Language Model Serving with PagedAttention — Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, Ion Stoica, 2023

https://scholar.google.com/scholar?q=Efficient+Memory+Management+for+Large+Language+Model+Serving+with+PagedAttention

3. Jenga: Effective Memory Management for Serving LLM with Heterogeneity — Chen Zhang, Kuntai Du, Shu Liu, Woosuk Kwon, Xiangxi Mo, Yufeng Wang, Xiaoxuan Liu, Kaichao You, Zhuohan Li, Mingsheng Long, Jidong Zhai, Joseph Gonzalez, Ion Stoica, 2025

https://scholar.google.com/scholar?q=Jenga%3A+Effective+Memory+Management+for+Serving+LLM+with+Heterogeneity

4. Jamba: A Hybrid Transformer-Mamba Language Model — Opher Lieber et al. (AI21 Labs), 2024

https://scholar.google.com/scholar?q=Jamba%3A+A+Hybrid+Transformer-Mamba+Language+Model

5. Hymba: A Hybrid-head Architecture for Small Language Models — Xin Dong et al. (NVIDIA), 2024

https://scholar.google.com/scholar?q=Hymba%3A+A+Hybrid-head+Architecture+for+Small+Language+Models