llama.cpp, prompt lookup 최적화로 초안 생성 최대 140배 가속

Hayder Tirmazi가 llama.cpp의 prompt lookup decoding(n-gram speculation) 성능 최적화를 발표했다. 초기 최적화로 초안 생성 속도가 최대 42배 빨라지고 메모리 사용량이 2.6배 감소했다. Daniel Lemire의 추가 PR로 총 속도 향상은 최대 140배에 달했다.

핵심 최적화

추론 가속을 위한 n-gram 캐시(context, dynamic, static)를 개선하기 위해 네 가지 단계를 거쳤다.

• 지도 복사 중단 (PR #2): 내부 지도를 복사 대신 참조로 읽도록 변경해 초안 생성 속도를 4.5배~25.6배 향상시켰다.
• 외부 지도를 평면 해시 맵으로 변경 (PR #5): std::unordered_map을 ankerl::unordered_dense::segmented_map으로 교체해 캐시 친화성을 높이고 로딩 시간을 1.41배~1.65배 단축했다.
• 내부 지도를 정렬 벡터로 변경 (PR #10): 2-gram의 64%가 후속 토큰이 하나뿐이라는 점에 착안해 내부 해시 맵을 정렬된 std::vector로 교체했다. 이를 통해 피크 메모리를 최대 1.97배 줄이고 초안 생성 속도를 2.09배 개선했다(static cache 제외).
• Static Cache를 constmap으로 변경: Daniel Lemire의 constmap을 적용해 static cache 로딩 시간을 6.32배~16.12배 단축하고 피크 메모리를 최대 1.30배 줄였다.

성능 지표

Apple M4 Pro(14코어, 48GB RAM)에서 WikiText-103 데이터셋으로 벤치마크를 진행했다.

• 초안 생성 지연: 541MB 코퍼스 기준 지연 시간이 기본값 165.48 µs에서 Lemire의 PR #12 적용 시 1.18 µs로 감소했다. 0MB 코퍼스에서는 0.45 µs까지 떨어졌다.
• 로딩 시간: 541MB 코퍼스의 static cache 로딩 시간이 정렬 벡터 기준 3.76 s에서 constmap 적용 시 0.23 s로 개선됐다.
• 메모리 사용량: 541MB 코퍼스의 피크 메모리가 기본값 3.55 GB에서 constmap과 Lemire PR 적용 시 1.31 GB로 감소했다.

기술 구현

constmap 구현체는 2-gram을 (token, count) 쌍의 연속 배열로 패킹한다. Daniel Lemire의 추가 최적화(PR #12)는 개별 후보 점수를 계산하기 전에 전체 n-gram 개수를 먼저 확인해 임계값을 충족하지 못할 경우 불필요한 계산을 생략하는 방식으로 속도를 높였다.