LycheeDecode: Aceleração da Inferência em LLMs de Contexto Longo via Decodação Esparsa de Cabeças Híbridas

LycheeDecode: Accelerating Long-Context LLM Inference via Hybrid-Head Sparse Decoding

February 4, 2026

Autores: Gang Lin, Dongfang Li, Zhuoen Chen, Yukun Shi, Xuhui Chen, Baotian Hu, Min Zhang

cs.AI

Resumo

A proliferação de modelos de linguagem de grande escala (LLMs) de contexto longo expõe um gargalo crítico: a rápida expansão da cache de chave-valor durante a decodificação, que impõe custos significativos de memória e latência. Embora abordagens recentes tentem mitigar este problema compartilhando um único conjunto de tokens cruciais entre as camadas, esse compartilhamento de granularidade grossa prejudica o desempenho do modelo ao negligenciar a diversidade funcional das cabeças de atenção. Para resolver isso, propomos o LycheeDecode, um método de decodificação eficiente centrado em um mecanismo de atenção híbrido de granularidade fina que emprega uma estratégia de seleção top-k eficiente em hardware. Especificamente, o novo mecanismo baseado em HardKuma particiona as cabeças de atenção em um pequeno subconjunto de cabeças de recuperação que identificam dinamicamente tokens cruciais e uma maioria de cabeças esparsas que os reutilizam para computação eficiente. Por meio de extensivos experimentos em modelos líderes como Llama3 e Qwen3 em diversos benchmarks para compreensão de contexto longo (por exemplo, LongBench, RULER) e raciocínio complexo (por exemplo, AIME24, OlympiadBench), demonstramos que o LycheeDecode alcança uma qualidade gerativa comparável e, por vezes, até superior à linha de base de atenção completa. Crucialmente, isso é alcançado com uma aceleração de até 2,7x em um contexto de 128K tokens. Ao preservar a diversidade funcional das cabeças de atenção, nossa estratégia de granularidade fina supera os gargalos de desempenho dos métodos existentes, fornecendo um caminho validado e robusto para uma inferência eficiente e de alta qualidade em LLMs de contexto longo.

English

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.