KV Cache from scratch in nanoVLM
Quick Summary
나노브이엘엠에 계층별 키·값 캐시와 사전 채움·순차 해독 구조를 직접 구현해, 자기회귀 생성의 중복 계산을 줄이고 생성 속도를 38% 높인 과정과 원리를 설명한다.
🖼️ 인포그래픽
🖼️ 4컷 인포그래픽
💡 한 줄 요약
나노브이엘엠에 계층별 키·값 캐시와 사전 채움·순차 해독 구조를 직접 구현해, 자기회귀 생성의 중복 계산을 줄이고 생성 속도를 38% 높인 과정과 원리를 설명한다.
📌 핵심 요약
- 자기회귀 언어 모델은 토큰을 하나씩 생성하며, 캐시가 없으면 새 토큰을 예측할 때마다 변하지 않은 이전 토큰들의 질의·키·값을 전체 계층에서 다시 계산한다.
- 키·값 캐시는 최초 입력을 처리할 때 각 계층의 키와 값을 저장하고, 이후에는 새 토큰의 질의·키·값만 계산한 뒤 새 키와 값을 기존 캐시에 이어 붙인다.
- 나노브이엘엠 구현은 주의집중 블록의 캐시 갱신, 언어 모델의 계층별 캐시 관리와 시작 위치 처리, 생성 루프의 사전 채움과 순차 해독 분리라는 세 부분으로 구성된다.
- 회전 위치 임베딩에는 현재 토큰의 실제 위치를 나타내는 시작 위치가 전달되며, 이를 통해 캐시를 사용할 때도 새 토큰의 위치 정보를 올바르게 계산한다.
- 이 구현은 나노브이엘엠의 생성 속도를 38% 높였지만, 캐시가 문맥과 계층 수에 따라 커지는 메모리 비용, 코드 복잡성 증가, 빔 탐색 같은 생성 방식의 제약을 동반한다.
🧩 주요 포인트
- 자기회귀 언어 모델은 토큰을 하나씩 생성하며, 캐시가 없으면 새 토큰을 예측할 때마다 변하지 않은 이전 토큰들의 질의·키·값을 전체 계층에서 다시 계산한다.
- 키·값 캐시는 최초 입력을 처리할 때 각 계층의 키와 값을 저장하고, 이후에는 새 토큰의 질의·키·값만 계산한 뒤 새 키와 값을 기존 캐시에 이어 붙인다.
- 나노브이엘엠 구현은 주의집중 블록의 캐시 갱신, 언어 모델의 계층별 캐시 관리와 시작 위치 처리, 생성 루프의 사전 채움과 순차 해독 분리라는 세 부분으로 구성된다.
- 회전 위치 임베딩에는 현재 토큰의 실제 위치를 나타내는 시작 위치가 전달되며, 이를 통해 캐시를 사용할 때도 새 토큰의 위치 정보를 올바르게 계산한다.
- 이 구현은 나노브이엘엠의 생성 속도를 38% 높였지만, 캐시가 문맥과 계층 수에 따라 커지는 메모리 비용, 코드 복잡성 증가, 빔 탐색 같은 생성 방식의 제약을 동반한다.
🧠 상세 정리
1. 자기회귀 생성에서 발생하는 반복 계산
자기회귀 언어 모델은 주어진 입력으로 다음 토큰을 예측하고, 그 토큰을 기존 문장 뒤에 붙인 다음 같은 절차를 반복한다. 새로운 토큰 하나를 만들 때마다 지금까지의 모든 토큰을 고려해야 하므로, 생성 과정 자체는 앞 단계의 결과에 의존하는 순차적 구조를 가진다. 트랜스포머 내부 연산은 병렬화할 수 있지만, 캐시가 없는 생성에서는 문장이 한 토큰 늘어날 때마다 전체 시퀀스를 모든 계층에 다시 통과시킨다. 그 결과 이미 처리한 토큰에 대한 연산까지 반복되며, 시퀀스가 길어질수록 메모리와 계산 부담이 빠르게 커진다. 글은 이러한 중복을 제거하는 방법으로 키·값 캐시를 제시하고, 작은 순수 파이토치 기반 시각언어모델인 나노브이엘엠에 이를 직접 구현한 결과 생성 속도가 38% 향상됐다고 보고한다.
2. 자기주의집중의 기본 계산 구조
트랜스포머의 각 계층은 다중 머리 자기주의집중, 순방향 신경망, 잔차 연결, 계층 정규화로 구성되며, 키·값 캐시가 직접 작용하는 부분은 자기주의집중이다. 입력 임베딩 행렬에 서로 다른 투영 행렬을 곱해 질의, 키, 값을 만들고, 질의와 키의 내적을 키 차원의 제곱근으로 나눈 점수에 인과 마스크를 적용한다. 인과 마스크는 각 위치가 미래 토큰을 보지 못하게 하며, 마스킹된 점수에 소프트맥스를 적용한 가중치와 값을 곱해 최종 출력을 얻는다. 글은 작은 무작위 텐서와 파이토치 코드를 이용해 이 과정을 단계별로 보여 주며, 주의집중 계산에서 어떤 값이 매 생성 단계마다 새로 필요하고 어떤 값은 재사용할 수 있는지 확인할 토대를 마련한다. 특히 이전 토큰의 표현이 바뀌지 않는 자기회귀 해독에서는 그 토큰들에서 계산된 키와 값 역시 변하지 않는다는 점이 핵심이다.
3. 키와 값의 중복 재계산 확인
기존 시퀀스에 새 토큰 임베딩 하나를 추가해 질의, 키, 값을 다시 계산하면, 마지막 토큰을 제외한 앞부분의 키와 값은 이전 단계에서 계산한 결과와 동일하다. 글은 확장 전 키·값 텐서와 확장 후 텐서의 기존 시퀀스 구간을 직접 비교하는 검사를 통해 두 결과가 같음을 확인한다. 그런데 캐시가 없는 구현은 동일한 값이라는 사실을 활용하지 못하고, 새 토큰이 들어올 때마다 전체 시퀀스의 투영과 주의집중 연산을 반복한다. 시각화에서는 기존에 계산된 영역 대부분이 다시 처리되고 실제로 새로 필요한 부분은 마지막 위치뿐이라는 점을 드러낸다. 따라서 생성 길이가 늘어날수록 반복되는 영역도 함께 커지며, 이는 긴 문맥과 긴 출력에서 추론 효율을 저해하는 주요 원인이 된다.
4. 키·값 캐시의 핵심 원리
키·값 캐시는 최초 입력을 처리한 뒤 각 트랜스포머 계층에서 계산된 키와 값을 보존하는 방식으로 중복을 제거한다. 이후 생성 단계에서는 현재의 새 토큰에 대해서만 질의, 키, 값을 계산하고, 새 키와 값을 기존 캐시의 시퀀스 축 뒤에 이어 붙인다. 현재 토큰의 질의는 이렇게 누적된 모든 키 및 값과 함께 사용되므로, 이전 문맥 전체를 참고하면서도 이전 토큰의 키와 값을 다시 투영할 필요가 없다. 캐시는 계층마다 별도로 존재하며, 각 항목은 키와 값으로 구성되고 배치 크기, 머리 수, 누적 시퀀스 길이, 머리 차원의 형태를 가진다. 캐시 길이는 토큰이 생성될 때마다 한 위치씩 증가하며, 전체 시퀀스 재계산을 가벼운 증분 갱신으로 바꾸는 것이 이 최적화의 본질이다.
5. 주의집중 블록의 캐시 갱신
나노브이엘엠의 주의집중 클래스는 순방향 함수가 해당 계층의 키·값 캐시를 입력받고 갱신된 캐시를 반환하도록 변경됐다. 최초 입력을 처리하는 사전 채움 단계에서는 기존 캐시가 없으므로 전체 입력에서 질의·키·값을 계산하고, 회전 위치 임베딩이 적용된 키와 값을 새 캐시로 저장한다. 이후 순차 해독 단계에서는 현재 토큰에서 얻은 새 키를 저장된 키 뒤에 연결하고, 새 값도 같은 방식으로 저장된 값 뒤에 연결한다. 주의집중 출력은 현재 질의와 누적 키·값을 이용해 계산되며, 함수는 출력과 함께 확장된 캐시를 다음 단계로 넘긴다. 이 변경으로 주의집중 블록은 이전 상태를 재사용할 수 있게 되었고, 매 단계 전체 문맥의 키와 값을 다시 만드는 기존 동작을 대체한다.
6. 계층별 캐시와 위치 정보 관리
언어 모델 수준에서는 트랜스포머 계층마다 하나씩 키·값 캐시를 유지하도록 목록 형태의 상태를 도입한다. 모델은 각 블록을 순회하면서 그 계층에 대응하는 캐시만 전달하고, 블록이 반환한 갱신 결과를 같은 위치에 다시 저장한다. 또한 현재 입력이 전체 시퀀스의 어느 위치에서 시작하는지를 나타내는 시작 위치를 순방향 함수에 추가해, 새 토큰에 적용할 회전 위치 임베딩의 위치 번호를 계산한다. 현재 입력 길이가 하나라면 시작 위치부터 한 개의 위치 번호를 만들고, 여러 토큰을 한꺼번에 처리하는 최초 단계에서는 해당 범위 전체의 위치 번호를 만든다. 이 처리가 빠지면 캐시에 저장된 과거 토큰과 새 토큰의 위치 관계가 어긋날 수 있으므로, 시작 위치는 캐시 사용 중에도 위치 인코딩의 정확성을 유지하는 핵심 요소다.
7. 사전 채움과 순차 해독의 분리
가장 큰 구조적 변화는 생성 함수를 사전 채움 단계와 순차 해독 단계로 분리한 것이다. 사전 채움에서는 전체 입력 프롬프트를 한 번에 모델에 전달해 첫 출력과 모든 계층의 초기 키·값 캐시를 만든다. 순차 해독에서는 직전 출력에서 다음 토큰을 선택한 뒤 그 토큰 하나만 모델에 넣고, 기존 캐시와 현재 위치를 함께 전달해 다음 출력을 얻는다. 이렇게 갱신된 출력과 캐시는 다시 다음 생성 단계의 입력 상태로 사용되며, 종료 조건이나 최대 생성 토큰 수에 도달할 때까지 같은 절차를 반복한다. 긴 프롬프트를 처음 한 번 처리한 후에는 매 단계 새 토큰만 계산하므로, 전체 프롬프트를 반복 입력하던 단일 단계 생성 방식보다 특히 긴 입력에서 큰 효율 향상을 얻는다.
8. 효율 향상과 비용의 균형
키·값 캐시를 적용하면 토큰마다 전체 시퀀스를 다시 계산하는 대신 누적 문맥을 활용한 선형 규모의 주의집중 계산으로 전환되며, 나노브이엘엠에서는 실제 생성 속도가 38% 향상됐다. 캐시는 새 토큰마다 한 위치씩 커지고 시작 위치를 통해 위치 인코딩의 정합성도 유지하므로, 긴 시퀀스와 실시간 생성에서 특히 유용하다. 다만 속도 향상은 무료가 아니며, 각 계층의 키와 값을 계속 보관해야 하므로 문맥 길이가 늘어날수록 메모리 사용량도 증가한다. 캐시 상태를 계층별로 전달하고 관리해야 해 코드가 복잡해지며, 글은 빔 탐색과 같은 더 복잡한 추론 방식을 적용할 때 제약이 생길 수 있다고 지적한다. 그럼에도 불필요한 반복 연산을 제거하는 효과가 커서, 키·값 캐시는 대규모 언어 모델의 추론을 빠르게 하고 소비자용 하드웨어에서도 실행 가능성을 높이는 대표적인 최적화로 평가된다.
🧾 핵심 주장 / 시사점
- 핵심 최적화 대상은 과거 토큰 전체가 아니라, 생성 단계가 바뀌어도 값이 변하지 않는 계층별 키와 값이다. 무엇이 불변인지 확인한 뒤 그 결과만 보존하는 것이 캐시 설계의 출발점이다.
- 키·값 캐시의 정확성은 저장 자체뿐 아니라 사전 채움과 순차 해독의 분리, 계층별 상태 대응, 회전 위치 임베딩을 위한 시작 위치 관리가 함께 맞물릴 때 확보된다.
- 이 구현은 계산량을 메모리 사용량과 교환한다. 따라서 긴 문맥에서 얻는 속도 이점이 크지만, 캐시 크기와 생성 방식의 복잡성까지 고려해 적용해야 한다.
✅ 액션 아이템
- 자기회귀 생성 경로에서 매 토큰 재계산되는 이전 토큰 질의·키·값을 줄이기 위해 주의집중 블록 캐시 갱신을 동일한 방식으로 구현한다.
- 계층별 캐시를 층 단위로 관리하고 start_pos를 정확히 전달한 뒤 사전 채움과 순차 해독을 분리해 위치 임베딩 정합성을 보장한다.
- 속도 38% 개선 효과를 문맥 길이·층 수별 메모리 비용, 코드 복잡도 증가, 빔 탐색 제약과 함께 측정해 캐시 적용 범위를 정한다.
❓ 열린 질문
- 문맥 길이가 증가할수록 캐시 메모리 사용량이 급증해 병목이 되는 임계 지점은 어느 정도인가?
- 회전 위치 임베딩에서 start_pos를 갱신할 때 캐시 연속성이 깨지는 경우를 어떤 조건으로 판별할 것인가?
- 캐시 적용 시 빔 탐색 제약이 성능 이득을 상쇄하는 구간은 어떤 경우에 전환 전략을 둬야 할 것인가?