Source-linked AI summary
You Need an Encoder for Native Position-Independent Caching
Shiju Zhao, Junhao Hu, Jiaqi Zheng, Guihai Chen
TL;DR
Prefix-based KV caching is inefficient when retrieved contexts arrive in arbitrary orders, while existing PIC methods can reduce accuracy. The paper introduces native PIC through a PIC-aware encoder and COMB caching system, achieving up to 94% lower TTFT and 3× higher throughput with comparable accuracy. COMB remains a plug-in that can be enabled or disabled on demand.
Problem
Prefix caching requires exact prefix matches, while existing PIC approaches can incur accuracy degradation when reusing static contexts independently of position.
Method
COMB reintroduces an encoder into decoder-only LLMs, freezes the decoder, trains the encoder for PIC-compatible generation, and integrates PIC management with existing inference frameworks.
Results
Up to 94% lower TTFT and 3× higher throughput were achieved with cache hits while matching or exceeding prefix-based attention accuracy.
Takeaways & Limitations
COMB provides native PIC with high accuracy, improved inference efficiency, and on-demand compatibility with standard prefix caching.
Takeaways & Limitations
COMB’s complexity analysis assumes an encoder depth LE < LD for lower KV memory, while larger LE increases expressivity and potentially accuracy.
Abstract
from arXiv · showhide
The Key-Value (KV) cache of Large Language Models (LLMs) is prefix-based, making it highly inefficient for processing contexts retrieved in arbitrary order. Position-Independent Caching (PIC) has been proposed to enable KV reuse without positional constraints; however, existing approaches often incur substantial accuracy degradation, limiting their practical adoption. To address this issue, we propose native PIC by reintroducing the encoder to prevalent decoder-only LLMs and explicitly training it to support PIC. We further develop COMB, a PIC-aware caching system that integrates seamlessly with existing inference frameworks. Experimental results show that COMB reduces Time-to-First-Token (TTFT) by 51-94% and increases throughput by 3$\times$ with comparable accuracy. Furthermore, the quality improvement when using DeepSeek-V2-Lite-Chat demonstrates the applicability of COMB to other types of decoder-only LLMs. Our code is available at https://github.com/shijuzhao/Comb.
1 Introduction
Prefix caching reuses KV vectors only for exact common prefixes, limiting reuse when static content appears after varying prefixes. COMB addresses this with native PIC, a trained encoder plug-in and integrated caching system, while reporting strong accuracy and efficiency results.
- Motivation: Prefix caching reuses KV vectors of the longest common prefix but limits reuse when static chunks follow varying prefixes.This affects settings such as few-shot learning and retrieval-augmented generation.
- Position-Independent Caching: PIC compiles static chunks from position zero, then links their KV vectors in arbitrary order or position, but can degrade accuracy.Its main challenge is accurate recovery because it deviates from standard attention mechanisms.
- COMB: COMB reintroduces an encoder plug-in into decoder-only LLMs and trains only that encoder to generate PIC-compatible representations.The decoder remains unchanged, combining native PIC support with a removable component.
- COMB: COMB combines higher accuracy with flexibility by making its PIC-specific encoder removable without affecting standard decoding.This contrasts with post-training PIC’s reduced accuracy and training-aware PIC’s permanent model changes.
- Results: Up to 94% lower TTFT and 3× higher throughput were reported for cache hits while matching or exceeding prefix-based attention accuracy.The results were obtained on LongBench and COMB can revert to standard prefix caching when disabled.
2 Background
This background frames PIC as a position-independent alternative to prefix caching and places it among broader transformer inference stages and prior PIC methods. Existing approaches divide into post-training recomputation and training-aware adaptation.
- Transformer and Caching Primer: LLM generation has prefill and decode stages, with prefill computing prompt-token KV vectors and initiating decoding with the first output token.TTFT measures the time required to generate that first token.
- Position-Independent Caching: PIC reuses cached KVs for static prompt segments independently of their positions or surrounding context.The method follows a compilation-and-linking framework for cached context reuse.
- Prior PIC Approaches: Post-training PIC keeps the decoder fixed and recovers accuracy through selective recomputation during linking.EPIC recomputes initial chunk tokens, while CacheBlend recomputes tokens with the largest discrepancies.
- Prior PIC Approaches: Training-aware PIC explicitly incorporates PIC into model training to reduce or avoid link-stage recomputation.The supplied passage introduces this category but does not provide further details.
3 System Overview
COMB separates reusable contexts from the query, retrieves or builds position-independent caches, and passes those caches to an existing inference engine. The system therefore supports arbitrary-position context reuse through cache management around standard serving frameworks.
- Inputs and Cache Lookup: COMB accepts a question plus zero or more contexts intended for reuse at arbitrary positions.Users explicitly distinguish contexts from the query before serving.
- Inputs and Cache Lookup: For multiple contexts, COMB checks a hash table for existing Position-Independent Caches.This lookup determines which contexts can be reused immediately.
- Cache Generation: Contexts without caches are processed into KV caches, stored, and added to the hash table.The chunk processor performs these cache-generation and management steps.
- Inference Integration: COMB fetches PICaches and transfers them to inference engines such as HuggingFace transformers, vLLM, or SGLang.The question is then passed directly to the inference engine for response generation using those PICaches.
4 Model Design
COMB augments a frozen decoder-only LLM with a shallower encoder that independently processes documents into reusable PICaches, while decoder queries access them through cross-attention. This design reduces attention and KV-memory costs relative to standard decoder-only processing.
- Architecture: COMB disables its encoder and cross-attention layers to recover the original decoder-only LLM, while keeping decoder parameters frozen during training.The encoder is trained for PIC without permanently changing the underlying decoder behavior.
- Architecture: The encoder processes document tokens and produces PICaches, whereas the decoder handles queries and generated tokens.The encoder has fewer layers than the decoder, such as 8 interleaved encoder layers versus 32 decoder layers.
- Attention design: Decoder tokens attend to encoder-produced document KVs through cross-attention, whose output enters the decoder pathway through residual connections.Cross-attention mixes lquery decoder queries with ldoc document KVs at per-layer cost O(lquery ldoc).
- PICache construction: COMB stores cross-attention KVs because this is more efficient than storing encoder hidden states, and independently compiled document PICaches can later be concatenated.Multiple documents can be generated independently from position zero during compilation and reused across requests.
- Complexity analysis: COMB separates attention over documents and queries: encoder self-attention costs O(ldoc^2), decoder self-attention costs O(lquery^2), and cross-attention costs O(lquery ldoc).The architecture uses LE encoder layers and LD decoder layers, with Lcross cross-attention layers.
- Deployment implications: Independent document prefilling and existing chunked-prefill procedures make cold-start PIC compilation compatible with long-prompt inference.The paper states that cache misses typically do not introduce noticeable overhead beyond the baseline system.
- Complexity analysis: COMB stores document KVs across fewer encoder layers and query KVs across decoder layers, yielding lower KV memory than a decoder-only model when LE < LD.The paper sets LE = 8 to balance accuracy and computational resources.
5 Implementation
COMB uses frozen open-source decoder backbones with a trainable encoder and is implemented as a plug-in system atop existing inference frameworks. Training pairs static contexts and queries with supervised responses so the encoder learns PIC-aware representations.
- Model setup: COMB augments frozen decoder-only LLMs with a trainable encoder, using Llama-3.1-8B-Instruct and DeepSeek-V2-Lite-Chat as decoder backbones.The two backbones represent standard and Multi-head Latent Attention KV-cache designs, respectively.
- Training data and recipe: Training examples contain a static context D, query or instruction Q, and ground-truth response Y, with documents independently encoded into document-side KVs.The frozen decoder generates outputs conditioned on encoder KVs and query tokens.
- Training data and recipe: COMB optimizes token-level cross-entropy between decoder outputs and target sequence Y using teacher forcing.The training datasets are SQuAD, Natural-Instructions, XSum, and Super-Natural-Instructions.
- Training data and recipe: Llama-generated outputs supervise both models, substantially improving DeepSeek accuracy because its native generations are lower quality.All models are trained on four NVIDIA A100 80GB GPUs with tensor parallelism set to 4.
- System implementation: The system is built on Hugging Face Transformers and vLLM with four components: PIC manager, PIC allocator, Chunk processor, and Inference engine.The PIC manager reserves GPU memory, initializes the allocator, and launches the processing and inference components.
6 Evaluation
COMB is evaluated across LongBench tasks, model families, PIC baselines, accuracy, TTFT, memory use, and online serving load. It maintains strong accuracy while reducing latency and memory demands, with especially favorable results under cache hits and increasing request rates.
- Experimental Setup: COMB is compared with Prefix caching, CacheBlend, EPIC, and BlockAttention using TTFT, F1, and Rouge-L metrics.CacheBlend and EPIC are post-training PIC approaches, while BlockAttention is training-aware; CacheBlend recomputes 20% of tokens and EPIC recomputes 64 tokens per document.
- Accuracy under PIC: Across all datasets and both Llama and DeepSeek model families, COMB attains high accuracy among evaluated PIC approaches and can recover or exceed prefix-based accuracy.The authors caution that absolute accuracy comparisons are not fully controlled because methods differ in PIC awareness, trainable parameters, and training steps.
- Non-Intrusiveness: COMB remains non-intrusive: disabling its encoder restores the underlying decoder-only model’s prefix-caching behavior and accuracy without degradation.By contrast, BlockAttention can suffer severe accuracy degradation when inputs are not carefully chunked and processed with its intended block structure.
- TTFT under Cache Hits and Misses: Cache hits substantially reduce TTFT for PIC approaches, and COMB achieves the lowest TTFT through document-only encoder processing, query-only decoder self-attention, and cross-attention.PICaches are compiled once per document and reused across subsequent requests; even cache misses remain competitive because COMB’s encoder is shallower and aligns with chunked prefill.
- Online Latency and Throughput: Across the evaluated load range, COMB maintains the lowest TTFT and highest sustained throughput while saving KV memory by 75% for Llama-3.1-8B-Instruct and 78% for DeepSeek-V2-Lite.Its memory design stores encoder KVs for static document tokens and decoder KVs for shorter query sequences, freeing HBM for more concurrent requests.
7 Discussion
The discussion argues that native PIC benefits from reintroducing an encoder because retrieved contexts often arrive in arbitrary orders, where prefix caching is inefficient. COMB separates retrieved content from the decoder’s question and reasoning, supporting this use case.
- The Potential of PIC: The proposed encoder-decoder direction draws on the Transformer’s original architecture and validation in multimodal LLMs.The passage identifies Whisper and Mllama as examples of this paradigm’s prior validation.
- The Potential of PIC: Retrieval makes PIC essential because arbitrary permutations allow prefix caching to reuse only the first retrieved item.Subsequent retrieved items fail to reuse the KV cache under this pattern.
- The Potential of PIC: Native PIC uses an encoder to process arbitrarily ordered contexts while leaving questions and reasoning in the decoder.The design is motivated by encoder comprehension capabilities and the separation of retrieved content from generation.
- The Potential of PIC: The encoder can discard previously retrieved content and replace it with newly retrieved items when the agent needs new information.The question and the model’s reasoning remain in the decoder during this replacement.
8 Related Work
The related work places COMB within LLM serving optimization and context-caching research, distinguishing position-dependent caching from position-independent approaches and post-training from training-aware methods.
- Context Caching: COMB is presented as native position-independent caching within the broader design space of LLM serving and context caching.The related-work discussion separately surveys serving systems, context-caching categories, and training trajectories.
- LLM Serving Optimizations: LLM serving systems and schedulers include vLLM, SGLang, disaggregated prefill and decode, continuous batching, and speculative decoding.These works target serving efficiency rather than specifically defining PIC.
- Context Caching: Context caching comprises position-dependent methods, including prefix-based and relative position-dependent caching, and position-independent caching.Gemini and DeepSeek are cited as vendors incorporating explicit prefix-based caching features by mid-2024.
- Post-Training vs. Training-Aware Approaches: Post-training adjustments commonly precede training-aware integration, as illustrated by the progression from inference-time sparsity heuristics to Native Sparse Attention.The passage frames this as a recurring trajectory in LLM research.
9 Conclusion
The paper proposes native PIC by adding and training an encoder for decoder-only LLMs, and presents COMB as an integrated PIC-aware caching system. Experiments report improved accuracy and inference efficiency, while preserving optional plug-in use.
- Conclusion: COMB reintroduces an encoder into decoder-only LLMs and explicitly trains it for PIC-compatible generation.This is the paper’s proposed native PIC approach.
- Conclusion: COMB integrates PIC-aware caching with existing inference frameworks.The system is presented alongside the model architecture as part of the paper’s contribution.
- Conclusion: Up to 94% lower TTFT and up to 3× higher throughput are reported, alongside the highest accuracy under PIC.These are the conclusion’s headline efficiency and accuracy claims.
- Conclusion: COMB can be enabled or disabled on demand, reverting to standard prefix caching when PIC is unused.This preserves an optional deployment mode within the reported system design.