HyLRA: Hybrid Layer Reuse Attention for Efficient Long-Context Inference

TL;DR

This paper introduces HyLRA, a novel framework driven by layer-wise sparsity profiling that allows LLMs to restrict computation to the most critical tokens, effectively overcoming the quadratic bottleneck of dense attention.

Abstract

Long-context inference in Large Language Models (LLMs) is bottlenecked by the quadratic computation complexity of attention and the substantial memory footprint of Key-Value (KV) caches. While existing sparse attention mechanisms attempt to mitigate this by exploiting inherent sparsity, they often rely on rigid patterns or aggressive pruning, failing to achieve an optimal balance between efficiency and accuracy. In this paper, we introduce {\bf HyLRA} ({\bf Hy}brid {\bf L}ayer {\bf R}euse {\bf A}ttention), a novel framework driven by layer-wise sparsity profiling. Our empirical analysis uncovers a dual characteristic in attention mechanics: \textit{intra-layer sensitivity}, where specific layers necessitate full attention to prevent feature distortion, and \textit{inter-layer similarity}, where consecutive layers share substantial critical tokens. Based on these observations, HyLRA employs an offline dynamic programming approach to derive an optimal layer-wise policy. This hybrid strategy retains full attention for sensitive layers to ensure robustness, while enabling tolerant layers to bypass quadratic calculations by directly reusing top-$k$ indices from preceding layers. This approach allows LLMs to restrict computation to the most critical tokens, effectively overcoming the quadratic bottleneck of dense attention. Extensive evaluations demonstrate that HyLRA improves inference throughput by 6\%--46\% while maintaining comparable performance (with $<1\%$ accuracy degradation), consistently outperforming state-of-the-art sparse attention methods. HyLRA is open source at \href{https://anonymous.4open.science/r/unified-cache-management-CF80/}{\texttt{/r/unified-cache-management-CF80/}}

Figures (3)

• Figure 1: Sensitivity analysis of each layer in Qwen3-32B. We measure the sensitivity of the $l$-th layer by applying Top-$k$ (k = 2,048 tokens) sparse attention to it, while maintaining full attention for all preceding layers ($0$ to $l-1$). The impact is evaluated by measuring the deviation in the input to layer $l+1$. (a) Layer 0 exhibits significant sensitivity (high KL divergence with RNMSE = 0.25). Layer 7 demonstrates low sensitivity (low KL divergence with RNMSE = 0.016).
• Figure 2: Visualization of layer-wise attention similarity. We calculate the overlap ratio of the top-$k$ ($k$ = 2,048) tokens with the highest attention scores between any pair of layers in DeepSeek-R1. The heatmap shows the pairwise overlap ratio, where the x-axis and y-axis represent layer indices. The dark regions along the diagonal indicate that consecutive layers tend to share a large number of critical tokens, suggesting a high degree of local consistency in attention distributions.
• Figure 3: Visualization of offline policy optimization using a toy example with 10 layers. The grid maps the current layer (y-axis) against the reuse source (x-axis). The cyan path traces the optimal policy: diagonal steps to dark gray cells (e.g., layers 0, 3, 6) denote "Reset Jumps" (full attention refresh), while vertical steps to light gray cells denote "Index Reuse" (inheriting indices). This trajectory globally minimizes full attention frequency while maximizing cumulative similarity to ensure accuracy.