[vllm] vLLM의 MoE Permute 최적화: 버퍼 사전 할당을 통한 성능 향상

PR 링크: vllm-project/vllm#43014 상태: Merged | 변경: +446 / -58

들어가며

vLLM은 대규모 언어 모델(LLM) 서빙을 위한 고성능 엔진입니다. 특히 Mixture-of-Experts(MoE) 모델을 처리할 때, 토큰을 각 전문가(Expert)에게 할당하기 위해 permute 및 unpermute 연산이 필수적입니다. 기존 구현에서는 매 연산마다 임시 버퍼를 torch.empty로 할당하고 해제하는 과정이 반복되었는데, 이는 특히 배치 사이즈가 작은 경우 오버헤드로 작용했습니다. 본 PR은 이러한 메모리 할당 오버헤드를 제거하기 위해 scratch 버퍼를 사전에 할당하여 재사용하는 방식을 도입했습니다.

코드 분석

1. C++ 커널 레벨의 변경 (csrc/moe/moe_permute_unpermute_op.cu)

가장 핵심적인 변화는 maybe_allocate_tensor 헬퍼 함수를 도입하여, 이미 할당된 버퍼가 있다면 이를 재사용하고, 없다면 새로 할당하는 로직을 구현한 것입니다.

// Before: 매번 새로운 텐서 할당
auto sort_workspace = torch::empty({sorter_size}, torch::dtype(torch::kInt8).device(torch::kCUDA));

// After: maybe_allocate_tensor를 통한 재사용
auto sort_workspace = maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8, device, "sort_workspace");

또한, moe_permute_impl 함수를 통해 기존 moe_permute 인터페이스를 유지하면서, 선택적으로 scratch 버퍼를 받을 수 있도록 확장했습니다. 이는 기존 코드와의 하위 호환성을 보장하면서 성능 최적화를 적용하는 좋은 패턴입니다.

2. 파이썬 레이어의 변경 (vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py)

파이썬 단에서는 MoEPermuteScratch 클래스를 도입하여 필요한 버퍼들을 관리합니다. 벤치마크 코드에서는 다음과 같이 사용됩니다.

# After: scratch 객체를 생성하여 전달
scratch = MoEPermuteScratch(
    max_num_tokens=num_tokens,
    topk=topk,
    num_experts=num_experts,
    # ...
)

# 연산 시 scratch 전달
moe_permute(..., scratch=scratch)

왜 이게 좋은가

성능 향상

벤치마크 결과에 따르면, 배치 사이즈가 작을수록(18) 약 914%의 성능 향상이 관찰되었습니다. 이는 GPU 커널 실행 시간에서 메모리 할당 및 동기화 오버헤드가 차지하는 비중이 작은 배치에서 매우 크다는 것을 의미합니다.

Batch Size	Improvement
1	14.44%
8	10.13%

교훈

1. 메모리 재사용의 중요성: 빈번하게 호출되는 커널에서 cudaMalloc과 같은 메모리 할당은 성능의 병목이 됩니다. 가능한 경우 workspace 버퍼를 외부에서 관리하여 재사용하는 것이 좋습니다.
2. 하위 호환성 유지: std::optional과 같은 C++ 기능을 활용하여 기존 API를 깨뜨리지 않고 최적화 경로를 추가하는 설계는 대규모 프로젝트에서 필수적입니다.
3. 방어적 프로그래밍: 리뷰 과정에서 논의되었듯이, 재사용되는 버퍼의 dtype, device, shape을 검증하는 로직(maybe_allocate_tensor 내부의 TORCH_CHECK)을 포함하여 안정성을 확보했습니다.

결론

이번 최적화는 vLLM이 소규모 배치 환경에서도 효율적으로 동작하도록 돕는 중요한 개선입니다. 특히 실시간 서빙 환경에서 지연 시간(Latency)을 줄이는 데 큰 기여를 할 것으로 기대됩니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.empty.html
• https://pytorch.org/docs/stable/cpp_api.html

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