llama.cpp 프롬프트 룩업 초안 생성을 최대 42배 빠르게

AI 요약

Hayder Tirmazi는 llama.cpp의 프롬프트 룩업 디코딩(prompt lookup decoding)에서 초안 토큰을 찾는 과정을 최적화했습니다. WikiText-103으로 정적 n-그램 캐시를 만들고 테스트 문장을 재생하는 방식으로 성능을 비교했습니다. 실험은 Apple M4 Pro, 14코어, 메모리 48GB 환경에서 진행했으며, 결과는 3회 측정값의 중앙값입니다. 변경 전후의 초안 수락률이 거의 같도록 확인하면서 토큰당 초안 생성 시간, 정적 캐시 로딩 시간, 메모리 사용량을 측정했습니다.

프롬프트 룩업과 llama.cpp의 캐시

프롬프트 룩업 디코딩은 현재까지 나온 토큰에서 반복되는 n-그램을 찾아 다음 토큰 후보를 초안으로 제시하는 추측 디코딩(speculative decoding) 방식입니다. llama.cpp는 현재 대화의 토큰을 담는 컨텍스트 캐시, 이전 실행 기록을 담는 동적 캐시, 별도 말뭉치에서 만든 정적 캐시를 사용합니다. n-그램과 뒤따르는 토큰의 빈도를 보고 후보를 고릅니다. 후보의 등장 횟수와 전체 등장 횟수 중 차지하는 비율이 설정 기준을 넘어야 초안으로 채택합니다.

저자는 알고리즘이나 수락 기준을 바꾸지 않고 자료구조와 검색 구현을 개선했습니다. 따라서 이 글에서 성능 변화가 나타나는 지점은 초안 생성 지연, 정적 캐시 로딩, 메모리 사용량입니다. 실험에는 약 541MB인 WikiText-103 전체 말뭉치 외에도 25MB, 50MB, 100MB, 200MB 크기의 자료를 사용했습니다. 정적 캐시를 쓰지 않는 측정도 포함했습니다.

불필요한 복사와 해시 맵 개선

기존 구현은 중첩된 std::unordered_map으로 n-그램마다 뒤따르는 토큰과 빈도를 저장했습니다. 초안 생성 과정에서 내부 맵을 불필요하게 복사하는 부분을 참조로 바꾸자, 말뭉치 크기에 따라 토큰 초안 생성이 4.5~25.6배 빨라졌습니다.

다음으로 바깥쪽 맵을 ankerl::unordered_dense::segmented_map으로 교체했습니다. 저자는 표준 unordered_map의 연결 리스트 기반 충돌 처리 방식이 메모리 접근에 불리하다고 설명합니다. 새 맵은 정적 캐시 로딩을 1.41~1.65배 빠르게 했고, 초안 생성은 1.02~1.13배 빨라졌습니다. 정적 캐시 메모리 사용량은 1.07~1.11배 줄었습니다. 기본 맵은 내부 벡터가 커질 때 전체 크기를 두 배로 늘려야 해 최대 메모리가 증가했습니다. 4096바이트 단위로 확장하는 분할형 맵은 그 문제를 피했습니다.

내부 자료구조와 이진 탐색 최적화

내부 맵까지 해시 맵으로 두는 방식은 메모리를 낭비했습니다. 정적 캐시에서 초안 생성에 쓰이는 2-그램 중 64%는 뒤따르는 토큰이 하나뿐이었습니다. 저자는 내부 맵을 정렬된 벡터로 바꿨습니다. 단순한 std::lower_bound는 자주 등장해 뒤따르는 토큰이 수천 개인 2-그램에서 느렸습니다. 비교 결과에 따라 반복 횟수까지 달라져, 메모리에서 값을 기다리는 동안 다음 작업을 진행하기 어려웠기 때문입니다.

수정한 이진 탐색은 비교 결과와 무관하게 탐색 구간의 길이를 일정한 방식으로 줄입니다. 예를 들어 8개 항목을 찾을 때 반복 횟수는 항상 3회입니다. 저자는 CPU가 현재 메모리 읽기를 기다리는 동안 다음 후보 검색을 진행할 여지를 만든다고 설명합니다. 이 변경으로 정적 캐시가 없을 때 초안 생성은 2.09배 빨라졌습니다. 정적 캐시를 쓸 때는 1.19~1.25배 빨라졌고, 최대 메모리는 1.97배 감소했습니다.

불변 맵으로 정적 캐시 로딩 단축

마지막으로 로딩 뒤에는 내용이 바뀌지 않는 정적 캐시의 바깥쪽 맵을 constmap으로 교체했습니다. n-그램별 토큰과 빈도는 연속된 배열에 저장하고, 맵에는 배열의 시작 위치와 항목 수를 묶어 기록합니다. 파일을 하나의 버퍼로 읽은 뒤 그 버퍼 안에서 맵을 열어 별도 변환 작업을 줄였습니다.

정렬 벡터 방식과 비교하면 정적 캐시 로딩은 6.32~16.12배 빨라졌습니다. 541MB 말뭉치 기준 로딩 시간은 3.76초에서 0.23초로 줄었습니다. 467MB 파일에 대해 캐시는 463MB를 사용했고, 최대 메모리는 1.71GB에서 1.31GB로 감소했습니다. 정적 캐시를 쓸 때 초안 생성은 1.06~1.20배 빨라졌으며, 수락률은 정렬 벡터 방식과 모든 말뭉치에서 같았습니다. 글의 제목에 나온 최대 42배 개선은 여러 최적화를 적용한 초안 생성 단계의 수치입니다.

Reddit 반응

• @u/quantum_being_1 — 저는 한동안 메인라인에 필요한 PR만 몇 개 골라 붙이는 방식으로 써 왔습니다. 제가 본 바로는 프롬프트 룩업 초안 생성의 이득은 작업에 따라 다릅니다. 시스템 프롬프트나 템플릿이 반복되면 n-그램 캐시가 거의 공짜 속도 향상을 주지만, 매번 다른 프롬프트를 쓰면 거의 작동하지 않습니다. 초안 생성 단계에서 42배는 인상적이지만, 전체 추론 과정 중 초안 생성은 일부라서 전체 속도 향상은 그보다 작습니다. 그래도 품질 손실이 거의 없는 드문 최적화라 메인라인에 들어가면 좋겠습니다.
  • @u/Available_Pressure47 — 좋은 지적입니다. 말씀하신 것처럼 프롬프트 특성에 따라 프롬프트 룩업이나 일반적인 추측 디코딩의 이득이 달라집니다. 어떤 프롬프트에서는 크게 빨라지지만, 다른 프롬프트에서는 속도 향상이 10% 미만일 수도 있습니다.
• @u/New_Comfortable7240 — 언젠가 메인 llama.cpp에 병합되면 좋겠습니다. 또 하나의 포크가 늘어나는 점은 조금 걱정됩니다.
  • @u/Available_Pressure47 — 감사합니다. 제가 포크 저장소의 PR에서 기여자 한 명을 실수로 태그했습니다. 실제 PR이 아닌데도 승인된 뒤, 시간을 낭비하게 했다는 이유로 PR 작성이 차단됐습니다. 더 귀찮게 할까 봐 어떻게 해야 할지 모르겠습니다.
• @u/trying4k — Johannes의 반응은 정말 이해하기 어렵습니다. 한 번 언급했다고 왜 그렇게 반응했을까요? 다만 그 행동이 정당화되는 건 아니지만, 자기 저장소에 PR을 만든 점은 제게도 조금 이상해 보입니다. 메인 저장소에 기여하는 길이 다시 열리길 바랍니다. 개선 내용은 좋아 보입니다.
  • @u/Available_Pressure47 — 맞습니다. PR은 글에서 최적화를 하나씩 연결해 보여주려고 만들었습니다. 개발만이 목적이었다면 각 변경을 따로 문서화하는 방식이 최선은 아닐 수 있습니다. 최근에는 차단됐다는 이유로 PR을 만들지 못했습니다. 잘 풀리기를 바란다는 말씀 감사합니다.
• @u/fallingdowndizzyvr — 포크를 걱정하기엔 이미 늦었습니다. 포크가 미래입니다. llama.cpp는 범용이라 쓸 곳이 있지만, 그만큼 성능이 좋지 않습니다. Strix Halo에서 Qwen FN을 쓸 때 일반 llama.cpp는 프롬프트 처리 속도가 약 300토큰/초였습니다. Strix Halo 전용 포크에서는 600토큰/초, 처음부터 해당 장비용으로 만든 다른 패키지에서는 1200토큰/초 이상을 얻었습니다.
  • @u/miversen33 — 저는 메인라인 위에 패치를 얹어 제 방식대로 llama.cpp를 만들고 있습니다.
• @u/ThatsALovelyShirt — 이 변경을 로컬 포크에 병합하려면 ngram-cache-constmap 브랜치와 ngram-cache-no-copy-upstream 중 무엇을 써야 하나요?
  • @u/Available_Pressure47 — 이 브랜치를 쓰시면 됩니다. 최적화 변경 전체가 포함돼 있습니다.

원문: Jadid Bourbaki / 번역·요약: Trawling