Source-linked AI summary
MILO: Efficient Many-shot In-Context Learning with Block-wise Low-rank Compression
Youpeng Zhao, Tian Tan, Liqian Peng, Jun Wang, Alec Go
TL;DR
Many-shot ICL shifts inference pressure to the memory required by its growing KV cache. MILO compresses that cache block-wise and allocates rank by entropy; on Qwen2.5 models, it reduces memory footprint and improves speed while preserving accuracy.
Problem
Many-shot ICL creates inference memory and compute demands because long prompts require storing a KV cache that grows linearly with sequence length.
Source: §1 · p. 1, §4.1 · p. 5Method
MILO compresses the KV cache by blocks and dynamically assigns rank budgets based on each block's entropy.
Source: §1 · p. 2, §4.3 · p. 6Results
50% reduction in KV cache memory footprint and up to 1.8× speedup upon existing baseline methods, while preserving model accuracy.
Source: §6 · p. 11Takeaways & Limitations
MILO provides a practical path toward efficient many-shot ICL inference by combining block-wise low-rank compression with customized fused Triton kernels.
Source: §6 · p. 11
Abstract
from arXiv · showhide
Many-shot in-context learning (ICL) enables large language models (LLMs) to adapt to complex tasks by conditioning on thousands of demonstration examples, but this paradigm shifts the inference efficiency bottleneck to the key-value (KV) cache memory. Due to the linear scaling behavior of the KV cache, storing these intermediate tensors has become a paramount challenge for both online serving and on-device deployment. To address this issue, we propose a novel compression framework, termed MILO, that exploits the low-rank redundancy inherent in many-shot contexts. Specifically, MILO features a block-wise low-rank compression strategy that compresses the KV cache at the block granularity, where each block contains multiple many-shot examples. Furthermore, to handle the heterogeneous context density across different blocks, MILO dynamically allocates rank budgets based on the information entropy, preserving the fidelity of critical blocks while aggressively compressing redundant ones. Experimental results on Qwen2.5 models demonstrate that our method achieves up to 50% reduction in KV cache memory and 1.8x throughput improvement, with negligible performance degradation on classification and reasoning benchmarks, significantly outperforming prior baselines.
1. Introduction
Many-shot ICL can improve task performance with hundreds or thousands of demonstrations, but the resulting KV cache creates memory and inference-efficiency challenges. MILO compresses that cache block-wise and adapts rank budgets to block entropy.
1. Introduction
- Many-shot ICL conditions LLMs on hundreds or thousands of examples, while long-context inference raises memory and compute requirements.The KV cache for 90k examples requires an additional 11.1 GB for 8B dense models.Source: §1 · pp. 1–2
- Prior approaches target compute or KV-cache reuse, while storage and memory efficiency for cached many-shot examples remain unaddressed.The cited examples include dynamic sparse attention and semantic-aware KV-cache reuse.Source: §1 · pp. 1–2
- For pre-encoded many-shot ICL, MILO exploits the KV cache’s low-rank redundancy by storing it in low-dimensional space, achieving up to 50% storage cost reduction and 1.8× speedup against prior works with minimal accuracy degradation.Source: §1 · p. 2
- MILO compresses KV caches in blocks containing multiple examples to balance representation quality, memory efficiency, and decompression cost.Global compression risks losing fine-grained information and adding decompression overhead; example-level compression can lose global semantics.Source: §1 · p. 2
- MILO assigns higher rank budgets to higher-entropy blocks and compresses redundant blocks more aggressively, with customized system optimizations.The method is designed to preserve critical-block fidelity while reducing aggregate cache memory.Source: §1 · p. 2
2. Background
The background describes how expanded context windows enable many-shot ICL while making cached intermediate states a persistent memory challenge. Existing approaches address inference through attention, reuse, low-rank compression, sparsity, or quantization, but do not resolve the many-shot cache-storage gap.
2. Background
- Many-shot ICL extends demonstration-based task adaptation from few-shot prompts to contexts with hundreds or thousands of examples as context windows grow.Prior context windows were often 2k or 4k tokens; long-context models can process up to 10M tokens.Source: §2 · p. 3
- Inference remains constrained by quadratic self-attention and a KV cache that grows linearly with prompt length, leaving cached states for thousands of examples as an open memory problem.DBSA targets compute latency with dynamic sparse attention, while Adapshot uses context-budget allocation and KV-cache reuse.Source: §2 · p. 3
- Prior low-rank KV-cache methods focus on limited context windows or data-dependent compression, whereas this work targets low-rank structure across many-shot examples.The cited methods include projection, inter-layer compression, and head-level compression approaches.Source: §2 · p. 3
- Sparse attention can evict less important tokens, while quantization reduces numerical precision; the paper presents low-rank representation compression as orthogonal to these approaches.The passage states that low-rank compression can be combined with quantization and sparse attention.Source: §2 · p. 3
- Figure 2 samples three MATHQA blocks of 64 examples and reports effective rank as the square of singular values, with mean and standard deviation across attention layers.The caption identifies key and value caches across blocks as the quantities observed.Source: Figure 2 · p. 4
3. Observation
The observations identify low-rank structure in many-shot KV caches, but show that it varies by cache component, shot count, block, and layer. These properties motivate exploiting redundancy while accounting for heterogeneous representation needs.
3. Observation
- As shot count increases, KV storage grows while low-rank redundancy also increases.Figure 1 visualizes the effective rank of pre-RoPE keys and values over different shot counts.Source: Figure 1 · p. 4
- The key cache exhibits a stronger low-rank nature than the value cache in many-shot ICL settings.The passage says this matches prior findings that pre-RoPE keys are better suited to low-rank compression.Source: §3 · p. 4
- Low-rank behavior is heterogeneous across blocks and layers, so different blocks have different spectral complexity and representation needs.The passage says these observations motivate exploiting low-rank structure in many-shot contexts.Source: §3 · p. 4
4. Methodology
MILO addresses many-shot ICL’s KV-cache memory bottleneck with block-wise low-rank compression, entropy-driven rank allocation, and system optimizations.
4.1. Background
- The KV cache grows linearly with sequence length and can exceed GPU HBM as shots increase, motivating compression that preserves model accuracy.MILO’s stated objective is to minimize memory use while preserving accuracy.Source: §4.1 · p. 5
4.2. Block-wise Low-Rank Compression
- BLC partitions the KV cache into multi-example blocks and compresses each independently, balancing global compression’s detail loss against accumulated per-example information loss.MILO applies low-rank approximation separately to each block’s keys and values.Source: §4.2 · p. 5, Figure 3 · p. 6
4.3. Dynamic Rank Allocation
- Under a fixed total rank budget, MILO assigns block-specific ranks using entropy, favoring slowly decaying spectra and allocating rank by marginal information gain.The allocation begins at a minimum rank and greedily assigns remaining rank units to blocks with the largest marginal gain.Source: §4.3 · pp. 6–7
4.4. MILO Workflow
- During encoding, MILO partitions the context, analyzes block entropy, and stores low-rank caches; at inference, it retrieves selected blocks, reconstructs them in parallel, and applies sparse attention.The workflow description assumes equal key and value ranks for simplicity, though they may differ in practice.Source: §4.4 · p. 7
4.5. System Optimizations
- System optimizations fuse retrieval, scaling, and reconstruction in Triton kernels, parallelize independent blocks with CUDA streams, and use CUDA Graphs to reduce repeated launch overhead.The fused kernel reduces intermediate HBM traffic and kernel-launch overhead; streams synchronize before sparse attention.Source: §4.5 · pp. 7–8
5. Evaluation
The evaluation compares MILO with full caching and low-rank baselines on Qwen2.5 models and many-shot ICL tasks. It reports accuracy, end-to-end throughput, and analyses of compression, block size, kernels, rank allocation, and quantization.
5.1. Experimental Setup
- Experiments use Qwen2.5-3B and Qwen2.5-7B on six NLU and reasoning datasets, comparing full caching, ASVD, and Palu on a single L4 GPU.Models run in FP16; system results include end-to-end memory and throughput.Source: §5.1 · p. 8
5.2. Main Results
- At a 50% rank budget, MILO improves accuracy by up to 4.5% over prior low-rank methods and averages 1.6× speedup over full KV offloading.On Qwen2.5-7B, it also sustains up to 3.2× more throughput for 512-shot inputs.Source: §5.2 · pp. 8–9, Table 1 · p. 9, Figure 4 · p. 9
5.3. Performance Analysis
- MILO maintains higher accuracy than ASVD and Palu across compression ratios; accuracy and efficiency initially improve with block size but decline beyond 32 examples per block.Larger blocks expose more cross-example redundancy, while overly large blocks reduce compression and retrieval granularity and increase reconstruction overhead.Source: §5.3 · pp. 9–10, Figure 5 · p. 10
- Fused Triton kernels reduce latency by approximately 2.5×, rising to approximately 3.7× over PyTorch when combined with CUDA streams and CUDA graphs.The fused implementation primarily reduces latency by eliminating intermediate HBM traffic.Source: Figure 5 · p. 10, §5.3 · p. 10
- Entropy-based rank allocation achieves the highest accuracy against uniform and random allocation under the same total rank budget, though its advantage narrows as shot count increases.The comparison uses Qwen2.5-7B.Source: Figure 5 · p. 10, §5.3 · p. 10
- MILO combines with asymmetric INT8 and KIVI quantization to provide complementary memory savings while preserving model accuracy.Source: §5.3 · p. 11
6. Conclusion
MILO compresses KV caches block-wise, allocating rank budgets dynamically to handle heterogeneous redundancy. The conclusion reports reduced memory with preserved accuracy and faster performance with customized fused Triton kernels.
6. Conclusion
- MILO is a block-wise low-rank KV-cache compression framework that exploits heterogeneous redundancy across demonstration blocks and dynamically allocates a global rank budget.Source: §6 · p. 11
- MILO reduces KV-cache memory footprint by 50% while preserving model accuracy.Source: §6 · p. 11
- Customized fused Triton kernels help MILO achieve up to 1.8× speedup over existing baseline methods.Source: §6 · p. 11