LycheeDecode: Ускорение вывода LLM с длинным контекстом с помощью гибридного разреженного декодирования заголовков

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

February 4, 2026

Авторы: Gang Lin, Dongfang Li, Zhuoen Chen, Yukun Shi, Xuhui Chen, Baotian Hu, Min Zhang

cs.AI

Аннотация

Распространение больших языковых моделей (LLM) с поддержкой длинного контекста выявило ключевую проблему: быстрорастущий кэш ключей-значений в процессе декодирования, который приводит к значительным затратам памяти и увеличению задержек. В то время как современные подходы пытаются решить эту проблему, используя общий набор ключевых токенов для всех слоев, такое грубое объединение снижает производительность модели, игнорируя функциональное разнообразие голов внимания. Для решения этой проблемы мы предлагаем LycheeDecode — эффективный метод декодирования, основанный на механизме гибридного внимания с детализированным разделением голов, который использует аппаратно-эффективную стратегию выбора top-k. В частности, новый механизм на основе HardKuma разделяет головы внимания на небольшое подмножество голов поиска, которые динамически идентифицируют ключевые токены, и большинство разреженных голов, которые повторно используют их для эффективных вычислений. В результате масштабных экспериментов с ведущими моделями, такими как Llama3 и Qwen3, на различных наборах данных для оценки понимания длинного контекста (например, LongBench, RULER) и сложных рассуждений (например, AIME24, OlympiadBench), мы демонстрируем, что LycheeDecode обеспечивает качество генерации, сопоставимое, а в некоторых случаях превосходящее даже базовый вариант с полным вниманием. При этом достигается ускорение до 2.7 раза при длине контекста 128K токенов. Сохраняя функциональное разнообразие голов внимания, наша детализированная стратегия преодолевает ограничения производительности существующих методов, предлагая проверенный путь к эффективному и качественному выводу LLM с длинным контекстом.

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.