AI 추론 워크로드에서
필요한 메모리를 역산합니다

토큰 하나당 KV cache 크기는 어텐션 방식마다 계산식이 다릅니다. MHA 와 GQA 와 MQA 는 K 와 V 를 각각 저장하지만, MLA 는 저차원 잠재 표현 하나만 저장하므로 계산식 자체가 달라집니다. 이 차이를 모델 프로파일에 반영해 두고 시작합니다.

모델 프로파일과 가속기 사양과 메모리 계층표 세 가지를 놓고 워크로드를 입력하면, 가중치와 KV cache 용량, 필요한 가속기 대수, 그 가운데 KV cache 때문에 증가한 대수, roofline 상의 병목 종류, 초과분을 어느 계층에 배치할 수 있는지가 산출됩니다.

값보다 중요한 것은 병목의 종류가 고객마다 다르다는 사실입니다. 같은 엔진에 MLA 계열의 대형 모델을 입력하면 KV cache 와 가중치의 비율이 0.86배로 역전되어 병목이 KV cache 가 아니라 가중치가 됩니다. 그 고객에게 계층 분리를 제안하면 초점을 잘못 읽은 것이 됩니다.

사이클마다 페이지 캐시 점유량과 major page fault 와 프로세스 상주 메모리와 커널의 메모리 압력 지표를 기록합니다. 계산기가 산출한 숫자를 같은 기계에서 검증하는 것이 목적입니다.

메모리 매핑으로 색인을 열면 프로세스의 메모리 사용량은 거의 증가하지 않습니다. 그래서 메모리를 사용하지 않는 것처럼 보이지만 페이지 캐시를 통해 DRAM 을 실제로 점유하고 있습니다. E3 가 측정하는 대상이 정확히 이 지점입니다.

하이퍼스케일러 자체 실리콘 조직, 범용 가속기 공급사, 신흥 클라우드 사업자, 스마트폰 제조사, 엣지 및 로보틱스 다섯 종입니다. 각각에 누가 결재하는지, 무엇으로 평가받는지, 무엇을 양보할 수 없는지, 의사결정 주기가 어떠한지를 기술하였습니다.

성격과 말투는 기술하지 않았습니다. 구매 조직의 구조가 문제점의 성격을 결정하기 때문입니다. 각 아키타입에는 그 고객이 제기할 반론과 공급 측이 그 고객을 놓치는 이유가 함께 들어갑니다.

출처를 확인하지 못한 주장 13건은 문제점 목록에서 분리하여 별도로 관리하고 도구가 산출하는 결과에는 포함하지 않습니다. 정량 근거를 확보하지 못한 아키타입에는 그 사실을 표시해 두고 수치를 제시하지 않습니다. 엣지 및 로보틱스가 현재 그 상태입니다.

5 GB 벡터 색인을 메모리 매핑으로 열어 두고 무작위 4KB 읽기로 탐색 패턴을 모사합니다. 그 상태에서 동거 프로세스가 메모리를 256 MB 씩 증가시키며 점유하고, 사이클마다 검색 지연과 major page fault 를 기록합니다. 호스트 메모리는 15.7 GB 입니다.

점유량이 8.00 GB 일 때까지 99분위 지연이 1 µs 에서 변하지 않습니다. 8.25 GB 에서 3,864 µs 로 상승하며 이는 기준값의 4,102배입니다. 이 시점에 남아 있는 메모리가 6.67 GB 이고 호스트 메모리의 42퍼센트에 해당합니다.

50분위 지연은 그보다 0.75 GB 뒤에 상승하기 시작합니다. 평균값으로 감시하면 그만큼 늦게 인지하게 됩니다.

색인 2.5 GB 에 호스트 메모리 7.8 GB 인 다른 환경에서도 같은 경향이 관측되었습니다. 그 환경에서는 3.86 GB, 즉 호스트 메모리의 49퍼센트가 남은 지점에서 상승이 시작되었습니다. 절대값은 환경마다 다르지만 경향은 동일합니다.

같은 모델을 F32 와 F16 과 Q8_0 과 Q5_0 과 Q4_K_M 다섯 수준으로 준비하고 llama-bench 로 처리량을 측정하였습니다. pp512 가 prefill 단계이고 tg128 이 decode 단계입니다. 두 값을 합치지 않았습니다. 합치는 순간 병목의 성격이 사라집니다.

F32 에서 Q4_K_M 으로 낮추면 모델 크기가 0.15배가 되는데 prefill 은 60퍼센트 개선되는 데 그치지만 decode 는 148퍼센트 개선됩니다. 네 단계 전부에서 decode 쪽의 개선 폭이 더 큽니다.

prefill 은 단조롭게 증가하지도 않습니다. F16 과 Q5_0 에서는 오히려 F32 보다 느려집니다. 이 프로세서는 정수 내적 명령을 지원하므로 Q8_0 과 Q4_K_M 은 그 경로를 사용하지만 F16 과 Q5_0 은 사용하지 못하기 때문입니다. prefill 의 성능을 결정하는 것이 모델 크기가 아니라 연산 경로라는 사실은, prefill 이 연산에 의해 제한된다는 진술을 오히려 강하게 뒷받침합니다.

decode 는 대역폭에 완전히 지배될 때의 이론선에도 미치지 못합니다. 모델 크기가 0.15배이므로 이론선은 6.5배이지만 실제 개선은 2.48배에 그칩니다. 역양자화에 연산이 소요되기 때문입니다.

어텐션 구조마다 KV cache 의 계산식 자체가 다릅니다. 표준 GQA 는 K 와 V 를 각각 저장하고, MLA 는 저차원 잠재 표현 하나만 저장하므로 계수 2가 없습니다. sliding window 혼합은 full 레이어만 컨텍스트 길이에 비례하고, 상태공간 모형 혼합은 어텐션 레이어만 비례하며 상태는 상수입니다.

네 구조를 같은 조건에 두면 토큰당 160 KiB 에서 6 KiB 까지 차이가 납니다. 컨텍스트 128K 에서 20 GB 를 사용하는 고객과 0.75 GB 를 사용하는 고객이 같은 가속기를 사용합니다. 고객이 어떤 모델을 서빙하는지에 따라 제안할 제품이 달라진다는 뜻입니다.

이 값은 계산값입니다. 가속기가 없어서 네 모델을 실제로 서빙하지 못하였습니다. 다만 기울기를 구하는 측정 절차는 검증하였습니다. 표준 GQA 구조의 모델에서 컨텍스트를 2,048 과 8,192 와 16,384 로 변경하며 엔진이 보고하는 메모리를 관측한 결과, 기울기가 48.0 KiB/token 으로 산출되었고 계산식이 주는 값과 100퍼센트 일치하였습니다.

“Attention Is All You Need” 의 encoder 와 decoder 를 PyTorch 기본 텐서 연산만으로 구현하였습니다. scaled dot-product attention 과 multi-head attention 과 encoder 와 decoder 를 연결하는 cross-attention 까지 전부 직접 작성하였습니다. GPT 계열인 ②에는 encoder 도 cross-attention 도 없습니다. 원형 전체를 다룬다는 점이 이 프로젝트의 차이입니다.

동작은 날짜 표기 변환으로 확인하였습니다. 번역의 가장 단순한 형태에 해당합니다. 입력과 출력의 길이와 어순과 표기가 모두 다르기 때문에, encoder 가 입력을 요약하고 decoder 가 cross-attention 으로 그것을 참조하지 않으면 풀리지 않습니다.

attention 블록부터 학습 루프까지 전부 직접 작성하였습니다. 여기에 KV cache 를 구현하였습니다. 생성할 때 이미 계산한 K 와 V 를 저장해 두고 새 토큰 하나만 계산하는 방식입니다. nanoGPT 계열의 참고 구현에는 없으므로 직접 설계하였습니다.

Part 1 의 계산기가 사용하는 2 × 층 수 × KV head × head dim × 정밀도 계산식이 여기에서 나왔습니다. 이 구현에서 측정한 캐시 크기가 계산식과 일치하는 것을 먼저 확인하고, 같은 식을 대형 모델에 적용합니다.

GPT-2(124M) 를 동적 양자화로 INT8 로 변환하였습니다. 모델의 내부 구조를 파악해야 하는 작업입니다. GPT-2 는 표준 Linear 대신 Conv1D 를 사용하므로 양자화 API 가 인식하도록 nn.Linear 로 먼저 변환해야 했습니다. 출력층인 lm_head 는 입력 임베딩과 가중치를 공유하고 있으므로 양자화 대상에서 분리하여 FP32 로 유지하였습니다.

그리고 스케일을 행렬 전체에 하나 부여하는 경우와 출력 채널마다 부여하는 경우를 같은 조건에서 비교하였습니다. 비트 수는 동일한데 결과가 달라집니다. Part 1 에서 정밀도를 워크로드의 입력값으로 받는 이유가 여기에 있습니다.

문서를 분할하고 임베딩하여 FAISS 색인에 넣고 유사도로 검색한 다음 그 결과를 프롬프트에 삽입하는 과정까지 전부 직접 작성하였습니다. LangChain 같은 프레임워크는 사용하지 않았습니다. 단계마다 무엇이 오가는지 코드에서 확인할 수 있도록 하기 위해서입니다.

색인이 메모리를 얼마나 점유하는지도 여기에서 나왔습니다. Part 1 의 E3 는 이 색인을 키운 다음, 메모리가 부족해지기 시작하면 검색이 어떻게 되는지를 측정합니다.