LycheeDecode: Aceleración de la Inferencia en LLM de Contexto Largo mediante Decodificación Dispersa de Cabezas 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

Resumen

La proliferación de modelos de lenguaje de gran tamaño (LLM) de contexto largo expone un cuello de botella clave: la rápida expansión de la caché clave-valor durante la decodificación, que impone elevados costes de memoria y latencia. Si bien enfoques recientes intentan mitigar esto compartiendo un único conjunto de tokens cruciales entre capas, dicho compartimiento de granularidad gruesa socava el rendimiento del modelo al ignorar la diversidad funcional de las cabezas de atención. Para abordar este problema, proponemos LycheeDecode, un método de decodificación eficiente centrado en un mecanismo de atención híbrido de granularidad fina que emplea una estrategia de selección top-k eficiente en hardware. Específicamente, el novedoso mecanismo basado en HardKuma divide las cabezas de atención en un pequeño subconjunto de cabezas de recuperación que identifican dinámicamente tokens cruciales y una mayoría de cabezas dispersas que los reutilizan para un cálculo eficiente. Mediante experimentos exhaustivos en modelos líderes como Llama3 y Qwen3 en diversos benchmarks para comprensión de contexto largo (por ejemplo, LongBench, RULER) y razonamiento complejo (por ejemplo, AIME24, OlympiadBench), demostramos que LycheeDecode logra una calidad generativa comparable y, en ocasiones, incluso superior a la línea base de atención completa. Crucialmente, esto se consigue con una aceleración de hasta 2.7x para una longitud de contexto de 128K. Al preservar la diversidad funcional de las cabezas de atención, nuestra estrategia de granularidad fina supera los cuellos de botella de rendimiento de los métodos existentes, proporcionando una vía potente y validada para una inferencia eficiente y de alta calidad en LLM de contexto largo.

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.