연속 배치의 비동기성 구현
TL;DR: CPU와 GPU 작업을 분리해 추론 성능을 크게 끌어올리는 방법을 설명합니다.
이 글은 효율적인 LLM 추론 시리즈의 두 번째 포스트입니다. 첫 번째 글(연속 배칭 입문)에서는 연속 배칭을 처음부터 설명했으며, KV 캐시, FlashAttention, 어텐션 마스크 등 여기서 활용할 개념들을 소개했습니다.
H200 한 대는 Inference Endpoints에서 시간당 약 5 달러에 이용할 수 있습니다. 한 시간만 쓰면 저렴하지만 하루 종일 쓰면 이미 120 달러가 됩니다. 이런 상황이라면 GPU를 최대한 활용하고 싶을 겁니다.
연속 배칭은 패딩 없이 빽빽하게 배치를 스케줄링해 GPU 활용도를 높여줍니다. 하지만 연속 배칭이 해결하지 못하는 두 번째 낭비가 있습니다: 기본적으로 동기 방식이기 때문입니다. 즉 CPU와 GPU가 번갈아 가며 일합니다. GPU가 연산하는 동안 CPU는 대기하고, CPU가 다음 배치를 준비하는 동안 GPU는 대기합니다. 초당 수백 번의 스텝을 수행하는 루프에서는 이러한 대기 시간이 누적돼 전체 실행 시간의 거의 ¼을 차지할 수 있습니다. GPU가 100 % 연산하도록 하려면 이 공백을 없애야 합니다.
이를 위해 비동기 배칭을 사용합니다. CPU 배치 준비와 GPU 배치 연산을 분리해 동시에 진행하게 하면 언제나 GPU가 바쁘게 동작합니다 🔥
동기 배칭
다음은 전형적인 동기 배칭 방식입니다:
CPU가 새 배치를 준비할 때 포함할 요청을 고르고, KV 캐시 테이블을 업데이트하며, 이전 라운드에서 완료된 요청을 내보내고, 빈 공간을 새로운 요청으로 채웁니다. 준비가 끝나면 입력을 GPU로 전송합니다. GPU는 포워드 패스를 실행하고 각 요청에 대해 새로운 토큰을 샘플링합니다. 결과가 CPU로 돌아오면 각 요청이 방금 생성한 토큰을 알 수 있고, 다시 전체 사이클이 반복됩니다.
오른쪽 빨간색 주석을 보면 GPU가 연산을 마친 뒤 대기 상태가 됩니다. 다음 배치는 CPU가 출력 토큰을 샘플링하고, 요청 상태를 업데이트하고, 배치를 재스케줄링하는 과정을 마쳐야 시작될 수 있습니다.
이것이 동기 배칭의 핵심 비효율성입니다: CPU와 GPU가 번갈아 가며 일합니다. GPU가 연산하는 동안 CPU는 대기하고, CPU가 업데이트하는 동안 GPU는 대기합니다. 두 장치가 동시에 유용한 작업을 수행하는 경우는 없습니다. 단일 포워드 패스에서는 큰 문제가 아닐 수 있지만, 초당 수백 스텝을 수행하는 연속 배칭 루프에서는 이러한 대기 구간이 실제 처리량 손실로 이어집니다.
이를 보여주기 위해 8 B 모델에 배치 크기 32로 8 K 토큰을 생성할 때 CPU와 GPU가 각각 얼마나 시간을 소비하는지 프로파일링했습니다:
동일한 그래프를 만들고 싶다면 연속 배칭 코드를 계측해 CPU·GPU 활동 구간을 덤프하고 이 스크립트를 사용하면 됩니다.
타임라인은 초록색( GPU 활성, CPU 대기)과 빨간색( CPU 활성, GPU 대기)으로 번갈아 나타나며 두 구간이 겹치지 않습니다. 전체 생성 시간은 300.6 초였으며, 그 중 24.0 %는 GPU가 CPU가 끝나기를