PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wanqi Yang, Yuexiao Ma, Mei Xie, Xiawu Zheng, Shiwei Liu

arXiv:2609.32831v1cs.LG

TL;DR

Growing multimodal contexts make KV-cache storage and access costly, while cache types and their importance vary across tasks and stages. UniCache assigns task- and type-aware policies and budgets, achieving 5× compression for understanding and editing and 2.5× for generation with negligible quality loss, alongside up to 1.78× long-context throughput.

  • Problem

    Growing multimodal contexts raise KV-cache costs, while existing compression methods generally target fixed tasks or cache types despite task-dependent cache composition and importance.

    Source: §1 · p. 2

  • Method

    UniCache identifies task-specific cache segments, assigns compression policies through offline calibration, and coordinates attention-guided budgets and task-aware temporal scheduling.

    Source: §1 · p. 3

  • Results

    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.

    Source: §1 · p. 3

  • Takeaways & Limitations

    UniCache improves the practicality of scaling unified multimodal models to longer contexts.

    Source: §5 · p. 9

Abstract

from arXiv · show

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\times$ KV cache compression for understanding and editing and $2.5\times$ for generation with negligible quality loss, while increasing throughput by up to $1.78\times$ in long-context settings, significantly improving the practicality of scaling unified multimodal models to longer context.

1 INTRODUCTION

Unified multimodal models face growing KV-cache costs, while cache types and their importance vary across tasks and inference stages. UniCache assigns task- and type-specific compression policies and budgets; on BAGEL, it reports substantial compression with negligible quality loss and faster long-context throughput.

1 INTRODUCTION

  • Growing multimodal contexts increase KV-cache storage and access costs, constraining inference efficiency.KV-cache compression is presented as a practical way to improve efficiency and support longer contexts.Source: §1 · p. 2
  • A single global compression policy struggles because tasks use different cache types, types differ in sensitivity and attention importance, and cache importance changes across tasks and inference stages.In image editing, eviction tends to lose source-image details, while quantization relatively preserves those details but can fail to apply the requested edit.Source: §1 · p. 2
  • UniCache identifies task-specific cache segments, assigns compression policies through offline attention calibration, and uses attention-guided allocation and task-aware scheduling to compress segments independently and in parallel.The overview describes type-specific budgets and parallel online compression.Source: §1 · p. 3, Figure 3 · p. 3
  • 5× KV cache compression for understanding and editing and 2.5× for generation with negligible quality loss, plus up to 1.78× throughput in long-context settings.Across models and benchmarks, UniCache also better preserves quality than uniform compression under comparable storage budgets.Source: §1 · p. 3

2 RELATED WORK

Unified-model inference accelerators adapt to differences between task execution modes, while KV-cache heterogeneity across task modes and cache types remains underexplored. Existing cache-compression approaches largely target autoregressive models; UniCache addresses the joint task- and type-specific management gap.

2 RELATED WORK

  • Hyper-Bagel (Lu et al., 2025) and Flash-Unified (Ke et al., 2026) use task-specific acceleration strategies, but heterogeneous KV caches across task modes and types remain underexplored.The cited approaches include speculative decoding and diffusion distillation, as well as pruning, layer skipping, visual-token pruning, and diffusion-head caching.Source: §2 · p. 3
  • Autoregressive cache methods use token retention or sparse access, including attention-sink retention, accumulated token importance, and attention-guided selection or budget allocation (Xiao et al., 2024; Liu et al., 2023; Li et al., 2024; Cai et al., 2025).DiffKV differentiates memory management for keys and values.Source: §2 · pp. 3–4
  • These designs do not jointly account for task mode and cache type; UniCache assigns per-type compression strategies and dynamically allocates a shared budget by task and inference stage.Source: §2 · p. 4

3 UNICACHE

UniCache coordinates task- and type-aware cache compression through task-specific segmentation, calibrated policy assignment, and attention-guided budget management. It also schedules total retention over inference and executes type-specific compression in parallel.

3.1 TASK-AWARE CACHE SEGMENTATION

  • UniCache segments each task’s attention context into conditioning-cache candidates and protected objects, keeping boundary and current-latent KVs outside the shared compression budget.In BAGEL, understanding uses instruction and source-ViT KVs, generation uses instruction KV, and editing uses instruction, source-ViT, and source-VAE KVs.Source: §3.1 · pp. 4–5

3.2 OFFLINE POLICY ASSIGNMENT

  • Offline calibration assigns a compression policy to each task–type pair using attention sparsity; policy assignments stay fixed while compression strength changes with online budget allocation.Figure 5 samples task–KV-type subsets comparing eviction and quantization at the same compression ratio.Source: §3.2 · p. 5, Figure 5 · p. 5

3.3 ATTENTION-GUIDED DYNAMIC BUDGET ALLOCATION

  • Importance statistics aggregate retained-key attention across chunks of layers and inference steps, then guide budget allocation for the next chunk.For generation and editing, initial statistics come from the first full-KV denoising step; for understanding, they come from final prefill queries.Source: §3.3 · pp. 5–6
  • Attention-guided allocation divides the shared storage budget across candidate cache types according to attention mass, with allocations capped by each type’s previous allocation.When a type reaches its cap, remaining budget is redistributed among uncapped types with positive attention mass; some may remain unused due to irreversible eviction.Source: §3.3 · p. 6
  • UniCache adjusts the total retention budget by task and inference stage: understanding uses a fixed task-level ratio, while generation and editing adapt it as conditioning attention declines.Understanding attention fluctuates across steps; generation and editing show declining conditioning attention at type-dependent rates.Source: §3.3 · p. 6, Figure 6 · p. 6

3.4 TYPE-AWARE PARALLEL COMPRESSION

  • Each assigned policy compresses only its own segments with independent runtime state, enabling parallel execution before compressed candidates and protected segments are reassembled in original token order.Eviction policies set the number of retained tokens; quantization policies set numerical precision and grouping granularity.Source: §3.4 · pp. 6–7

4 EMPIRICAL RESULTS

UniCache preserves balanced quality across understanding, generation, and editing, with ablations supporting its component choices and long-context tests showing efficiency gains. The evaluation covers two unified multimodal models and compares against full-KV and compressed baselines.

4.1 EXPERIMENTAL SETTING

  • The evaluation covers BAGEL and SenseNova-U1, using understanding, text-to-image generation, and image-editing benchmarks with official protocols.Compressed methods are compared with the full-KV model at approximately 80% memory reduction for understanding and editing and 60% for text-to-image generation.Source: §4.1 · p. 7

4.2 MAIN RESULTS

  • UniCache performs comparably to the original model on MME perception and cognition, while retaining nearly all original performance in generation and editing.For text-to-image generation, it outperforms existing methods in nearly every subcategory; existing editing methods tend to lose either structural details or semantic understanding and instruction following.Source: §4.2 · p. 7

4.3 ABLATION STUDIES

  • Removing boundary and current-latent KV protection degrades text-to-image generation, while attention-guided allocation is especially important for understanding and editing.Task-aware scheduling shifts capacity toward early generation stages as conditional-KV attention declines, and removing it slightly degrades generation.Source: §4.3 · p. 8, Table 3 · p. 8
  • Attention-guided decay matches linear decay on GenEval and achieves the best results across all PIE-Bench metrics.The comparison holds other components and compression policies fixed; constant-budget and linear-decay schedules serve as alternatives.Source: §4.3 · p. 8, Table 4 · p. 8
  • UniCache is more robust than H2O and KIVI at high compression ratios, maintaining balanced performance across both PIE-Bench metrics.Figure 7 compares methods across benchmarks and compression ratios.Source: §4.3 · p. 9, Figure 7 · p. 9

4.4 END-TO-END EFFICIENCY

  • At approximately 20K KV tokens per attention branch, UniCache improves throughput by up to 1.78× while reducing persistent managed KV cache memory by about 80%.Measurements use BAGEL-7B-MoT on one NVIDIA A100 GPU; Figure 8 reports throughput and memory across context lengths.Source: §4.4 · p. 9, Figure 8 · p. 9

5 CONCLUSION

The conclusion presents UniCache as a training-free response to heterogeneous cache types and task-dependent dynamics in MoT-based unified multimodal models. It reports preserved quality across tasks alongside compression and long-context throughput gains.

5 CONCLUSION

  • UniCache combines offline type-aware policy assignment with attention-guided budget allocation and task-aware temporal scheduling.It coordinates different compression policies across cache segments.Source: §5 · p. 9
  • UniCache preserves quality across understanding, generation, and editing, achieving up to 5× KV-cache compression with negligible performance loss.The conclusion reports up to 1.78× long-context throughput.Source: §5 · p. 9
  • The conclusion states that these results improve the practicality of scaling unified multimodal models to longer contexts.Source: §5 · p. 9

A ATTENTION CONCENTRATION AND POLICY PREFERENCE EXPERIMENTS

The appendix describes a controlled experiment comparing eviction and quantization preferences across task–KV-type pairs. It defines relative quality degradation and an eviction-advantage measure for comparing the policies.

A ATTENTION CONCENTRATION AND POLICY PREFERENCE EXPERIMENTS

B DETAILED EXPERIMENTAL SETTINGS

The appendix records both target and realized compression ratios and describes how calibration chooses policies for task–type pairs. Calibration also sets the minimum retention ratio to match a target average compression ratio.

B DETAILED EXPERIMENTAL SETTINGS

C ATTENTION ANALYSIS OF SENSENOVA-U1

SenseNova-U1 uses a single native visual KV type for source images, and UniCache assigns it eviction based on concentrated attention.

C ATTENTION ANALYSIS OF SENSENOVA-U1

  • Attention over SenseNova-U1’s single source-native-visual KV type concentrates on a small subset of tokens, so UniCache assigns it eviction.Offline calibration adapts policy assignment to the model’s cache structure and attention distribution.Source: C ATTENTION ANALYSIS OF SENSENOVA-U1 · p. 12

D SAMPLED RESULTS GENERATED BY UNICACHE

UniCache-generated samples illustrate image generation at 60% KV cache compression and image editing at 80%.

D SAMPLED RESULTS GENERATED BY UNICACHE

E EFFICIENT GPU IMPLEMENTATION

UniCache includes a PyTorch reference implementation and a GPU engine that physically compacts or packs KV entries to reduce persistent storage and attention overhead.

E EFFICIENT GPU IMPLEMENTATION

  • The PyTorch reference implementation uses attention masks and fake quantization, while the GPU engine physically compacts or packs KV entries.The GPU engine reduces persistent cache storage and attention overhead.Source: E EFFICIENT GPU IMPLEMENTATION · pp. 12–13
  • Figure 9 depicts layer-8, token-level attention probabilities averaged over heads and queries for one example, with tokens in original order.Source: Figure 9 · p. 13
  • The GPU engine stores retained entries and logical indices for eviction, packs quantized segments with metadata and a BFloat16 residual tail, and leaves protected segments unchanged.Source: E EFFICIENT GPU IMPLEMENTATION · p. 13
  • Batched GQA selection, heterogeneous-segment attention, chunk-level allocation, cached execution plans, and batched attention/CFG reduce runtime overhead.For separate attention chunks, log-sum-exp statistics preserve global normalization; Triton dequantizes packed KVs for FlashAttention.Source: E EFFICIENT GPU IMPLEMENTATION · p. 13
  • UniCache-optimized BAGEL produces image-generation samples at 60% compression, preserving generative performance with high-quality textures and correct object placements; editing samples follow instructions while preserving unrelated visual content.Source: Figure 10 · p. 14, Figure 11 · p. 15
  • Attention-guided allocation adds 0.45%–1.7% overhead depending on chunk size, small relative to the system’s throughput and memory savings.The engine targets BF16 inference on NVIDIA GPUs.Source: E EFFICIENT GPU IMPLEMENTATION · p. 14

Related papers

Loading 2609.32831v1…