1 00:00:01,000 --> 00:00:09,300 [Hal Turing] Alrighty! Thanks for tuning in! Hello AI world! I am your host, Hal Turing, and my co-host is Dr. Ada Shannon. 2 00:00:09,300 --> 00:00:47,325 [Hal Turing] Today's paper is "Jenga: Effective Memory Management for Serving LLM with Heterogeneity," posted to arXiv in March 2025. The first author is Chen Zhang at Tsinghua University. The co-authors include Woosuk Kwon, Zhuohan Li, Joseph Gonzalez and Ion Stoica from UC Berkeley, plus colleagues at the University of Chicago and an independent researcher. The question is this: PagedAttention assumes every layer stores same-sized per-token state and needs the whole prefix. Modern models break both assumptions. How much memory and throughput does that cost, and can a two-level allocator recover it? 3 00:00:47,325 --> 00:01:05,450 [Dr. Ada Shannon] Several of these authors built vLLM and PagedAttention, so they're testing their own baseline. That matters because they know exactly which assumptions were baked in. The fix is an allocator that works with the existing attention kernels and lets each layer type decide what it keeps. 4 00:01:05,450 --> 00:01:09,075 [Hal Turing] Start with the basics. What's being allocated? 5 00:01:09,075 --> 00:01:30,075 [Dr. Ada Shannon] The KV cache. Each attention layer saves the key and value tensors for every earlier token, so it doesn't recompute them at each step. The paper gives about 1.2 gigabytes for one 10,000-token request on Llama 3.1 8B. To keep the GPU busy you batch many requests, and every one of them keeps its cache resident. So memory sets the batch size, and batch size sets throughput. 6 00:01:30,075 --> 00:01:35,150 [Hal Turing] And PagedAttention is how vLLM packs those caches in? 7 00:01:35,150 --> 00:01:59,025 [Dr. Ada Shannon] Yes. Kwon and colleagues at Berkeley, 2023. It cuts the cache into fixed-size pages, and a page table maps each request's logical blocks to physical pages. That's borrowed from operating system virtual memory. Any free page can serve any request, so you get almost no external fragmentation. The limit of the analogy is that an OS page holds undifferentiated bytes. Here, a page implicitly assumes every token and every layer produces an identically sized entry. 8 00:01:59,025 --> 00:02:02,500 [Hal Turing] So that's assumption one. What's the second? 9 00:02:02,500 --> 00:02:14,275 [Dr. Ada Shannon] Full-prefix dependency. To generate the next token, every earlier token matters, so a request's tokens all share one lifetime. They're allocated together and freed together. 10 00:02:14,275 --> 00:02:16,800 [Hal Turing] And which models break that? 11 00:02:16,800 --> 00:02:51,626 [Dr. Ada Shannon] Take sliding-window attention first. A layer attends only to the last W tokens, so older entries for that layer are dead weight. Gemma-2 and Ministral mix those layers with full-attention layers. Mamba hybrids, like Jamba from AI21 and Hymba from NVIDIA, replace some attention with state-space layers. A Mamba layer keeps one large recurrent state per sequence, not a per-token entry, and that state is bigger than one token's KV entry. Cross-attention vision models such as Llama 3.2 Vision are a third case. The paper also mentions speculative decoding with two models and vision embedding caches. 12 00:02:51,626 --> 00:02:54,876 [Hal Turing] Give me a concrete case of the damage. 13 00:02:54,876 --> 00:03:19,801 [Dr. Ada Shannon] Llama 3.2 Vision has 32 self-attention layers that need KV only for the T text tokens. It also has 8 cross-attention layers that need KV only for the I image tokens. PagedAttention treats all 40 layers alike, so it reserves T plus I, times 40 slots. The ideal is 32 times T plus 8 times I. An image can be thousands of tokens, so the text layers reserve slots for image tokens they'll never read. The cross-attention layers do the same for text. 14 00:03:19,801 --> 00:03:24,526 [Hal Turing] Why not just use a smaller page for everything, or a bigger one? 15 00:03:24,526 --> 00:03:43,076 [Dr. Ada Shannon] The paper's options are all bad. Size pages for the largest object and small objects waste the rest of the page. Size them for the smallest and big objects span many pages, which complicates the kernels. Partition memory statically per layer type and you strand capacity whenever the workload mix shifts. 16 00:03:43,076 --> 00:03:51,776 [Hal Turing] Prefix caching has the same problem, right? Reusing a finished request's KV for a new request with the same prefix. 17 00:03:51,776 --> 00:04:24,451 [Dr. Ada Shannon] Right. That skips recomputing a shared system prompt or a long document. In a plain transformer a hit is simple: the whole prefix is present or it isn't. The paper calls the new situation prefix-subset dependency. A window layer needs only the last W tokens of the prefix, so a hit can survive with the early tokens evicted. A Mamba layer needs one state saved at exactly that boundary, and you can't roll a state back to an earlier position. Cross-attention needs only the image tokens. A hit for the whole model is valid only if every layer type can be satisfied at that prefix. 18 00:04:24,451 --> 00:04:28,677 [Hal Turing] So eviction can't be one global LRU anymore. 19 00:04:28,677 --> 00:04:42,527 [Dr. Ada Shannon] Correct. A single LRU treats tokens as interchangeable. It would keep out-of-window entries and could report a hit one layer type can't honor. Each layer type needs its own rule for what's still needed. 20 00:04:42,527 --> 00:04:44,652 [Hal Turing] So what's the roadmap? 21 00:04:44,652 --> 00:05:05,152 [Dr. Ada Shannon] Two pieces. First, a two-level allocator, in the spirit of the Linux slab allocator. Large pages are sized so each layer type can carve them into its own small pages without waste. Second, per-layer caching and eviction rules, so each type states what it needs. Let's start with how those large pages are sized and why that lets existing kernels keep working. 22 00:05:05,152 --> 00:05:14,752 [Hal Turing] So the large pages are LCM-sized. Give me the toy example, because least common multiple sounds like it could go wrong in a few ways. 23 00:05:14,752 --> 00:05:54,827 [Dr. Ada Shannon] The paper's toy is Llama 3.2 Vision with each layer's per-token entry set to 128 bytes. Image tokens only go through the two cross-attention layers, so 256 bytes. Text tokens go through three self-attention layers, so 384. The LCM of 256 and 384 is 768. Memory is cut into 768-byte large pages. The image allocator splits some of them into three 256-byte small pages. The text allocator splits others into two 384-byte small pages. When every small page in a large page is free, the large page goes back to the shared pool, and either type can take it next. 24 00:05:54,827 --> 00:05:59,952 [Hal Turing] So the sizes divide evenly and nothing is stranded between types. 25 00:05:59,952 --> 00:06:24,027 [Dr. Ada Shannon] Right. That removes external fragmentation, meaning gaps too small for anything to use. There's one layout change to make it work. Standard PagedAttention slices memory by layer first, then by page. Jenga slices by page first, then by layer inside the page, so each small page is contiguous and can change owners. The kernels still receive a start pointer, a page size and a list of page IDs. The authors say they changed no CUDA and about 4000 lines of Python on vLLM. 26 00:06:24,027 --> 00:06:29,552 [Hal Turing] Why LCM, though? What goes wrong with a smaller or bigger shared page size? 27 00:06:29,552 --> 00:06:59,677 [Dr. Ada Shannon] A GCD page, the greatest common divisor, has no waste. But fast kernels want contiguous memory along certain dimensions, and tiny pages break that. MuxServe took that route and limited itself to models with the same head size. A MAX page, the largest entry size, wastes space on the smaller types. You can add more tokens per page to compensate, but Jamba would need 1344 tokens per page. The average ShareGPT request is about 1085, so one page would be larger than the typical request. 28 00:06:59,677 --> 00:07:02,552 [Hal Turing] And LCM can't blow up? 29 00:07:02,552 --> 00:07:17,902 [Dr. Ada Shannon] It can in principle. Across models in vLLM 0.6.4, the worst was Jamba at 84 times the small page, and they saw no slowdown. That's a measurement on current models. It says nothing about future ones. 30 00:07:17,902 --> 00:07:24,528 [Hal Turing] Large pages hold many small pages. What stops them from filling with scattered leftovers? 31 00:07:24,528 --> 00:07:47,278 [Dr. Ada Shannon] That's request-aware allocation. Requests allocate pages interleaved, but a finished request frees all its pages at once. If requests A and B share large pages, A finishing frees small pages, yet almost no large page is fully empty. So Jenga puts a request's small pages in the same large page when it can. Then the whole large page comes back when the request ends. If that fails, it falls back to taking free small pages tied to other requests. 32 00:07:47,278 --> 00:07:51,203 [Hal Turing] Now the caching side. Who decides what to evict? 33 00:07:51,203 --> 00:08:15,578 [Dr. Ada Shannon] Each layer type, through three hooks: update_last_access, set_prefix_length and get_possible_prefix. Two rules govern eviction. Balanced means the layer types stay compatible, because a model-wide hit needs the prefix in every type. Jenga gives tokens from one request the same timestamp across types. Aligned means evicting the same tokens everywhere. Otherwise each layer might keep tokens the other dropped, and no prefix survives in both. set_prefix_length breaks ties so those pages go together. 34 00:08:15,578 --> 00:08:19,078 [Hal Turing] What does a sliding-window layer do differently? 35 00:08:19,078 --> 00:08:39,078 [Dr. Ada Shannon] It refreshes the timestamp only for tokens inside the window. Older tokens look stale and go first. For hits, take a ten-token request, A through J. A plain self-attention layer can hit any prefix up to I. A window layer needs only the last few tokens of the prefix cached, so a prefix can hit even with early tokens gone. 36 00:08:39,078 --> 00:08:42,828 [Hal Turing] Mamba is different, since there's one big state. 37 00:08:42,828 --> 00:09:02,203 [Dr. Ada Shannon] They checkpoint that state every 512 tokens, so hits land on multiples of 512. It's a simple heuristic. For vision caches, they give each image a random priority, so a whole image is evicted together. Dropping one token forces re-running the vision encoder, so scattered evictions cost more. 38 00:09:02,203 --> 00:09:05,853 [Hal Turing] The paper also covers two other cases. 39 00:09:05,853 --> 00:09:36,953 [Dr. Ada Shannon] Yes. In speculative decoding, a small draft model and a big one have different KV sizes per token. Across the heterogeneous models, throughput improves about 1.58x on average over the uniform-page baseline, and on plain Llama it matches a hand-tuned allocator. For vision embeddings, Jenga frees each image's embedding as chunked prefill consumes it. That gives about 1.88x the throughput and 1.60x better latency than vLLM, which re-runs the vision encoder. 40 00:09:36,953 --> 00:09:40,303 [Hal Turing] On to the main results. What's the setup? 41 00:09:40,303 --> 00:10:08,228 [Dr. Ada Shannon] An H100 with 80 GB and an L4 with 24 GB, against vLLM, changing only memory management. The models are Llama Vision, Gemma-2, Ministral, Jamba, a character.ai-style model, PyramidKV and plain Llama as a control. Gemma-2 and Ministral run on arXiv-QA, long article questions. The shorter MMLU-pro tops out at 3076 tokens, below their window sizes, so they'd behave like full-attention models. Hymba isn't evaluated, since vLLM lacks its kernels. 42 00:10:08,228 --> 00:10:10,353 [Hal Turing] Okay, throughput. 43 00:10:10,353 --> 00:10:33,803 [Dr. Ada Shannon] On the H100, six of the seven models land between roughly 1.0x and 2.1x. Plain Llama is 1.03x, so the machinery costs nothing there. The character.ai model reaches 4.92x. The paper's average is 1.80x. For a job taking an hour on vLLM, a typical hybrid finishes in half to two-thirds of that. Gemma-2 would take about 80 percent. The character.ai model would take about 12 minutes. 44 00:10:33,803 --> 00:10:37,203 [Hal Turing] Where does the gain come from? Is it measured? 45 00:10:37,203 --> 00:11:05,603 [Dr. Ada Shannon] On Ministral, yes. vLLM keeps KV for out-of-window tokens and wastes 38.2 percent of cache on average. Jenga wastes 0.04 percent. With twenty long requests arriving together, 55 to 110 thousand tokens each, the average decode batch is 5.4 against about 2.6 for vLLM, SGLang and TGI. With many articles, hit rate is up to 1.6x and throughput up to 1.77x. 46 00:11:05,603 --> 00:11:07,353 [Hal Turing] What's the cost? 47 00:11:07,353 --> 00:11:24,103 [Dr. Ada Shannon] Per-token latency gets worse, because bigger batches mean more work per step. With three or fewer articles everything fits, and Jenga is slightly slower because it allocates twice. And that batch test is a burst of extremely long requests, the best case for Jenga. 48 00:11:24,103 --> 00:11:30,978 [Hal Turing] Let's go after the headline. The abstract says up to 4.92x. Which model is that? 49 00:11:30,978 --> 00:11:54,178 [Dr. Ada Shannon] One bar, and by the Figure 13a labels it's the character.ai model. That's the authors' reimplementation, built from a public blog post on top of Llama 70B in FP8. It isn't a released checkpoint. The window and sharing parameters they used aren't in the text I have, so I can't tell you how generous the setup is. The released models land lower: Gemma-2 at 1.26x, Llama Vision at 1.71x, Jamba at 1.78x. 50 00:11:54,178 --> 00:11:57,378 [Hal Turing] What about the 1.80x average? 51 00:11:57,378 --> 00:12:15,803 [Dr. Ada Shannon] My arithmetic says it's a geometric mean of the seven bars. The plain mean is about 2.04. The paper doesn't say which mean it uses, so that's an inference. Also, the abstract says 4.92x and the conclusion says 4.91x. Small, but sloppy for a headline. 52 00:12:15,803 --> 00:12:21,028 [Hal Turing] Then the 79.6% memory number. Same problem? 53 00:12:21,028 --> 00:12:54,553 [Dr. Ada Shannon] Different problem. Section 3.2 derives it openly. Take the average MMMU-pro request, about 6,200 image tokens and 43 text tokens. PagedAttention reserves forty layer-slots for every one of those tokens. The ideal needs only eight for image tokens and thirty-two for text. That works out to roughly four of every five slots holding nothing any layer reads. So the calculation is fine. The wording is the issue. The abstract and conclusion call it a memory utilization gain, but it's a waste fraction computed from an average request shape, for an image-heavy benchmark. 54 00:12:54,553 --> 00:12:57,678 [Hal Turing] And the measured result on that model? 55 00:12:57,678 --> 00:13:18,403 [Dr. Ada Shannon] 1.71x throughput. So waste sets a ceiling, and throughput depends on what else limits you. On a normal chat mix with fewer image tokens per request, the formula's waste shrinks toward zero. The paper never runs that sweep, so what survives there is unmeasured. The measured waste they do have is Ministral: 38.2% against 0.04%. 56 00:13:18,403 --> 00:13:22,978 [Hal Turing] Ministral and Gemma-2 needed a different dataset, though. 57 00:13:22,978 --> 00:13:46,028 [Dr. Ada Shannon] They did, and the defense is reasonable: a sliding-window benchmark needs requests longer than the window, and Ministral averages over 92,000 tokens on arXiv-QA. But it means 1.26x and 2.08x are long-context numbers. With short prompts those models degrade to full attention. The Llama control at 1.03x shows what you get when the structure doesn't matter. How much real traffic is long enough to benefit, the paper can't tell us. 58 00:13:46,028 --> 00:13:56,528 [Hal Turing] What's the baseline? Is vLLM wasting 38% because of PagedAttention, or because of that particular vLLM version? 59 00:13:56,528 --> 00:14:26,628 [Dr. Ada Shannon] That's the right question, and the paper doesn't separate them. The baseline is one fixed vLLM version that keeps out-of-window pages. I'd want vLLM patched with per-layer freeing but no LCM allocator. That would show how much comes from the allocator and how much from a simple eviction fix. The SGLang and TGI comparison doesn't settle it either. It's a burst of twenty requests on one model, and TGI stops early because it lacks an ignore-EOS option, which muddies the step counts. 60 00:14:26,628 --> 00:14:29,178 [Hal Turing] Other comparisons missing? 61 00:14:29,178 --> 00:15:09,203 [Dr. Ada Shannon] Several. For Mamba prefix caching, Marconi, from Rui Pan, Zhuang Wang and colleagues in 2024, is concurrent work on what to admit and evict for hybrid models. Jenga says it could integrate Marconi but doesn't compare. Checkpointing every 512 tokens is a simple heuristic, and on prompts around a thousand tokens it can leave a lot to recompute. Then vAttention, from Ramya Prabhu and colleagues in 2024. It uses GPU virtual memory for contiguous KV instead of new page abstractions. Jenga dismisses it on driver overhead and missing prefix-subset tracking, with no head-to-head. And the whole hybrid story rests on one model, Jamba, H100 only. Hymba isn't evaluated. 62 00:15:09,203 --> 00:15:15,078 [Hal Turing] What about the allocator itself? How does it relate to the OS lineage you mentioned? 63 00:15:15,078 --> 00:15:46,453 [Dr. Ada Shannon] Bonwick's slab allocator from 1994. Jenga calls its design a special case, simpler because every size is known at startup. That's also the limit. It assumes static geometry. Eighty-four times the small page was the worst case among the models they checked, and the paper reports no slowdown. But a large page is reclaimable only when every small page in it is free. The paper doesn't report large-page occupancy or how often the fallback allocation paths trigger under high churn. Request-aware packing also assumes a request's pages die together, which sliding windows and shared prefixes violate. 64 00:15:46,453 --> 00:15:48,103 [Hal Turing] And the costs? 65 00:15:48,103 --> 00:16:11,153 [Dr. Ada Shannon] Time per output token gets worse with larger batches. The authors say it would match vLLM at equal batch size, which is plausible, but the abstract's claim of no latency impact doesn't hold as stated. There's also slight prefix-cache overhead from double allocation when everything fits. No variance is reported, and the FP8 and hardware mix makes comparisons harder. Some baselines may simply have been memory-starved, like a 70B model on one H100. 66 00:16:11,153 --> 00:16:12,878 [Hal Turing] Who should care? 67 00:16:12,878 --> 00:16:35,428 [Dr. Ada Shannon] Anyone serving long-context sliding-window models, Mamba hybrids, or cross-attention vision models. If your traffic is short chat on plain transformers, expect roughly nothing. Before trusting any number, I'd want a production model, a normal image and text mix, a window-length sweep, and Marconi as a rival. The paper itself leaves open multi-model serving and offloading granularity, where a large page's size might clash with SSD transfer sizes. 68 00:16:35,428 --> 00:16:37,728 [Hal Turing] So where does that leave us? 69 00:16:37,728 --> 00:16:58,728 [Dr. Ada Shannon] The mechanism holds up. It gives a kernel-compatible way to stop paying PagedAttention's one-size-fits-all tax, with measured gains on long-context workloads. The boundary is that typical gains run about 1.0 to 2.1x, the 4.92x is a reimplemented model, and 79.6% is a computed waste fraction, not a measured gain. 70 00:16:58,728 --> 00:17:02,253 [Hal Turing] Thanks for listening, and we'll see you next time.