LycheeDecode: 하이브리드 헤드 희소 디코딩을 통한 장문 컨텍스트 LLM 추론 가속화

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

장문 컨텍스트를 처리하는 대규모 언어 모델(LLM)의 확산은 중요한 병목 현상을 야기합니다. 바로 디코딩 과정에서 급격하게 증가하는 키-값 캐시로 인해 상당한 메모리 및 지연 시간이 발생합니다. 최근 연구에서는 레이어 간에 중요한 토큰 집합을 공유하여 이 문제를 완화하려는 시도가 있었지만, 이러한 거칠고 일반적인 공유 방식은 어텐션 헤드의 기능적 다양성을 무시하여 모델 성능을 저하시킵니다. 이러한 문제를 해결하기 위해, 우리는 하드웨어 효율적인 상위-k 선택 전략을 사용하는 세밀한 하이브리드 헤드 어텐션 메커니즘을 중심으로 하는 효율적인 디코딩 방법인 LycheeDecode를 제안합니다. 구체적으로, 새로운 HardKuma 기반 메커니즘은 어텐션 헤드를 중요한 토큰을 동적으로 식별하는 소수의 검색 헤드와, 효율적인 계산을 위해 이를 재사용하는 다수의 희소 헤드로 분할합니다. Llama3 및 Qwen3와 같은 주요 모델에 대한 광범위한 실험을 통해, LongBench, RULER와 같은 장문 컨텍스트 이해 벤치마크 및 AIME24, OlympiadBench와 같은 복잡한 추론 벤치마크에서 LycheeDecode가 전체 어텐션 기준 성능과 동등하거나 그 이상의 생성 품질을 달성함을 확인했습니다. 더욱 중요한 점은, 128K 컨텍스트 길이에서 최대 2.7배의 속도 향상을 달성했습니다. 우리의 세밀한 전략은 어텐션 헤드의 기능적 다양성을 유지함으로써 기존 방법의 성능 병목 현상을 극복하고, 효율적이고 고품질의 장문 컨텍스트 LLM 추론을 위한 강력하고 검증된 방법을 제공합니다.

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.