핵심 개념
DEFT는 트리 구조의 KV 캐시와 부분 결과(QK⊤, Softmax)에 대한 IO 최적화를 통해 대규모 언어 모델 추론의 효율성을 크게 향상시킨다.
초록
이 논문은 대규모 언어 모델(LLM) 추론에서 트리 검색 알고리즘의 효율성을 높이기 위한 DEFT라는 IO 인식 트리 주의 알고리즘을 제안한다.
기존 트리 기반 디코딩 방식은 메모리 저장, 계산, 메모리 접근 측면에서 중복성이 있었다. 특히 KV 캐시 IO가 병목 현상을 일으켰다.
DEFT는 두 단계로 구성된다:
QKV 준비 단계: KV 기반 트리 분할 전략을 통해 KV 캐시 IO를 크게 줄인다.
주의 계산 단계: 퓨전 커널과 트리 토폴로지 인식 전역 감소 전략을 사용하여 부분 결과(QK⊤, Softmax)의 IO를 제거한다.
DEFT는 기존 방식 대비 1.7-2.4배 빠른 추론 속도를 달성했다. 이는 KV 캐시 IO를 3.6-4.5배, QK⊤와 Softmax IO를 25% 줄였기 때문이다.
DEFT는 트리 마스크가 필요 없어 단순하고, 쿼리 수에 민감하지 않아 큰 검색 공간을 지원할 수 있다.
통계
DEFT는 기존 방식 대비 KV 캐시 IO를 3.6-4.5배 줄였다.
DEFT는 기존 방식 대비 QK⊤와 Softmax IO를 25% 줄였다.