[vllm] vLLM의 성능 병목 해결: Host-to-Device 복사 최적화로 비동기 실행 보장하기

PR 링크: vllm-project/vllm#51841 상태: Merged | 변경: +23 / -5

들어가며

고성능 추론 엔진인 vLLM을 운영할 때, GPU 연산 파이프라인의 핵심 경로에서 발생하는 지연 시간은 전체 처리량(throughput)에 치명적인 영향을 미칩니다. 특히 non_blocking=True 옵션을 사용했음에도 불구하고 cudaMemcpyAsync 호출이 호출 스레드를 차단(blocking)하는 현상이 발생한다면, 이는 비동기 스케줄링의 이점을 완전히 상쇄합니다. 본 글에서는 vLLM의 ViT(Vision Transformer) 모델 처리 과정에서 발견된 호스트-디바이스(H2D) 복사 병목 현상을 어떻게 해결했는지 살펴봅니다.

코드 분석

이번 최적화는 크게 두 가지 영역에서 이루어졌습니다.

1. M-RoPE 포지션 데이터 복사 최적화 (vllm/v1/worker/gpu_model_runner.py)

기존 코드에서는 mrope_positions 텐서가 torch.compile을 위해 의도적으로 비연속적(non-contiguous)인 구조로 할당되어 있었습니다. 이로 인해 cpu[:, :N] 슬라이싱이 스트라이드된 뷰(strided view)를 생성하게 되었고, copy_() 호출 시 내부적으로 페이지 가능한 임시 버퍼로 데이터를 모으는 과정에서 동기화가 발생했습니다.

Before:

self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(
    self.mrope_positions.cpu[:, :total_num_scheduled_tokens],
    non_blocking=True,
)

After:

for row in range(self.mrope_positions.gpu.shape[0]):
    self.mrope_positions.gpu[row, :total_num_scheduled_tokens].copy_(
        self.mrope_positions.cpu[row, :total_num_scheduled_tokens],
        non_blocking=True,
    )

각 행(row)은 메모리상에서 연속적이기 때문에, 행 단위로 나누어 복사함으로써 pinned memory의 이점을 살리고 불필요한 동기화를 제거했습니다.

2. Vision Position IDs 복사 최적화 (vllm/model_executor/models/qwen3_vl.py)

pos_ids를 생성하는 과정에서 데이터가 pinned memory에 있지 않아 H2D 복사 시 동기화가 발생했습니다. 이를 torch.empty(..., pin_memory=True)를 사용하여 미리 할당된 핀 메모리에 직접 cat 연산을 수행하도록 개선했습니다.

Before:

pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True)

After:

pinned = torch.empty(
    (num_pos, pos_ids[0].shape[1]),
    dtype=pos_ids[0].dtype,
    pin_memory=True,
)
pos_ids = torch.cat(pos_ids, dim=0, out=pinned).to(
    self.device, non_blocking=True
)

왜 이게 좋은가

이번 최적화의 핵심은 '비동기 실행의 보장'입니다. cudaMemcpyAsync는 소스 메모리가 pinned memory이고 데이터가 연속적일 때만 진정한 비동기 동작을 수행합니다.

1. 스트라이드 제거: 비연속적인 텐서 슬라이싱은 드라이버 레벨에서 추가적인 복사를 유발하여 non_blocking 옵션을 무력화합니다. 행 단위 복사로 이를 우회했습니다.
2. Pinned Memory 활용: pin_memory=True를 통해 호스트 메모리를 페이지 고정하여 DMA(Direct Memory Access) 전송이 가능하게 함으로써 CPU의 개입을 최소화했습니다.

이러한 개선을 통해 프로파일링 결과, 기존에 발생하던 긴 차단(long-blocking) 현상이 완전히 제거되었습니다. 일반적인 교훈은 "비동기 복사를 사용할 때는 데이터의 연속성(Contiguity)과 핀 메모리(Pinned Memory) 여부를 반드시 확인해야 한다"는 것입니다.

리뷰 피드백 반영

njhill은 torch.cat의 out 파라미터를 활용하여 추가적인 호스트 측 메모리 복사를 방지하는 최적화 제안을 주었습니다. 이는 메모리 할당과 복사 횟수를 줄여 성능을 극대화하는 좋은 예시입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.pin_memory.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html

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