Jenga: Two-Level Memory Allocation for Heterogeneous LLM Serving

Chen Zhang, Kuntai Du, Shu Liu, Woosuk Kwon, … Joseph Gonzalez, Ion Stoica · 2025
arXiv:2503.18292 Interactive page

The one-size-fits-all tax

Llama 3.2 Vision: 32 self-attention layers read only text-token KV; 8 cross-attention layers read only image-token KV. PagedAttention reserves (T+I)×40 slots; the ideal is 32T+8I.
-slots reserved
-slots needed
-wasted fraction (computed)
self-attn × text (read)cross-attn × image (read)reserved, never read
Hover a cell.

Two-level allocator

image-type small pagestext-type small pagesToy: 256 B and 384 B → 768 B large pages

Per-layer prefix hits Illustrative

Click KV tokens or Mamba checkpoints to evict/restore them. A model-wide hit needs every layer type satisfied at the same prefix. (The paper checkpoints every 512 tokens; here every 4.)
hitmissimage token (C, D)
hooks: update_last_accessset_prefix_lengthget_possible_prefix

Reported results

Hover a bar.

Headline audit

Waste fraction vs image-token share, from the episode's slot formula (reserved 40(T+I), needed 32T+8I). Computed, not measured
Measured on this model: 1.71× throughput. At 0% image tokens the formula still gives 20% (cross-attention layers reserve text slots); the episode notes the chat-mix sweep was never run.

What the paper doesn't show

Hover a tile.