[vllm] vLLM Gemma4 모델의 GPU/CPU 동기화 병목 현상 해결하기: non_blocking 전송의 중요성

PR 링크: vllm-project/vllm#39234 상태: Merged | 변경: +None / -None

들어가며

고성능 LLM 추론 엔진인 vLLM 프로젝트에서 최근 Gemma4 모델의 성능을 한 단계 끌어올린 중요한 최적화가 진행되었습니다. 이번 포스트에서는 embed_input_ids 함수 내에서 발생하던 GPU/CPU 동기화(Sync) 병목 현상을 어떻게 발견하고, 단 한 줄의 코드 수정으로 이를 해결했는지 깊이 있게 살펴보겠습니다.

딥러닝 모델의 추론 속도는 단순히 연산량(FLOPs)에만 의존하지 않습니다. CPU와 GPU 사이의 데이터 흐름이 원활하지 않아 발생하는 'Wait' 시간은 전체 처리량(Throughput)을 갉아먹는 주범입니다. 특히 멀티모달 모델인 Gemma4의 경우, 텍스트와 이미지 데이터를 처리하는 과정에서 발생하는 데이터 전송 방식이 성능의 핵심이 됩니다.

문제 상황: 암시적인 동기화의 함정

기존 Gemma4 구현에서는 멀티모달 입력을 처리할 때, 특정 토큰이 이미지인지 텍스트인지 구분하는 is_multimodal 마스크를 사용합니다. 문제는 이 마스크 텐서가 CPU에 머물러 있는 상태에서 GPU 연산에 참여하기 위해 디바이스를 이동할 때 발생했습니다.

리뷰어인 lgeiger가 언급했듯이, 이전 PR(#34246) 이후 is_multimodal 텐서는 의도적으로 CPU에 유지되도록 설계되었습니다. 하지만 이를 GPU 텐서인 input_ids와 함께 torch.where 연산에 사용할 때, 명시적인 비동기 설정이 없으면 PyTorch는 해당 전송이 완료될 때까지 CPU 호스트를 블로킹(Blocking)하게 됩니다.

코드 분석: Before & After

문제가 된 vllm/model_executor/models/gemma4_mm.py 파일의 변경 사항을 살펴보겠습니다.

Before

# vllm/model_executor/models/gemma4_mm.py

if is_multimodal is not None:
    is_multimodal = is_multimodal.to(input_ids.device)
    ple_input_ids = torch.where(
        is_multimodal, torch.zeros_like(input_ids), input_ids
    )

기존 코드에서는 is_multimodal.to(input_ids.device)를 호출할 때 별도의 옵션을 주지 않았습니다. 이 경우 기본값은 non_blocking=False입니다. 즉, CPU는 is_multimodal 텐서가 GPU 메모리로 완전히 복사될 때까지 다음 명령을 실행하지 못하고 기다리게 됩니다. 이는 GPU 커널이 실행될 준비가 되었음에도 불구하고 CPU가 명령을 늦게 내려 GPU가 노는 시간을 발생시킵니다.

After

# vllm/model_executor/models/gemma4_mm.py

if is_multimodal is not None:
    # is_multimodal = is_multimodal.to(input_ids.device) <- 제거됨
    ple_input_ids = torch.where(
        is_multimodal.to(input_ids.device, non_blocking=True),
        torch.zeros_like(input_ids),
        input_ids,
    )

개선된 코드에서는 두 가지 변화가 있습니다.

1. non_blocking=True 사용: 텐서 전송을 비동기적으로 요청합니다. CPU는 전송이 시작되자마자 바로 다음 코드로 넘어가며, 실제 연산은 GPU 스트림에서 전송이 완료된 후 순차적으로 진행됩니다.
2. 인라인 처리: 별도의 할당 단계 대신 torch.where 내부에서 직접 전송을 수행하여 불필요한 중간 단계와 잠재적인 동기화 지점을 최소화했습니다.

왜 이게 좋은 최적화인가?

1. GPU Utilization 극대화

PR에 첨부된 프로파일링 결과를 보면 차이가 명확합니다. 'Before' 상태에서는 CPU와 GPU 사이의 동기화로 인해 타임라인에 빈 공간(Gap)이 존재하지만, 'After' 상태에서는 커널들이 끊김 없이 연속적으로 실행되는 것을 확인할 수 있습니다. 이는 특히 작은 배치 사이즈나 지연 시간(Latency)이 중요한 실시간 서비스에서 큰 차이를 만듭니다.

2. Host-Device Overlap

non_blocking=True를 사용하면 데이터 전송(Copy)과 CPU의 다른 로직 수행이 겹쳐질(Overlap) 수 있습니다. is_multimodal 텐서가 고정된 메모리(Pinned Memory)에 있다면, 이 효과는 더욱 극대화되어 데이터 전송 중에도 CPU는 다음 연산을 위한 준비를 마칠 수 있습니다.

3. 유지보수성

리뷰어의 피드백을 반영하여 is_multimodal이 CPU에 남아있어야 한다는 의도를 해치지 않으면서도, 연산 시점에만 효율적으로 GPU로 넘겨주는 깔끔한 구조를 유지했습니다.

교훈: 시니어 엔지니어의 관점

이번 최적화는 복잡한 알고리즘의 개선이 아니라, 프레임워크의 동작 원리에 대한 깊은 이해에서 비롯되었습니다.

• 프로파일링 우선: 막연히 코드를 고치는 대신, torch.profiler를 통해 실제 동기화 지점을 찾아냈습니다.
• 기본값의 위험성: to() 메서드의 기본값이 블로킹 방식이라는 점을 인지하고, 성능 민감도가 높은 추론 루프 내에서는 항상 비동기 옵션을 고려해야 합니다.
• 컨텍스트 유지: 리뷰어와의 소통을 통해 해당 변수가 왜 CPU에 있어야 하는지 파악하고, 그 제약 조건 내에서 최선의 해결책을 도출했습니다.

vLLM과 같은 대규모 프로젝트에서 이런 작은 디테일의 차이가 모여 세계 최고의 성능을 만들어냅니다. 여러분의 코드에서도 to(device) 뒤에 숨겨진 성능 저하 요인이 없는지 다시 한번 확인해 보시기 바랍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.to.html
• https://pytorch.org/docs/stable/generated/torch.where.html
• https://pytorch.org/docs/stable/notes/cuda.html#host-device-synchronization

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.