Shallow Prefill, Deep Decoding: Efficient Long-Context Inference via Layer-Asymmetric KV Visibility

Researchers introduce SPEED, a novel inference optimization technique for long-context language models that reduces computational cost by materializing key-value cache states only in lower layers during the prefill phase while maintaining full-depth processing during decoding. Testing on Llama-3.1-8B demonstrates 33% improvement in time-to-first-token, 22% improvement in tokens-per-second, and 25% reduction in KV memory with minimal quality degradation, suggesting that prompt tokens don't require persistent full-depth caching.

The research addresses a fundamental bottleneck in modern language model inference: the exponential cost of processing long prompts through transformer architectures. Current systems cache key-value states across all layers for every prompt token, creating substantial memory overhead and computational latency that scales with context length. SPEED's contribution lies in its asymmetric approach—rather than attempting to make upper-layer caching cheaper through pruning or approximation methods, the authors eliminate upper-layer prefill entirely while preserving full-depth processing for decode-phase tokens.

This work reflects a broader industry shift toward understanding transformer layer heterogeneity. Recent studies have demonstrated that different transformer layers perform distinct functions: lower layers handle prompt selection and semantic representation, while upper layers focus on token prediction. SPEED operationalizes this insight by aligning cache visibility with functional requirements, keeping only a minimal begin-of-sequence anchor in upper layers during prefill.

The practical implications are substantial for deployed systems. The reported 33% reduction in time-to-first-token addresses a critical user experience metric, while 25% memory savings enable longer context handling on resource-constrained infrastructure. The minimal 0.2-point benchmark degradation (51.2 vs 51.4) on instruction-following tasks suggests the approach generalizes across domains. This becomes particularly relevant for applications requiring extended context—retrieval-augmented generation, document analysis, and code understanding—where prompt processing dominates latency.

Future research should explore whether this layer-cutoff pattern holds across diverse model architectures and whether the approach combines effectively with other optimization techniques like speculative decoding or adaptive computation. The work suggests that efficient long-context inference needn't require fundamental architectural changes.

• →SPEED reduces time-to-first-token by 33% and active KV memory by 25% on 128K context lengths with minimal quality loss
• →Layer-asymmetric KV visibility proves effective because lower transformer layers handle prompt representation while upper layers focus on token prediction
• →Only 75% layer coverage for prefill tokens maintains baseline benchmark performance, indicating upper-layer prompt caching is unnecessary
• →The approach is orthogonal to other optimization techniques and could integrate with speculative decoding or quantization methods
• →Results suggest efficient long-context inference requires rethinking cache architecture rather than incremental improvements to existing methods