🤖 AI

Reducing Peak Memory Usage for Modern Multimodal Large Language Model Pipelines

This paper proposes a sequential, structure-aware key-value cache compression mechanism that enforces a fixed memory budget during the prefill stage of multimodal large language models, significantly reducing peak memory usage while maintaining generative performance.

Original authors: Junwan Kim, Hyunkyung Bae

Published 2026-04-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Junwan Kim, Hyunkyung Bae

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Overfilled Backpack"

Imagine you are a detective (the AI) trying to solve a mystery by looking at a massive photo album (the visual input) and reading a list of clues (the text prompt).

In modern AI, to understand high-resolution photos or long videos, the system has to break them down into thousands of tiny puzzle pieces called tokens. To solve the mystery, the AI needs to remember every single piece it has seen so far. It keeps these memories in a special "backpack" called the KV Cache.

The Issue:
If you show the AI a 4K image or a 10-minute video, the number of puzzle pieces explodes. The backpack gets so heavy and full that it bursts before the AI even starts solving the mystery. This is called an Out-of-Memory (OOM) error.

Existing solutions try to fix this by letting the AI fill the backpack completely, then throwing away the "less important" items after the whole album is loaded. But by then, it's too late—the backpack already exploded.

The Solution: The "Smart Packing" Strategy

The authors of this paper propose a new way to pack: Don't wait until the end to throw things away. Pack and purge as you go.

They call this a "Prefill-Aware" approach. Instead of filling the backpack to the brim and then compressing it, they enforce a strict memory budget (a maximum weight limit) from the very first step.

Here is how their method works, broken down into three simple concepts:

1. The "Block-by-Block" Approach

Instead of dumping the entire photo album into the backpack at once, the AI looks at the images in small chunks (blocks).

Old Way: Dump 1,000 photos in. Backpack breaks.
New Way: Look at 10 photos. Put them in the backpack. Check the weight. If it's too heavy, immediately throw out the least useful photos before adding the next 10.
Result: The backpack never gets too heavy, no matter how big the album is.

2. Two Ways to Decide What to Throw Away

The AI needs to know which puzzle pieces to keep and which to toss. They use two different "filters":

The "Question-Aware" Filter (For Single Tasks):
If you ask a specific question (e.g., "Where is the cat?"), the AI looks at your question first. It uses your question as a magnet to pull in the most relevant parts of the image and push away the boring background (like the sky or a wall). It keeps the "cat" and throws away the "sky."
- Analogy: Like a librarian who only keeps books on your specific topic and shoves everything else off the cart immediately.
The "Question-Agnostic" Filter (For Future Chats):
Sometimes, you haven't asked a question yet, or you are in the middle of a long conversation where the context changes. In this case, the AI looks for variety. It keeps the most unique and strange-looking puzzle pieces and throws away the repetitive ones (like 50 identical pixels of a blue sky).
- Analogy: Like a chef who keeps the most unique ingredients (truffles, saffron) and throws away the 50th bag of plain flour, ensuring the final dish has flavor even if they don't know the recipe yet.

The Results: A Lighter Backpack, Same Smarts

The researchers tested this on powerful AI models (like InternVL and Qwen) using high-resolution images and long videos.

Memory Usage: The new method kept the memory usage flat and low, even when the input size grew huge. The old method crashed (ran out of memory) as soon as the images got too big.
Performance: Surprisingly, throwing away 90% of the memory didn't make the AI much dumber. It still answered questions correctly because it kept the important pieces and only discarded the redundant ones.
The Trade-off: The only downside is speed. Because the AI has to pack and unpack in small chunks rather than doing it all at once, it takes a tiny bit longer to start answering. However, this is a small price to pay for being able to handle huge images without crashing.

The Big Picture

Think of this paper as inventing a smart, self-cleaning suitcase for AI.

Previously, if you tried to pack a whole wardrobe into a suitcase, it would rip open. Now, this new method lets you pack the wardrobe piece by piece, instantly discarding the socks you don't need as you go, so the suitcase never gets too heavy to carry. This allows AI to "see" and understand massive, high-definition videos and images without needing a supercomputer's worth of memory.

1. Problem Statement

Modern Multimodal Large Language Models (MLLMs) are increasingly capable of processing high-resolution images and long video sequences. However, this scalability introduces a critical bottleneck: peak memory consumption during the prefill stage.

The KV Cache Explosion: MLLMs rely on Key-Value (KV) caching to accelerate autoregressive decoding. While this reduces decoding complexity from $O(N^2)$ to $O(N)$ , the cache size grows linearly with the sequence length.
Visual Token Overhead: High-resolution images (tiled) and dense video sampling generate thousands of vision tokens before generation begins. For example, a 4K image can yield over 42,000 vision tokens.
The Prefill Spike: Existing KV-cache compression methods (e.g., SnapKV, H2O) typically operate after the full multimodal context is encoded. Consequently, the GPU memory must first accommodate the full, uncompressed KV cache of the entire input, leading to Out-of-Memory (OOM) errors during the prefill stage, even if the final cache could theoretically fit.
Limitations of Current Solutions:
- Post-prefill compression: Fails to address the initial memory spike.
- Token pruning: Discards input tokens at the input level, often ignoring the heterogeneous roles of tokens across different layers and attention heads, risking the loss of structurally important information.

2. Methodology

The authors propose a prefill-aware, structure-aware KV-cache compression framework that enforces a fixed memory budget during input processing rather than after.

Core Concept: Block-wise Prefill with Online Eviction

Instead of constructing the full cache and then compressing it, the method processes the input sequence in contiguous blocks.

Sequential Processing: The input sequence is partitioned into blocks of size $b$ .
Incremental Encoding: For each block, the model computes the KV pairs and appends them to the cache.
Immediate Budget Enforcement: If the cache size exceeds a fixed memory budget $M$ , an eviction strategy is triggered immediately to prune the cache back to $M$ before processing the next block.
Structural Alignment: Block boundaries are aligned with the structural regularities of visual data (e.g., spatial tiles in images, temporal frame groups in videos) to ensure eviction happens at semantically meaningful granularities.

Eviction Strategies

The framework offers two strategies depending on the availability of query information:

Query-Aware Eviction (Single-Turn):
- Based on SnapKV.
- Uses proxy query tokens extracted from the textual prompt to compute cross-attention scores against cached keys.
- Retains tokens with high attention scores (visually salient regions relevant to the prompt) and evicts low-scoring redundant tokens.
Query-Agnostic Eviction (Multi-Turn/No Query):
- Based on KeyDiff.
- Does not rely on a specific query. Instead, it preserves representational diversity.
- Calculates the mean vector ( $\mu$ ) of cached keys and prioritizes retaining keys that deviate most from this mean (outliers), ensuring rare visual features are kept while redundant, average features are removed.

3. Key Contributions

Shift in Pipeline Paradigm: Moves from a "process first, compress later" approach to a "compress as you prefill" approach, effectively bounding peak memory usage throughout the entire input processing phase.
Structure-Aware Compression: Exploits the inherent spatial and temporal redundancies in visual data by aligning block processing with visual structures (tiles/frames), rather than treating tokens as a flat sequence.
Dual-Strategy Framework: Introduces both query-aware (for task-specific relevance) and query-agnostic (for general diversity) eviction mechanisms to handle diverse interaction patterns (single-turn vs. multi-turn).
Hybrid Execution Optimization: Proposes a hybrid strategy that performs "bulk forward" passes for the initial tokens to minimize latency, switching to block-wise processing only when necessary to stay within memory limits.

4. Experimental Results

The method was evaluated on InternVL3.5-8B/14B and Qwen2.5-VL-7B/32B across image (ImageNeedleInHaystack, V*) and video (MLVU, Video-MME) benchmarks.

Memory Efficiency:
- The method maintains a nearly constant peak memory footprint regardless of the number of input image tiles.
- Baseline methods (Full Cache) grow linearly and fail (OOM) beyond ~36 tiles.
- Achieves up to ~90% KV-cache compression (reducing budget to 1024 tokens from full context) without OOM errors.
Performance Preservation:
- Minimal Accuracy Degradation: At a budget of 1024 tokens, the average accuracy drop is minimal (e.g., ~0.78 points for InternVL3.5-8B with Query-Aware).
- Non-Monotonic Gains: In some video tasks, aggressive compression actually improved performance slightly, suggesting the removal of redundant temporal noise leads to more focused representations.
Trade-offs:
- Latency: Sequential block processing increases Time-To-First-Token (TTFT) compared to bulk processing. However, this is a controllable trade-off that enables the processing of inputs that would otherwise be impossible (OOM).
- Block Size Sensitivity: Performance is highly dependent on aligning block sizes with the model's native visual tokenization (e.g., 784 tokens for a 28x28 grid). Misalignment causes performance drops.

5. Significance and Conclusion

This work addresses a fundamental scalability barrier in MLLMs: the inability to process high-resolution or long-duration visual inputs due to prefill memory spikes.

Practical Impact: It enables the deployment of MLLMs on hardware with fixed memory constraints (e.g., consumer GPUs) for tasks requiring high-resolution or long-context visual understanding.
Theoretical Insight: It demonstrates that how visual context is processed during prefill is as critical as the final cache size. By respecting the structural regularities of visual data, compression can be applied early without sacrificing reasoning capabilities.
Future Direction: The paper suggests that future models could be explicitly trained with prefill-stage compression in mind to further improve robustness, moving beyond inference-time optimization.

In summary, the proposed framework transforms the memory-latency trade-off from a hard failure (OOM) into a manageable parameter, making large-scale multimodal inference significantly more practical and accessible.