UniCache

Task- and Type-Aware KV Cache Compression for Unified Multimodal Models

Paper Code Gallery

Generated by UniCache × Bagel

Wanqi Yang^1,2,3Yuexiao Ma⁴Mei Xie⁵Xiawu Zheng⁶Shiwei Liu^1,2,3

¹ 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

Image Generation Samples

60% KV cache compression

Rose-colored sandstone arch with a pale moon over a desert — **Desert moon**Sculptural rose-colored sandstone arches in a vast desert at blue hour, a large pale moon visible through the central arch...

Red cabin on a frozen fjord under green aurora — Aurora cabin

Flower shop on a rainy Paris street — Rainy flower shop

Portrait of an aviator with goggles and aircraft — Aviator

Violinist playing in warm window light — Violinist

Friendly brass robot tending sunflowers — Clockwork garden

Hummingbird beside an orange hibiscus flower — Hummingbird

Translucent blue glass octopus sculpture — Glass octopus

Embroidered mountain landscape with flowers — Embroidered landscape

Golden stars and crescent moon above an illustrated indigo ocean — Celestial ocean

Three pastries on a marble counter — Patisserie

Botanical garden inside a space station with Earth visible through the window — Space station garden

Colorful hot air balloons rising over a rocky valley at sunrise — Hot air balloons

Image Editing Samples

80% KV cache compression

A rose on a rail, before editing — Source

“a white tulip on the rail”

A white tulip on a rail, UniCache edited output — UniCache

Round orange-frosted cake before editing — Source

“a heart shape cake with orange frosting on a painted wooden plate”

Heart-shaped orange-frosted cake, UniCache edited output — UniCache

Chair and table by a castle-facing window before editing — Source

“anime of a chair and table in front of a window overlooking a castle”

Anime-styled chair, table and castle view, UniCache edited output — UniCache

Task-Dependent Cache Composition

In MoT-based unified multimodal models, the information required by different tasks is represented by distinct KV cache types.

Image understanding

Image understanding uses instruction and source-ViT KVs to represent the instruction and the source image's semantic information, respectively.

Text-to-image generation

Text-to-image generation conditions on instruction KV.

Image editing

Image editing uses instruction, source-ViT, and source-VAE KVs together, with source-VAE KV providing source-image details.

Three tasks in a unified multimodal model activate different combinations of instruction, source-ViT, source-VAE, generated-text, and current-latent KV — Examples of deploying a Mixture-of-Transformers unified multimodal model for multi-task inference and the corresponding task-dependent KV cache.

UniCache

UniCache identifies the KV cache types present in each task and separates compression candidates from protected segments. Through offline calibration of attention distributions, UniCache assigns suitable compression policies to different candidate types.

Task-Aware Cache Segmentation

UniCache partitions the attention context into segments based on the current task and semantic role.

Offline Policy Assignment

UniCache assigns a compression policy to each candidate task--type pair using attention sparsity measured through offline calibration.

Attention-Guided Dynamic Budget Allocation

UniCache estimates cache importance by aggregating attention probabilities over chunks of layers and inference steps, reducing scheduling frequency.

Type-Aware Parallel Compression

Once configured, each policy maintains its own runtime state and compresses only its assigned segments, allowing the segment compression to execute independently and in parallel.

UniCache framework: offline policy assignment, task-aware segmentation, attention-guided budget allocation, parallel compression, and KV reassembly — UniCache 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.

Empirical Results

UniCache effectively mitigates the failure modes encountered by prior KV cache compression methods and achieves balanced quality in all evaluated tasks through task-aware and type-aware KV cache compression.

BAGEL

Method	MME ↑	GenEval ↑	PIE Structure ↓	PIE PSNR ↑
Full KV	2373.21	0.781	0.101	18.805
Global H₂O	2372.60	0.777	0.125	14.663
Global KIVI	2362.31	0.777	0.098	18.576
UniCache	2377.36	0.780	0.100	18.969

SenseNova-U1

Method	MME ↑	MMMU ↑	GenEval ↑	PIE Structure ↓
Full KV	2396.50	47.67	0.887	0.099
Global H₂O	2290.27	47.67	0.875	0.142
Global KIVI	2377.65	47.00	0.888	0.131
UniCache	2383.72	48.00	0.892	0.100

Up to 5×logical KV compression in understanding and editing

Up to 1.78×throughput in the reported long-context, single-A100 physical-engine setting

Get Started

The repository includes a PyTorch reference implementation, a physical CUDA engine, and a full-KV baseline. Model weights are obtained separately.

Installation and backend details ↗

Image understanding · PyTorch backend

python infer.py --backend torch --task understanding \
  --model-path checkpoints/BAGEL-7B-MoT \
  --image /path/to/source.jpg \
  --prompt "Describe the image." \
  --output-dir outputs/understanding