The proliferation of long-context large language models (LLMs) exposes a key bottleneck: the rapidly expanding key-value cache during decoding, which imposes heavy memory and latency costs. While recent approaches attempt to alleviate this by sharing a single set of crucial tokens across layers, such coarse-grained sharing undermines model performance by neglecting the functional diversity of attention heads. To address this, we propose LycheeDecode, an efficient decoding method centered on a fine-grained hybrid-head attention mechanism that employs a hardware-efficient top-k selection strategy. Specifically, the novel HardKuma-based mechanism partitions attention heads into a small subset of retrieval heads that dynamically identify crucial tokens and a majority of sparse heads that reuse them for efficient computation. Through extensive experiments on leading models like Llama3 and Qwen3 across diverse benchmarks for long-context understanding (e.g., LongBench, RULER) and complex reasoning (e.g., AIME24, OlympiadBench), we demonstrate that LycheeDecode achieves generative quality comparable to, and at times surpassing even the full-attention baseline. Crucially, this is accomplished with up to a 2.7x speedup at a 128K context length. By preserving the functional diversity of attention heads, our fine-grained strategy overcomes the performance bottlenecks of existing methods, providing a powerful and validated pathway to both efficient and high-quality long-context LLM inference.

翻译：长上下文大语言模型的广泛应用揭示了一个关键瓶颈：解码过程中快速膨胀的键值缓存带来了沉重的内存与延迟开销。现有方法尝试通过在层间共享一组关键令牌来缓解此问题，但这种粗粒度的共享忽视了注意力头的功能多样性，从而损害了模型性能。为此，我们提出LycheeDecode，一种高效的解码方法，其核心是一种细粒度的混合注意力头机制，并采用了硬件高效的前k项选择策略。具体而言，新颖的基于HardKuma的机制将注意力头划分为两个部分：一小部分检索头动态识别关键令牌，以及占多数的稀疏头复用这些令牌以进行高效计算。通过在Llama3、Qwen3等主流模型上，针对长上下文理解与复杂推理的多样化基准测试进行广泛实验，我们证明LycheeDecode在生成质量上可媲美甚至有时超越全注意力基线。关键在于，该方法在128K上下文长度下实现了高达2.7倍的加速。通过保留注意力头的功能多样性，我们的细粒度策略克服了现有方法的性能瓶颈，为高效且高质量的长上下文大语言模型推理提供了一条强有力且经过验证的路径。