UniCache: Task- and Type-Aware KV Cache Compression for Unified Multimodal Models
Authors: Wanqi Yang, Yuexiao Ma, Mei Xie, Xiawu Zheng, Shiwei Liu
Organizations: Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen · Tübingen AI Center · Nanyang Technological University · Independent Researcher · Key Laboratory of Multimedia Trusted Perception and Efficient Computing, Xiamen University
Unified multimodal models combine understanding, generation, and editing within a single network, offering a promising foundation for versatile multimodal applications. However, growing multimodal contexts make KV cache storage and access increasingly costly. Existing KV cache compression methods are typically tailored to specific tasks and single-modality caches, while overlooking changes in cache importance across tasks and timesteps. However, in unified multimodal models, each task involves multiple KV cache types, and both their composition and dynamics differ across tasks. As a result, a single compression policy overlooks task- and type-specific requirements, leading to the loss of critical information and degraded quality across tasks. Based on these findings, we propose UniCache, a training-free framework for task- and type-aware KV cache compression. UniCache identifies the cache segments activated by each task and assigns suitable compression policies through offline calibration. It coordinates their parallel execution under a shared storage budget through attention-guided allocation and task-aware temporal scheduling. Experiments show that UniCache achieves 5× KV cache compression for understanding and editing and 2.5× for generation with negligible quality loss, while increasing throughput by up to 1.78× in long-context settings, significantly improving the practicality of scaling unified multimodal models to longer context.
Figures & tables
Figure 1: UniCache versus uniform KV-cache compression in image editing. At approximately 80% compression, eviction tends to lose source-image details while quantization weakens instruction following. UniCache preserves both.
Figure 2: Left : Examples of deploying a Mixture-of-Transformers (MoT) unified multimodal model for multi-task inference: image understanding, text-to-image generation and image editing. Right : UniCache adopts task- and type-aware KV cache compression and significantly increases throughput and reduces memory overhead for KV cache.
Figure 3: Overview of UniCache. It identifies the cache segments activated by each task, assigns suitable compression policies via offline calibration and allocates type-specific compression budgets from the corresponding attention statistics, compressing different cache types online and in parallel.
Figure 4: Task- and type-dependent KV cache attention patterns in BAGEL at layer 8. (a) Attention mass over the complete editing KV sequence, preserving the original segment order, where B denotes boundary KV. (b) Current-latent K/V recomputed at each denoising step. (c–e) Attention distributions of different cache types across image understanding, text-to-image generation, and image editing. We report attention probability averaged over heads and queries.
Figure 5: Each point represents a sampled subset for a task–KV-type pair with eviction or quantization at the same compression ratio. Details are provided in Appendix A .
Figure 6: Attention mass across inference steps. Understanding exhibits step-to-step fluctuations, while generation and editing show overall declines at type-dependent rates.
Model
Method
Understanding
Generation
Editing
MME ↑
MMBench ↑
MMMU ↑
GenEval ↑
GEdit QO↑
PIE Structure ↓
PIE PSNR ↑
PIE CLIP Target ↑
BAGEL
-
2373.21
85.14
53.33
0.781
6.947
0.101
18.805
28.174
H 2 O
2372.60
85.31
52.78
0.777
6.860
0.125
14.663
28.567
KIVI
2362.31
83.68
51.78
0.777
6.929
0.098
18.576
27.812
StreamingLLM
2372.60
85.31
52.78
0.313
6.347
0.151
13.021
29.118
SnapKV
2371.01
85.22
51.91
0.558
6.551
0.149
12.932
28.942
Table 1: Cross-task results on BAGEL and SenseNova-U1. MMBench reports circular accuracy (%), MMMU reports validation accuracy (%), and GEdit reports QO . Best and second-best compressed results within each model are bold and underlined .
Method
MME
GenEval
Comp.
Perception
Cognition
Total
Comp.
Single
Two
Count
Color
Position
Binding
Image Acc.
Prompt Acc.
Overall
BF16
0%
1680.35
692.86
2373.21
0%
0.997
0.922
0.781
0.846
0.520
0.623
0.773
0.897
0.781
H 2 O
80.07%
1679.74
692.86
2372.60
59.47%
0.997
0.909
0.788
0.864
0.505
0.600
0.768
0.902
0.777
KIVI
80.69%
1685.53
676.79
2362.31
59.12%
1.000
0.912
0.778
0.840
0.505
0.628
0.768
0.897
0.777
StreamingLLM
79.94%
1679.74
692.86
2372.60
58.01%
0.863
0.051
0.419
0.540
0.003
0.003
0.287
0.405
0.313
SnapKV
79.94%
1679.58
691.43
2371.01
61.25%
0.997
0.503
0.659
0.832
0.095
0.263
0.536
0.673
0.558
Table 2: Detailed MME and GenEval results on BAGEL. Comp. denotes KV-cache compression. Higher is better for all quality metrics. Best and second-best results among compressed methods are bold and underlined , respectively.
MME
GenEval
PIE-Bench
Setting
Perception ↑
Cognition ↑
Total ↑
Overall ↑
Structure ↓
PSNR ↑
CLIP Target ↑
BAGEL
1680.35
692.86
2373.21
0.781
0.101
18.805
28.174
Full UniCache
1681.65
695.71
2377.36
0.780
0.101
19.100
28.171
w/o Task/Type Awareness
1679.74
692.86
2372.60
0.767
0.103
17.875
28.221
w/o Protection
1680.47
692.86
2373.32
0.567
0.102
18.959
27.235
w/o Attention-Guided Allocation
1680.32
688.01
2368.33
—
0.121
17.277
27.036
Table 3: Method component ablations across tasks. Targets are approximately 80% compression on MME and 60% on GenEval and PIE-Bench. Dashes indicate inapplicable settings.
Schedule
GenEval ↑
Structure ↓
LPIPS ↓
SSIM ↑
CLIP Target ↑
CLIP Edit ↑
Constant
0.776
0.102
0.184
0.757
28.010
27.070
Linear decay
0.780
0.102
0.184
0.757
28.187
27.285
Attention-guided decay
0.780
0.101
0.179
0.763
28.211
27.293
Table 4: Ablation of temporal budget schedules. We change the schedule methods for global budget scheduling and report results in GenEval and PIE-Bench.
Figure 7: Ablation of compression ratio across different methods and benchmarks. Across different benchmarks, UniCache is robust to high compression ratio and maintains balanced performance at very high compression ratios.
Figure 8: Runtime efficiency of UniCache across context lengths. Throughput is measured in tokens per second for understanding and denoising steps per second for generation and editing. Memory reports the persistent managed KV cache memory summed over all layers.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Understanding
Text-to-Image
Image Editing
Target compression
80.00
60.00
80.00
H 2 O
80.07
59.47
80.09
KIVI
80.69
59.12
79.59
StreamingLLM
79.94
58.01
79.94
SnapKV
79.94
61.25
79.93
PyramidKV
79.94
61.25
79.93
Appendix
Table 5: Target and realized KV cache compression ratios (%) on BAGEL.
Figure 9: Task- and type-dependent KV-cache attention patterns in SenseNova-U1 at layer 8. Each panel shows token-level attention probability averaged over heads and queries for one example, with tokens kept in their original order. Source-native-visual KV denotes SenseNova-U1’s native source-image representation.
Figure 10: Image generation samples produced by our UniCache optimized BAGEL model. The compression rate is 60%. The model successfully handles complex descriptions, rendering high-quality textures and correct object placements, demonstrating that our KV cache compression framework preserves the generative performance of the model.
Figure 11: Image editing samples produced by our UniCache-optimized BAGEL model. Each group shows the source image, editing prompt, and edited output from left to right. UniCache preserves the original model’s ability to follow instructions while preserving visual content unrelated to the requested edits.
Multimodal Large Language Models (MLLMs) have advanced unified reasoning over text, images, and videos, but their inference is hindered by the rapid growth of key-value (KV) caches. Each visual input expands into thousands of tokens, causing caches to scale linearly with context length and remain resident in GPU memory throughout decoding, which leads to prohibitive memory overhead and latency even on high-end GPUs. A common solution is to compress caches under a fixed allocated budget at different granularities: token-level uniformly discards less important tokens, layer-level varies retention across layers, and head-level redistributes budgets across heads. Yet these approaches stop at allocation and overlook the heterogeneous behaviors of attention heads that require distinct compression strategies. We propose HybridKV, a hybrid KV cache compression framework that integrates complementary strategies in three stages: heads are first classified into static or dynamic types using text-centric attention; then a top-down budget allocation scheme hierarchically assigns KV budgets; finally, static heads are compressed by text-prior pruning and dynamic heads by chunk-wise retrieval. Experiments on 11 multimodal benchmarks with Qwen2.5-VL-7B show that HybridKV reduces KV cache memory by up to 7.9× and achieves 1.52× faster decoding, with almost no performance drop or even higher relative to the full-cache MLLM.
Bowen Zeng, Feiyang Ren, Jun Zhang +4
The State Key Laboratory of Blockchain and Data Security, Zhejiang University · Hangzhou High-Tech Zone (Binjiang) Institute of Blockchain and Data Security · Hangzhou Dianzi University, Hangzhou, China
Multimodal Large Language Models face severe challenges in computational efficiency and memory consumption due to the substantial expansion of the visual KV cache when processing long visual contexts. Existing KV cache compression methods typically rely on the "persistence of importance" hypothesis to prune tokens. However, this approach proves fragile in multimodal settings due to two key issues: 1) Visual tokens display "deferred importance," initially exhibiting low salience but becoming pivotal during later decoding, which can lead to premature eviction. 2) Discrete pruning disrupts the inherent spatial continuity of visual cues. To address these challenges, we propose RetentiveKV, an entropy-driven KV cache optimization method that reformulates KV eviction from "discrete context truncation" to "continuous memory evolution" based on State Space Models. Our method leverages information entropy to quantify the information potential of low-attention tokens and integrates tokens scheduled for eviction into a continuous state space through entropy-guided state transitions, enabling their dynamic reactivation when semantic relevance arises during subsequent decoding. Extensive experiments on multimodal benchmarks demonstrate that RetentiveKV achieves 5.0 times KV cache compression and 1.5 times decoding acceleration.
Sihao Liu, YuFan Xiong, Zhonghua Jiang +2
Zhejiang University · Alibaba · Shanghai Institute for Advanced Study of Zhejiang University
KV caches are a major bottleneck in long-context inference and long-form generation with large language models. Existing training-free eviction policies largely rely on proxy importance signals, such as attention mass, to decide which past tokens to retain. We argue that cache compression should instead preserve the predictive behavior of the full-cache model, retaining entries whose removal would substantially change the model's output distribution. We propose Behavior-Preserving KV Cache Compression, a training-free framework that scores candidate evictions by estimating the compressed-cache logits induced by their removal and evaluating the resulting KL to the full-cache next-token distribution. Using pre-eviction forward statistics, the method avoids running separate masked forward passes for each candidate. Across diverse architectures and both prefill-time and generation-time compression, our method delivers substantial gains in downstream task quality over lightweight attention-based heuristics at matched retained-KV budgets, with the largest gains under aggressive compression. It achieves these gains with additional compression-time computation while retaining an end-to-end speedup over full-cache inference in our evaluated settings.
Doo Hwan Hwang, Junyoung Jang, Junho Na +2
Kim Jaechul Graduate School of AI, KAIST, Daejeon, Republic of Korea · KT Corporation, Seoul, Republic of Korea