Hydragen: 공유 프리픽스를 활용한 고처리량 LLM 추론

초록

Transformer 기반의 대규모 언어 모델(LLM)은 현재 수억 명의 사용자에게 배포되고 있다. LLM 추론은 일반적으로 몇 가지 예시나 챗봇 시스템 프롬프트와 같은 공통 접두사를 공유하는 시퀀스 배치에서 수행된다. 이러한 대규모 배치 설정에서 디코딩은 주의력 연산(attention operation)에 의해 병목 현상이 발생할 수 있으며, 이 연산은 메모리에서 대규모 키-값(KV) 캐시를 읽고 배치 내 모든 시퀀스에 대해 비효율적인 행렬-벡터 곱셈을 계산한다. 본 연구에서는 공유 접두사를 가진 주의력 연산을 하드웨어를 고려한 정확한 방식으로 구현한 Hydragen을 소개한다. Hydragen은 공유 접두사와 고유 접미사에 대한 주의력 연산을 별도로 계산한다. 이 분해는 시퀀스 간 쿼리를 함께 배치 처리함으로써 접두사 주의력 연산을 효율적으로 수행하며, 중복된 메모리 읽기를 줄이고 하드웨어 친화적인 행렬 곱셈을 가능하게 한다. 우리의 방법은 경쟁력 있는 기준선 대비 최대 32배까지 LLM의 종단 간 처리량을 향상시킬 수 있으며, 배치 크기와 공유 접두사 길이가 증가함에 따라 속도 향상이 커진다. Hydragen은 또한 매우 긴 공유 컨텍스트 사용을 가능하게 한다: 높은 배치 크기에서 접두사 길이를 1K 토큰에서 16K 토큰으로 증가시켰을 때, Hydragen의 처리량은 15% 미만으로 감소하는 반면, 기준선의 처리량은 90% 이상 감소한다. Hydragen은 단순한 접두사-접미사 분해를 넘어 트리 기반 프롬프트 공유 패턴에도 적용할 수 있어, 경쟁 프로그래밍 문제에서 추론 시간을 추가로 55% 줄일 수 있다.

English

Transformer-based large language models (LLMs) are now deployed to hundreds of millions of users. LLM inference is commonly performed on batches of sequences that share a prefix, such as few-shot examples or a chatbot system prompt. Decoding in this large-batch setting can be bottlenecked by the attention operation, which reads large key-value (KV) caches from memory and computes inefficient matrix-vector products for every sequence in the batch. In this work, we introduce Hydragen, a hardware-aware exact implementation of attention with shared prefixes. Hydragen computes attention over the shared prefix and unique suffixes separately. This decomposition enables efficient prefix attention by batching queries together across sequences, reducing redundant memory reads and enabling the use of hardware-friendly matrix multiplications. Our method can improve end-to-end LLM throughput by up to 32x against competitive baselines, with speedup growing with the batch size and shared prefix length. Hydragen also enables the use of very long shared contexts: with a high batch size, increasing the prefix length from 1K to 16K tokens decreases Hydragen throughput by less than 15%, while the throughput of baselines drops by over 90%. Hydragen generalizes beyond simple prefix-suffix decomposition and can be applied to tree-based prompt sharing patterns, allowing us to further reduce inference time on competitive programming problems by 55%.

Hydragen: 공유 프리픽스를 활용한 고처리량 LLM 추론

Hydragen: High-Throughput LLM Inference with Shared Prefixes

초록

Support