[vllm] vLLM 성능 최적화: cuMemcpyBatchAsync를 활용한 KV 캐시 스왑 효율화

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

들어가며

vLLM은 대규모 언어 모델(LLM) 서빙을 위한 고성능 엔진입니다. 모델의 KV 캐시를 CPU와 GPU 사이에서 효율적으로 이동시키는 것은 긴 문맥(long-context) 처리에 필수적입니다. 기존 vLLM 구현에서는 swap_blocks 함수를 통해 레이어별, 블록별로 데이터를 하나씩 복사했습니다. 하지만 이 방식은 수많은 작은 복사 요청을 드라이버에 개별적으로 전달해야 하므로, 드라이버 호출 오버헤드(submission overhead)가 전체 성능을 저하시키는 병목 현상을 일으켰습니다. 본 PR은 이 문제를 해결하기 위해 여러 복사 작업을 하나의 배치 작업으로 통합하는 swap_blocks_batch를 도입했습니다.

코드 분석

1. C++ 커널 구현 (csrc/cache_kernels.cu)

핵심 변경 사항은 CUDA 12.8에서 도입된 cuMemcpyBatchAsync를 활용하는 것입니다. 기존에는 루프를 돌며 cudaMemcpyAsync를 호출했지만, 이제는 배치 API를 사용하여 단 한 번의 드라이버 호출로 모든 복사를 처리합니다.

// Before: 개별 루프 호출
for (int64_t i = 0; i < n; i++) {
  cudaMemcpyAsync(..., stream);
}

// After: cuMemcpyBatchAsync 활용
CUresult result = cuMemcpyBatchAsync(
    reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
    reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
    reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
    static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx, 
    static_cast<CUstream>(stream));

2. Python 레이어 최적화 (vllm/v1/kv_offload/worker/cpu_gpu.py)

Python 단에서는 복사할 데이터의 포인터와 크기를 미리 계산하여 numpy 배열로 준비한 뒤, 이를 torch.from_numpy를 통해 C++ 영역으로 전달합니다. 이렇게 함으로써 Python 루프 오버헤드를 최소화하고 C++ 레벨에서 일괄 처리하게 됩니다.

# Pre-compute base pointers
self._src_base_ptrs = np.array([t.data_ptr() for t in self.src_tensors], dtype=np.int64)

# ... (중략) ...

# Batch call
if total > 0:
    ops.swap_blocks_batch(batch_src, batch_dst, batch_sizes)

왜 이게 좋은가

이 최적화는 특히 작은 블록 단위의 복사가 빈번할 때 극적인 효과를 발휘합니다. 벤치마크 결과에 따르면, Qwen2.5-0.5B 모델에서 7.4배의 속도 향상을 보였으며, E2E 서빙 환경에서도 처리량(throughput)이 최대 18% 증가했습니다.

핵심 교훈:

1. 드라이버 호출 오버헤드 최소화: GPU 작업 시 개별 커널/복사 호출은 생각보다 큰 오버헤드를 가집니다. 가능하다면 배치 API를 사용하는 것이 좋습니다.
2. 하드웨어 버전 대응: cuMemcpyBatchAsync는 최신 CUDA 버전에서만 지원되므로, 이전 버전과의 호환성을 위해 if defined(CUDA_VERSION) && CUDA_VERSION >= 12080과 같은 전처리기 매크로를 통한 폴백(fallback) 로직이 필수적입니다.

리뷰어 피드백과 논의

리뷰 과정에서 CUDA 12.8 미만 환경에서 undefined symbol: cuMemcpyBatchAsync 오류가 발생하는 이슈가 제기되었습니다. 이는 런타임 환경과 드라이버 버전의 불일치로 인해 발생하며, 현재는 컴파일 타임 체크를 통해 이를 방지하고 있습니다. 향후 cuGetProcAddress를 사용하여 런타임에 API 존재 여부를 확인하는 더 방어적인 접근 방식도 논의되었습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f
• https://pytorch.org/docs/stable/generated/torch.Tensor.data_ptr.html

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