LycheeDecode : Accélération de l'inférence des LLM à contexte long par décodage parcimonieux à têtes hybrides

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

February 4, 2026

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

cs.AI

Résumé

La prolifération des modèles de langage à grand contexte (LLM) révèle un goulot d'étranglement majeur : l'expansion rapide du cache clé-valeur lors du décodage, qui entraîne des coûts élevés en mémoire et en latence. Si les approches récentes tentent de remédier à ce problème en partageant un ensemble unique de tokens cruciaux entre les couches, un partage aussi granulaire nuit aux performances du modèle en négligeant la diversité fonctionnelle des têtes d'attention. Pour résoudre cela, nous proposons LycheeDecode, une méthode de décodage efficace centrée sur un mécanisme d'attention hybride à granularité fine qui utilise une stratégie de sélection top-k optimisée pour le matériel. Concrètement, ce nouveau mécanisme basé sur HardKuma partitionne les têtes d'attention en un petit sous-ensemble de têtes de récupération qui identifient dynamiquement les tokens cruciaux et une majorité de têtes éparses qui les réutilisent pour un calcul efficace. Grâce à des expériences poussées sur des modèles leaders comme Llama3 et Qwen3, évalués sur divers benchmarks de compréhension de contexte long (LongBench, RULER) et de raisonnement complexe (AIME24, OlympiadBench), nous démontrons que LycheeDecode atteint une qualité générative comparable, voire supérieure dans certains cas, à la baseline utilisant l'attention complète. Fait crucial, ces résultats sont obtenus avec une accélération allant jusqu'à 2,7x pour un contexte de 128K tokens. En préservant la diversité fonctionnelle des têtes d'attention, notre stratégie granulaire surmonte les limitations des méthodes existantes, offrant ainsi une voie puissante et validée pour une inférence LLM à la fois efficace et de haute qualité sur de longs contextes.

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.