SinkRouter: 효율적인 장문 텍스트 디코딩을 위한 싱크(Sink)-인식 라우팅 - 대규모 언어 및 다중 모달 모델

SinkRouter: Sink-Aware Routing for Efficient Long-Context Decoding in Large Language and Multimodal Models

대규모 언어 모델(LLM) 및 다중 모달 모델(LMM)에서 장문 텍스트 디코딩 과정에서 어텐션 연산은 메모리 병목 현상을 야기하는데, 이는 각 디코딩 단계에서 GPU 메모리에서 많은 양의 KV 캐시 데이터를 로드해야 하기 때문입니다. 기존의 가속화 전략들은 종종 효율성을 높이기 위해 정확도를 희생하며, 유용한 정보를 제거할 수 있는 휴리스틱 기반의 가지치기를 사용합니다. 또한, 이러한 방법들은 종종 모든 높은 점수 토큰을 무분별하게 유지하거나, 초기 토큰을 필수적인 고정점(anchor)으로 취급하거나, 휴리스틱 기반의 헤드 라우팅에 의존하며, 이는 어텐션 싱크(attention sink) 현상에 대한 충분한 메커니즘적 이해를 반영하지 못합니다. 본 논문에서는 어텐션 싱크 현상이 학습 과정에서 생성되는 안정적이고 도달 가능하며, 오류 제어가 가능한 고정점(fixed point)과 일치한다는 것을 보여줍니다. 이러한 통찰력을 바탕으로, 우리는 SinkRouter라는 학습이 필요 없는 선택적 라우팅 프레임워크를 제안합니다. 이 프레임워크는 싱크 신호를 감지하고, 그렇지 않으면 거의 0에 가까운 출력을 생성하는 연산을 건너뜁니다. 이 메커니즘을 실제 가속화로 구현하기 위해, 우리는 블록 수준의 분기(branching) 및 Split-K 병렬 처리를 지원하는 하드웨어 친화적인 Triton 커널을 개발했습니다. 우리는 Llama-3.1-8B, Llama-3.1-70B, Yi-9B-200K, LLaVA-1.5-7B, LLaVA-1.5-13B와 같은 텍스트 전용 및 다중 모달 백본을 사용하여 LongBench, InfiniteBench, CVBench, MileBench, MMVP 등 다양한 장문 텍스트 벤치마크에서 광범위한 평가를 수행했습니다. 이러한 다양한 설정에서, SinkRouter는 일관되게 디코딩 효율성을 향상시키면서 경쟁력 있는 정확도를 유지하며, 512K 컨텍스트에서 최대 2.03배의 속도 향상을 달성했습니다.

In long-context decoding for LLMs and LMMs, attention becomes increasingly memory-bound because each decoding step must load a large amount of KV-cache data from GPU memory. Existing acceleration strategies often trade efficiency for accuracy by relying on heuristic pruning that may discard useful information. At a deeper level, they also tend to indiscriminately preserve all high-scoring tokens, treat early tokens as indispensable anchors, or rely on heuristic head routing, reflecting an insufficient mechanistic understanding of the attention sink phenomenon. In this paper, we show that the attention sink phenomenon corresponds to a stable, reachable, and error-controllable fixed point constructed during training. Based on this insight, we propose SinkRouter, a training-free selective routing framework that detects the sink signal and skips computations that would otherwise produce near-zero output. To translate this mechanism into real-world acceleration, we develop a hardware-aware Triton kernel with block-level branching and Split-K parallelism. We conduct extensive evaluations on a diverse suite of long-context benchmarks, including LongBench, InfiniteBench, CVBench, MileBench, and MMVP, using both text-only and multimodal backbones such as Llama-3.1-8B, Llama-3.1-70B, Yi-9B-200K, LLaVA-1.5-7B, and LLaVA-1.5-13B. Across these settings, SinkRouter consistently improves decoding efficiency while maintaining competitive accuracy, and reaches 2.03x speedup with a 512K context.