연속 배치에서 비동기성 활용하기

동기식 배치에서는 GPU가 1/4 시간 동안 유휴 상태

동기식 연속 배치에서는 CPU와 GPU가 번갈아 가며 작업합니다: GPU가 순전파를 수행하는 동안 CPU는 대기하고, CPU가 다음 배치를 준비하는 동안 GPU는 대기합니다. 배치 크기 32로 8K 토큰을 생성하는 8B 모델을 프로파일링한 결과 전체 시간이 300.6초였으며, 그 중 24.0%가 CPU를 기다리는 GPU 유휴 시간으로 사용되었습니다. 이 유휴 간격은 연속 배치만으로는 해소되지 않는 처리량 손실을 의미합니다.

동시성을 만들려면 비기본 CUDA 스트림이 필요합니다

CPU와 GPU 작업을 겹치려면 CPU를 차단하지 않고 GPU 연산을 시작해야 합니다. PyTorch에서는 명시적인 스트림이 지정되지 않은 연산은 기본 스트림에서 실행되며, 이 스트림은 모든 다른 스트림과 동기화되어 CPU가 GPU가 작업을 마칠 때까지 기다리게 합니다. 비기본 스트림을 사용하면 CPU가 작업을 큐에 넣고 즉시 제어권을 되찾을 수 있습니다. 우리는 세 개의 별도 스트림이 필요합니다: 호스트‑투‑디바이스(H2D) 전송용, GPU 연산용, 디바이스‑투‑호스트(D2H) 전송용 스트림입니다. 이러한 연산은 독립적이며 서로 다른 스트림에 배치하면 동시에 실행될 수 있습니다.

스트림 간 순서를 강제하려면 CUDA 이벤트가 필요합니다

세 스트림을 독립적으로 시작하면 경쟁 상태가 발생합니다: H2D 전송이 끝나기 전에 연산이 시작될 수 있고, D2H가 연산이 끝나기 전에 결과를 전송할 수 있습니다. CUDA 이벤트는 스트림에 마커를 기록하며, 다른 스트림은 진행하기 전에 해당 마커를 기다릴 수 있습니다. 각 H2D 전송 후에 이벤트를 기록하고 연산 스트림에서 이를 기다리며, 연산 후에도 이벤트를 기록하고 D2H 스트림에서 기다리게 하면, GPU가 올바른 순서를 강제하고 CPU는 모든 작업을 큐에 넣은 뒤 자유롭게 남게 되는 파이프라인을 만들 수 있습니다.

데이터 손상을 방지하고 캐리‑오버를 처리하려면 이중 버퍼링과 마스크가 필요합니다

배치 N이 연산 중일 때 배치 N+1을 준비하면 GPU가 아직 읽고 있는 입력 버퍼를 덮어쓸 위험이 있습니다. 해결책은 호스트와 디바이스 측 텐서 두 세트를(슬롯 A와 B) 유지하고 교대로 사용하는 것으로, 메모리 사용량이 두 배가 되지만 경쟁 상태를 방지합니다. 두 배치에 모두 나타나는 요청은 다음 배치의 입력으로 새로 생성된 토큰이 필요하므로, 배치 N+1 입력 버퍼에 자리표시자 토큰(값 0)을 삽입하고 나중에 배치 N 출력에서 실제 토큰으로 교체합니다. 이 교체를 캐리‑오버라고 하며, 목적지 위치를 지정하는 캐리‑오버 마스크에 의해 수행됩니다; 마스크는 CUDA 그래프 내부에서 거의 비용 없이 적용됩니다.

전체 비동기 루프는 CPU 배치 준비와 GPU 연산을 겹칩니다

Step 0(콜드 스타트)에서는 동기 경로를 사용해 배치 0을 디스패치합니다. Step 1부터는 GPU가 슬롯 A에서 배치 N을 연산하는 동안 CPU가 슬롯 B에서 배치 N+1을 준비합니다: 완료된 요청을 제거하고, 새로운 요청을 받아들이며, KV 캐시를 업데이트하고, 캐리‑오버 마스크를 구축합니다. 입력이 준비되면 CPU는 H2D 스트림에 H2D 전송을 큐에 넣고, 이벤트를 기록한 뒤 연산 스트림에서 이를 기다립니다. 그런 다음 순전파를 실행하고 또 다른 이벤트를 기록한 뒤 D2H 스트림에서 이를 기다린 후 출력 전송을 시작합니다. CPU는 결과를 복사하기 위한 최종 동기화 이벤트에서만 차단되며, 이후 즉시 다음 배치를 준비합니다. 이 패턴이 반복되어 배치 N이 끝날 때 배치 N+1의 입력이 준비되어 있는 한 GPU는 계속 바쁘게 유지됩니다.

실험 결과 22% 속도 향상을 보여줍니다

동기식 기준(8K 토큰, 배치 크기 32, 8B 모델)과 동일한 작업을 비동기 파이프라인으로 실행하면 CPU와 GPU가 거의 지속적으로 겹치는 타임라인이 얻어집니다. GPU는 전체 실행 시간의 99.4% 동안 활성화되어, 동기식 경우의 76.0%에서 크게 증가했습니다. 전체 생성 시간은 300.6초에서 234.5초로 감소했으며, 약 22%의 속도 향상을 나타냅니다. 남은 차이는 CPU가 D2H 전송이 완료될 때까지 기다야 하는 불가피한 동기화 지점에 해당하며, 새로운 커널이나 모델 변경은 필요하지 않습니다.

결론

CUDA 스트림과 이벤트를 사용해 스케줄 기반 의존성을 데이터 기반 의존성으로 교체하고, 이중 버퍼링과 캐리‑오버 마스크로 경쟁 상태를 해결함으로써 연속 배치에서 CPU와 GPU 작업을 분리합니다. 이를 통해 GPU가 포화 상태를 유지하게 되어 LLM 추론의 처리량을 크게 향상시키면서 모델 정확도는 유지됩니다.