Continuous Batching에서 비동기성(asynchronicity) 구현하기

요약(TL;DR): 추론(inference) 성능을 대폭 향상시키기 위해 CPU와 GPU 워크로드를 분리하는 방법을 설명합니다.

이 글은 효율적인 LLM 추론에 관한 시리즈의 두 번째 포스트입니다. 첫 번째 포스트에서는 원리부터 시작하는 Continuous Batching을 다루었습니다. 여기서는 우리가 기반으로 삼을 몇 가지 개념, 즉 KV cache, FlashAttention, attention masks 등을 소개했습니다.

Inference Endpoints에서 H200의 비용은 시간당 약 5달러입니다. 한 시간 동안 사용하기에는 저렴하지만, 하루 동안 사용하면 이미 120달러를 지불하게 됩니다. 상황이 이렇다면, GPU를 최대한 활용하고 싶을 것입니다.

우리는 Continuous Batching이 빽빽하게 채워진 배치(batches)를 스케줄링함으로써 GPU 활용도(utilization)를 개선하고, 패딩(padding)에 계산 자원이 낭비되지 않도록 한다는 것을 확인했습니다. 하지만 Continuous Batching이 해결하지 못하는 두 번째 낭비 요인이 있습니다. 바로 기본적으로 이것이 동기적(synchronous)이라는 점입니다. 이는 CPU와 GPU가 교대로 작동함을 의미합니다. 즉, GPU가 계산하는 동안 CPU는 기다리고, CPU가 다음 배치를 준비하는 동안 GPU는 기다립니다. 초당 수백 단계가 실행되는 루프에서 이러한 유휴 간격(idle gaps)은 누적되며, 우리가 보여주겠지만 이는 전체 실행 시간의 거의 4분의 1을 차지할 수 있습니다. GPU가 100%의 시간 동안 계산에 매진할 수 있도록 하려면, 이러한 간격을 제거해야 합니다.

이를 달성하기 위해 우리는 **비동기 배치(asynchronous batching)**를 사용할 수 있습니다. CPU의 배치 준비와 GPU의 배치 계산을 분리하여, 두 작업이 병렬로 실행될 수 있도록 하고 항상 생산적인 GPU 상태를 유지할 것입니다 🔥

단순한 동기식 배치(naive synchronous batching)의 작동 방식은 다음과 같습니다:

CPU가 새로운 배치를 준비할 때, 어떤 요청(requests)을 포함할지 선택하고, KV cache 테이블을 업데이트하며, 이전 실행에서 종료된 요청을 제거(evict)하고, 확보된 공간을 채우기 위해 새로운 요청을 수락합니다. 이 작업이 완료되면 준비된 입력을 GPU로 전송합니다. GPU는 순전파(forward pass)를 실행하고 각 요청에 대해 새로운 토큰을 샘플링(즉, 선택)합니다. 결과가 CPU로 돌아오면 CPU는 각 요청이 방금 어떤 토큰을 생성했는지 알게 되며, 이후 전체 사이클이 다시 반복됩니다.

오른쪽의 빨간색 주석에 주목하십시오: GPU가 계산을 마치면 유휴(idle) 상태가 됩니다. CPU가 업데이트 단계(출력 토큰 샘플링, 요청 상태 업데이트, 배치 재스케줄링)를 거치기 전까지는 다음 배치가 시작될 수 없습니다.

이것이 동기식 배치(synchronous batching)의 핵심적인 비효율성입니다: CPU와 GPU가 차례를 주고받습니다. GPU가 계산하는 동안 CPU는 유휴 상태입니다. CPU가 업데이트하는 동안 GPU는 유휴 상태입니다. 어떤 상황에서도 두 장치가 동시에 유용한 작업을 수행하지 않습니다. 단일 순전파(forward pass)에서는 이것이 지불할 만한 작은 대가처럼 보일 수 있지만, 초당 수백 단계를 실행하는 연속 배치(continuous batching) 루프에서는 이러한 유휴 간극이 누적되어 실제 처리량(throughput) 손실로 이어집니다.

이를 보여주기 위해, 8B 모델을 사용하여 배치 크기(batch size) 32로 8K 토큰을 생성할 때 CPU와 GPU에서 소비되는 시간을 프로파일링(profile)했습니다.

동일한 종류의 그래프를 생성하고 싶다면, 연속 배치 코드를 인스트루먼테이션(instrument)하여 CPU 및 GPU 활동 구간(spans)을 덤프하고 이 스크립트를 사용할 수 있습니다.

타임라인은 초록색(GPU 활성, CPU 유휴)과 빨간색(CPU 활성, GPU 유휴) 사이를 교차하며, 두 색상은 절대 겹치지 않습니다. 총 생성 시간은 300.6초이며, 그중 24.0%는 GPU가 CPU가 완료되기를 기다리며 유휴 상태로 보낸 시간입니다. GPU의 관점에서 볼 때, 전체 생성 시간의 거의 4분의 1이 낭비되고 있습니다. 이것은 상황을 비관적으로 바라본 방식입니다.

낙관적인 방식은, 만약 우리가 CPU 오버헤드(overhead)를 완전히 제거할 수 있다면 생성 시간이 300초에서 228초로 줄어든다는 것입니다(무료로 24%의 속도 향상!). 이는 새로운 커널(kernel)이나 모델 변경 없이, 하드웨어의 세심한 조정만으로 가능합니다.

근본적으로 아이디어는 간단합니다: 배치 N이 계산되는 동안 배치 N+1을 위한 배치 준비를 어떻게 실행할지 알아내야 합니다. 하지만 이 간단한 아이디어 뒤에는 몇 가지 기술적인 어려움이 숨어 있습니다:

• 어떻게 GPU에서 무언가를 실행하고 CPU로 제어권을 다시 가져올 수 있을까요?
• 각 작업(task)이 실행될 때, CPU 또는 GPU 작업을 위한 데이터가 준비되어 있음을 어떻게 보장할 수 있을까요?
• 배치 N+1이 배치 N의 예측 결과에 기반한다면, 어떻게 배치 N+1을 준비할 수 있을까요?

이 질문들에 답함으로써, 우리는 비동기 배치(asynchronous batching)를 처음부터 구축해 나갈 것입니다. 우리는 transformers 라이브러리의 continuous batching의 일부로 이를 구현하기 위해 동일한 단계를 따랐습니다. 코드를 확인하고 비교해 보세요!

우리의 최종 목표는 CPU와 GPU 작업의 동시(concurrent) 실행을 구현하는 것입니다. 우리는 어떤 작업이 동시에 실행될 수 있는지 머신에 알려줄 수 있도록, 작업을 분류할 방법이 필요합니다. 우리는 이를 CUDA 스트림(CUDA streams)을 사용하여 달성할 수 있습니다.

CUDA가 작업을 어떻게 순서화하는지 이해하려면, **CUDA 스트림 (CUDA streams)**에 대해 이야기해야 합니다. 스트림은 제출된 순서대로 실행되는 GPU 작업(커널 실행(kernel launches), 메모리 복사(memory copies), 동기화 장벽(synchronization barriers))의 순차적인 큐(ordered queue)입니다. 모든 GPU 작업은 항상 스트림 내부에서 스케줄링됩니다. 동일한 스트림 내의 작업들은 순차적입니다: GPU는 이전 작업이 완료될 때까지 다음 작업을 시작하지 않습니다. 서로 다른 스트림에 있는 작업들은 서로 독립적이며 동시에 실행될 수 있습니다. 이를 설명하기 위해, 만약 3개의 서로 다른 스트림에 걸쳐 3개의 작업을 실행한다면, 실행 모습은 다음과 같습니다:

세 작업 모두 동시에 시작됩니다. 이는 약간의 단순화된 설명입니다: 모든 GPU 작업은 궁극적으로 CPU에 의해 시작되며, 그 시작에는 약간의 시간이 소요됩니다: 적절한 커널을 찾고, 호출을 발행하며, 명령을 CPU에서 GPU로 전달하는 등의 과정입니다. 이를 **CPU 실행 오버헤드 (CPU launch overhead)**라고 하며, 더 현실적인 다이어그램은 다음과 같습니다:

연산들은 여전히 병렬적으로(concurrent) 수행되지만, 각 CPU 런치(CPU launch) 비용만큼 시작 시간이 뒤로 밀리게 됩니다. 우리는 이러한 CPU 런치 이벤트가 실제 시간을 소모하기 때문에, 비동기 워크플로(asynchronous workflows)로 넘어갈 때 "언제 무엇이 런치되는지"를 추적하는 데 도움이 되도록 계속해서 이를 보여줄 것입니다. 예를 들어, 우리는 스트림(stream)이 플러시(flushed) 되었는지 자주 확인할 것입니다. 이는 스트림 내의 모든 연산이 실행되었음을 의미합니다.

만약 PyTorch에서 CUDA 스트림을 명시적으로 사용해 본 적이 없다면, 스트림이 존재한다는 사실 자체에 놀랄 수도 있습니다. 일반적인 PyTorch 스크립트는 스트림을 전혀 언급하지 않으며, GPU 연산이 비동기적으로 작동한다는 느낌을 주지 않습니다. CPU가 다음 단계로 넘어가기 전에 GPU가 작업을 마칠 때까지 기다리는 것처럼 보이기 때문입니다. 그러한 느낌은 정확하며, 이는 기본 스트림 (default stream) 때문에 발생합니다.

스트림을 지정하지 않고 PyTorch 연산을 호출하면 해당 연산은 기본 스트림에 할당됩니다. 기본 스트림에는 한 가지 특별한 속성이 있는데, 바로 **동기화(synchronizing)**를 한다는 점입니다. 만약 어떤 연산이 기본 스트림에 스케줄링되면, 해당 연산은 **모든 다른 스트림이 플러시(flushed)**될 때까지 기다립니다. 즉, 기본 스트림의 단일 연산이 시작되기 전에 GPU의 모든 작업이 완료되어야 합니다. 그 반대도 마찬가지입니다. 어떤 연산이든 스트림에 관계없이, 실행되기 전에 기본 스트림이 플러시될 때까지 기다립니다.

따라서 기본 스트림 연산의 결과를 CPU로 전송할 때, 설령 CPU에 대해 비차단(non-blocking) 방식으로 설계된 전송이라 할지라도, 연산들이 기본 스트림에 스케줄링되었기 때문에 모든 GPU 연산이 끝날 때까지 CPU는 여전히 차단(block)됩니다. 이는 병렬성(concurrency)을 구축하려는 모든 노력을 사실상 무너뜨립니다.

이것이 바로 우리가 기본 스트림이 아닌 스트림을 사용해야 하는 이유입니다. 커널 런치(kernel launch)나 비차단 메모리 복사(non-blocking memory copy)를 인큐(enqueue)하면 즉시 CPU로 제어권이 반환됩니다. GPU는 백그라운드에서 연산을 실행하지만, CPU는 기다리지 않습니다. 이것이 우리의 첫 번째 질문에 대한 답입니다. GPU 작업을 런치한 후 CPU 제어권을 다시 가져오기 위해서는, 기본 스트림이 아닌 스트림을 사용해야 합니다.

이 포스트의 나머지 부분에서는 한 장치에서 다른 장치로의 모든 메모리 전송(memory transfers)이 비차단(non-blocking) 방식이라고 가정하겠습니다. 따라서 우리는 이를 직접 동기화(synchronize)해야 합니다.

우리는 어떤 GPU 작업도 기본 스트림(default stream)에 할당되어서는 안 된다는 점을 확인했습니다. 하지만 질문은 여전히 남아 있습니다. 기본 스트림을 사용하지 않는다면, 어떤 스트림을 사용해야 할까요? 동기식 배치(synchronous batching) 그림으로 다시 돌아가 보겠습니다.

우리는 세 가지 별개의 GPU 작업을 식별할 수 있습니다:

• CPU에서 GPU로의 입력 전송 (Transfer of inputs from CPU to GPU)
• GPU에서의 연산 (Compute on the GPU)
• GPU에서 CPU로의 출력 전송 (Transfer of outputs from the GPU to the CPU)

이는 세 개의 스트림이 필요함을 의미합니다. 즉, 연산용 스트림 하나, CPU-to-GPU 전송용 스트림 하나, 그리고 GPU-to-CPU 전송용 스트림 하나가 필요합니다. 전송 작업들은 독립적이므로 이를 직렬화(serialize)할 이유가 없으며, 각 작업은 자신만의 스트림을 가집니다.

명칭에 관한 참고 사항: CPU와 GPU에 대해 이야기할 때, CUDA 문서 전반에서 사용되는 관례는 CPU를 host라 부르고 GPU를 device라고 부르는 것입니다. 우리는 앞으로 이 관례를 따를 것입니다. CPU-to-GPU 전송은 host-to-device (H2D) 전송이라 부르고, GPU-to-CPU 전송은 device-to-host (D2H) 전송이라 부릅니다. 따라서 세 개의 스트림은 H2D 스트림, 연산(compute) 스트림, 그리고 D2H 스트림입니다.

이제 스트림을 사용하여 GPU에서 배치를 비동기적으로 런치(launch)하고 CPU 제어권을 다시 가져오는 시도를 해보겠습니다. CPU에서는 다음과 같은 작업을 수행합니다:

• CPU에서 배치 입력 데이터 준비 (스트림 없음, CPU 전용 작업)
• GPU로 데이터 전송 (H2D 스트림 사용)
• GPU에서 연산 실행 (연산 스트림 사용)
• 배치 출력 결과 회수 (D2H 스트림 사용)
• 결과 확인 (스트림 없음)

만약 오직 CUDA 스트림만을 사용하여 이 과정을 수행한다면, 결과는 거의 즉시 나타나지만 잘못된 값일 것입니다. 왜 그런지 이해하기 위해, 어떤 일이 일어났는지 살펴보겠습니다:

스트림(streams)은 서로 독립적이기 때문에, 세 가지 GPU 연산이 거의 동시에 실행되었습니다. 연산(compute) 스트림은 H2D 전송이 완료될 때까지 기다리지 않았으므로, 순전파(forward pass)는 GPU 메모리에 이미 들어있던 데이터를 사용하여 실행되었습니다. D2H 스트림은 연산이 끝날 때까지 기다리지 않았으므로, 아직 계산되지 않은 결과를 전송했습니다. 5단계는 CPU를 차단(blocking)하는 요소가 없었기 때문에 즉시 반환되었습니다. 즉, 동기화할 기본 스트림(default stream)이 없었던 것입니다.

각 연산은 개별적으로는 모두 올바르게 실행되고 있습니다. 문제는 우리가 스트림들에게 서로를 기다리라고 명령하지 않았다는 점입니다. 우리는 H2D가 완료된 후에 연산이 시작되어야 하며, 연산이 완료된 후에 D2H가 시작되어야 한다는 것을 알고 있지만, 그 순서를 강제하지 않았습니다. 스트림 경계(stream boundaries)를 가로질러 "저 연산이 끝날 때까지 이 연산을 시작하지 마라"라고 말할 수 있는 메커니즘이 필요합니다.

스트림 간의 동기화(synchronization)를 강제하기 위해, 우리는 **CUDA 이벤트 (CUDA events)**를 사용할 것입니다.

A CUDA 이벤트는 스트림에 기록될 수 있는 마커(marker)입니다. GPU가 실행 중에 해당 마커에 도달하면, 해당 이벤트를 완료된 것으로 설정합니다. 그러면 다른 어떤 스트림이든 다음 연산을 시작하기 전에 해당 이벤트가 완료될 때까지 기다리도록 설정할 수 있습니다. 구체적으로 두 가지 연산이 있습니다: 현재 위치에 스트림에 마커를 삽입하는 stream.record(event), 그리고 이벤트가 완료로 표시될 때까지 스트림의 진행을 차단하는 stream.wait(event)입니다. 중요한 점은, wait는 CPU나 병렬로 실행 중인 다른 스트림이 아니라 스트림을 차단한다는 것입니다. 즉, CPU 호출은 즉시 반환되며, 오직 기다리는 스트림만 멈춰 있게 됩니다.

위 그림은 하나의 이벤트가 두 스트림을 동기화하는 모습을 보여줍니다. CPU는 세 가지 작업(세 개의 작은 블록)을 빠르게 연속해서 발행합니다: 스트림 1(stream 1)에서 입력 준비(input preparation)를 시작하고, 스트림 1에서 이벤트를 기록(record)한 다음, 스트림 2(stream 2)에게 이를 기다리라고 명령합니다. 그 후 CPU는 즉시 다음 작업으로 넘어갑니다. 스트림 1은 자신의 작업을 실행하며, 작업이 완료되면 이벤트가 설정(set)됩니다. 스트림 2는 대기 마커(wait marker)에서 계속 머물러 있다가, 이벤트가 완료로 표시되는 순간에만 연산(compute)을 시작합니다. CPU는 이 과정에 전혀 관여하지 않았습니다. 순서 제어는 전적으로 GPU 측에서 강제되었습니다.

우리의 사례에 적용하면, 해결 방법은 간단합니다. H2D 전송을 인큐(enqueue)한 후, h2d_stream.record(h2d_done)를 호출합니다. 이렇게 하면 전송이 완료될 때만 이벤트가 완료된 것으로 표시됩니다. 순전파(forward pass)를 인큐하기 전에 compute_stream.wait(h2d_done)를 호출하여, h2d_done이 설정될 때까지 연산 스트림이 시작되지 않도록 합니다. 연산과 D2H 사이에서도 동일하게 수행합니다: model.forward로 순전파를 실행한 후, compute_stream.record(compute_done)를 호출하고, 출력 전송을 인큐하기 전에 d2h_stream.wait(compute_done)를 호출합니다. 그 결과 명시적인 순서가 있는 파이프라인이 만들어집니다:

• H2D 전송은 h2d_stream에서 실행됩니다.
• compute_stream은 h2d_done을 기다린 후 순전파를 실행합니다.
• d2h_stream은 compute_done을 기다린 후 출력을 다시 전송합니다.

CPU는 이 모든 것을 순차적으로 인큐한 다음 바로 다음 단계로 넘어갑니다. 어떤 시점에서도 CPU는 차단(block)되지 않습니다. GPU는 이벤트를 통해 순서를 강제하며, 세 스트림 모두 의존성(dependency)이 충족되는 즉시 활성화됩니다.

위 그림은 이 과정이 어떻게 전개되는지 보여줍니다. CPU는 배치를 준비한 다음, record와 wait를 사용하여 H2D 전송, 순전파, D2H 전송 등 모든 GPU 작업을 빠르게 인큐합니다.