트랜스포머 KV 캐시 메모리 계산기

트랜스포머 KV 캐시 메모리 계산기 사용법

이 계산기는 트랜스포머 모델이 이전 토큰의 키와 값을 저장하는 데 필요한 KV 캐시 메모리를 추정합니다. 레이어 수, KV 헤드 수, 헤드 차원, 컨텍스트 길이와 데이터 타입을 입력하면 모델 설정에 따른 기본 메모리 요구량을 확인할 수 있어요.

GQA나 MQA 모델은 모든 어텐션 헤드가 아니라 KV를 저장하는 헤드 수가 더 적을 수 있습니다. 이런 경우에는 전체 헤드 수 대신 실제 KV 헤드 수를 입력하면 계산 결과가 모델 구조에 더 가까워집니다.

계산 방법

KV 캐시는 각 레이어에서 키와 값 두 종류의 텐서를 저장합니다. 따라서 전체 저장 요소 수는 레이어 수, KV 헤드 수, 헤드 차원, 컨텍스트 길이를 곱한 뒤 K와 V 두 종류를 반영합니다.

계산 결과는 바이트를 기준으로 하며, 화면에는 이해하기 쉽도록 MiB와 GiB 단위로 함께 표시합니다. FP16과 BF16은 일반적으로 요소당 2바이트, FP32는 4바이트를 사용합니다.

계산 예시

• 상황 1: 레이어 32개, KV 헤드 32개, 헤드 차원 128, 컨텍스트 길이 4,096, FP16인 경우 32 × 2 × 32 × 128 × 4,096 × 2로 계산하며 약 2.00GiB가 필요합니다.
• 상황 2: 레이어 64개, KV 헤드 40개, 헤드 차원 128, 컨텍스트 길이 8,192, BF16인 경우 약 10.00GiB가 필요합니다.
• 상황 3: 동일한 모델 설정에서 FP16 대신 INT8을 사용하면 데이터 타입 바이트 수가 절반으로 줄어 KV 캐시 메모리도 이론적으로 절반이 됩니다.

사용 시 주의사항 및 활용 팁

• 본 계산기 결과는 배치 크기 1 기준의 단순 추정 참고용 수치입니다.
• 여러 요청을 동시에 처리하는 배치 환경에서는 계산 결과에 배치 크기를 곱해야 합니다.
• 실제 GPU 메모리에는 모델 가중치, 활성화 값, CUDA 워크스페이스, 메모리 정렬과 프레임워크 오버헤드가 추가됩니다.
• 양자화된 KV 캐시를 사용할 때는 구현 방식에 따라 실제 저장 형식과 압축 효율이 달라질 수 있습니다.
• 긴 컨텍스트를 지원할수록 KV 캐시가 선형으로 증가하므로, 서비스 운영 시 최대 시퀀스 길이와 동시 요청 수를 함께 검토하세요.