Source-linked AI summary
HySparse: A Hybrid Sparse Attention Architecture with Oracle Token Selection and KV Cache Sharing
Yizhao Gao, Jianyu Wei, Qihao Zhang, Yu Cheng, Shimao Chen, Zhengju Tang, Zihan Jiang, Yifan Song, Hailin Zhang, Liang Zhao, Bo Yang, Gang Wang, Shijie Cao, Fuli Luo
TL;DR
Long-context self-attention faces quadratic cost, while existing sparse methods often depend on approximate token-selection proxies or retain full KV caches. HySparse uses preceding full-attention layers as token-selection oracles and shares their KV caches with sparse layers; it consistently outperforms full attention and hybrid SWA baselines, including with nearly 10× lower KV-cache storage in an 80B MoE setting.
Problem
Long-context attention is costly because standard self-attention scales quadratically, while sparse methods may rely on approximate token-selection proxies and retain full KV caches.
Method
HySparse interleaves full and sparse attention layers, reusing preceding full-layer token selections and KV caches for oracle-guided sparse attention.
Results
Across tasks and context lengths, HySparse consistently outperforms full attention and hybrid SWA baselines; in an 80B MoE model, 5 of 49 layers use full attention, reducing KV-cache storage by nearly 10×.
Takeaways & Limitations
HySparse substantially reduces full-attention usage while preserving strong long-context modeling and improving computation and memory efficiency.
Takeaways & Limitations
Completely eliminating full attention remains challenging because efficient sparse methods still rely on expensive global-computation or indexing components.
Abstract
from arXiv · showhide
This work introduces Hybrid Sparse Attention (HySparse), a new architecture that interleaves each full attention layer with several sparse attention layers. While conceptually simple, HySparse strategically derives each sparse layer's token selection and KV caches directly from the preceding full attention layer. This architecture resolves two fundamental limitations of prior sparse attention methods. First, conventional approaches typically rely on additional proxies to predict token importance, introducing extra complexity and potentially suboptimal performance. In contrast, HySparse uses the full attention layer as a precise oracle to identify important tokens. Second, existing sparse attention designs often reduce computation without saving KV cache. HySparse enables sparse attention layers to reuse the full attention KV cache, thereby reducing both computation and memory. We evaluate HySparse on both 7B dense and 80B MoE models. Across all settings, HySparse consistently outperforms both full attention and hybrid SWA baselines. Notably, in the 80B MoE model with 49 total layers, only 5 layers employ full attention, yet HySparse achieves substantial performance gains while reducing KV cache storage by nearly 10x.
1 Introduction
HySparse addresses sparse attention’s proxy-selection and KV-cache limitations by deriving token indices and caches from preceding full attention layers. Across dense and MoE settings, it consistently outperforms full attention and hybrid SWA baselines while reducing memory pressure.
- Quadratic self-attention makes computation increasingly costly as context length grows.
- Existing sparse methods use approximate proxies for token selection, which may miss true importance in long, evolving contexts.
- Dynamic sparse attention reduces computation but typically retains the full KV cache, limiting serving throughput and batch size.
- HySparse interleaves one full attention layer with multiple sparse layers that reuse oracle-guided token indices and KV caches from the preceding full layer.This removes auxiliary proxy modules and avoids additional per-layer KV overhead.
- Across tasks and context lengths, HySparse outperforms full attention and hybrid SWA baselines without additional KV-cache cost relative to hybrid SWA.In the 80B MoE model, 5 of 49 layers use full attention, yielding nearly 10× KV-cache reduction.
2 Background & Motivation
Prior efficient-attention work uses hybrid mechanisms, stable salient-token patterns, or cross-layer KV reuse to reduce long-context costs. These approaches motivate combining full-layer selection with cache sharing while retaining trainable sparse attention.
- Training-free sparse attention uses fixed patterns or heuristics, but inference-only sparsity can create training–inference mismatch and error accumulation.
- Hybrid attention combines local or linear mechanisms with softmax attention to reduce quadratic compute and KV-cache costs.
- Salient tokens tend to remain relatively stable across consecutive transformer layers, enabling later layers to reuse full-attention selections.
- Cross-layer KV sharing reuses keys and values from one layer in subsequent layers, substantially reducing cache memory with little to no reported accuracy degradation.
3 Methodology
HySparse builds repeated blocks around one full attention layer followed by multiple sparse layers. Full attention emits block-level importance scores and KV caches, while sparse layers reuse selected blocks and fuse global sparse retrieval with local SWA.
- Each hybrid block contains one full attention layer followed by N consecutive sparse attention layers.Both sparse token indices and sparse-layer KV caches derive from the preceding full layer.
- The full layer produces block-wise attention scores used to derive TopK block indices for subsequent sparse layers.
- HySparse materializes block-level maximum scores instead of the full attention matrix for token-block selection.The modified FlashAttention kernel reuses online-softmax intermediates to produce these scores with negligible overhead.
- With default k=1024 and block size B=64, selecting k tokens corresponds to selecting k/B TopK blocks.
- Sparse layers combine a shared-KV block-sparse branch with a local SWA branch using its own lightweight cache, then fuse both outputs through sigmoid gating.The SWA window size is 128, and experiments identify its independent cache as important for expressivity.
4 Experiments
Experiments evaluate HySparse against Full-Attn and Hybrid SWA on 7B dense and 80B MoE models, general benchmarks, long-context tasks, and architecture ablations.
- Experimental setup: The experiments compare Full-Attn, Hybrid SWA, and HySparse using 7B dense and 80B-A3B MoE Transformer configurations.The 7B model has 36 layers, while the 80B MoE model has 49 layers.
- General benchmarks: HySparse outperforms both Full-Attn and Hybrid SWA across nearly all 80B MoE benchmarks with only five full-attention layers.The exceptions versus Full-Attn are MMLU-Pro, DROP, and ARC-C, where HySparse is slightly lower.
- Long-context benchmarks: At 32k context, HySparse scores 87.4 versus 82.1 for Full-Attn in the 80B MoE model, while Hybrid SWA scores 69.5.For 7B models, HySparse reaches overall scores of 94.1 at 16k and 89.3 at 32k.
- Ablation study: Sharing KV caches only for SA improves MMLU-Pro from 23.2 to 29.0 (+5.8), whereas sharing them for both SA and SWA substantially degrades accuracy.The results support an independent SWA KV cache and cross-layer sharing for the sparse-attention branch.
5 Discussion & Future Works
HySparse suggests reducing the expensive full-attention component while addressing GPU KV-cache usage. It also points to a serving strategy that offloads full-attention caches and keeps selected sparse KV data on GPU.
- Discussion: HySparse remains part of a broader efficient-attention trend in which completely eliminating full-attention components remains challenging.The paper emphasizes balancing expensive global computation against cheaper local or sparse computation and GPU memory usage.
- Discussion: HySparse uses oracle token selection and cross-layer KV sharing to reduce the ratio of expensive global computation to sparse computation while preserving long-context capabilities.
- Future Works: Offloading full-attention KV caches to external memory while retaining selected sparse KV on GPU could reduce GPU cache footprint and support larger serving batch sizes.The proposed strategy pre-fetches the full cache before computation and keeps persistent selected/sparse KV for subsequent sparse layers.
6 Conclusion
HySparse interleaves full and sparse attention layers, deriving token selections and KV caches from preceding full-attention layers. This removes proxy-based selection and additional sparse-layer memory overhead while allowing fewer full-attention layers without compromising modeling capabilities.
- Conclusion: HySparse interleaves each full-attention layer with multiple sparse-attention layers.
- Conclusion: HySparse derives important token selections and KV caches from preceding full-attention layers, eliminating proxy-based selection and additional memory overhead.
- Conclusion: HySparse substantially reduces the number of full-attention layers in hybrid models without compromising modeling capabilities.