[vllm] vLLM MoE 성능 최적화: FlashInfer One-Sided Combine을 활용한 메모리 복사 제거

PR 링크: vllm-project/vllm#47156 상태: Merged | 변경: +61 / -2

들어가며

vLLM의 Mixture-of-Experts(MoE) 연산 과정에서, 분산 환경의 Expert 출력들을 하나로 합치는(combine) 과정은 성능 병목 중 하나였습니다. 기존에는 FlashInfer의 combine 함수가 새로운 텐서를 반환하면, 이를 vLLM이 미리 할당해둔 출력 버퍼로 다시 copy_()하는 비효율적인 과정이 포함되어 있었습니다. 본 PR은 FlashInfer의 최신 기능을 활용하여, combine 연산이 vLLM의 출력 버퍼에 직접 결과를 쓰도록(Direct Write) 변경함으로써 메모리 대역폭 낭비를 줄이고 성능을 향상시켰습니다.

코드 분석

1. vllm/distributed/device_communicators/all2all.py

핵심 변경 사항은 FlashInferNVLinkOneSidedManager 클래스에 combine_into 메서드를 추가한 것입니다. 이 메서드는 FlashInfer의 combine 함수가 output 인자를 지원하는지 여부를 사전에 확인하여, 지원할 경우 직접 쓰기를 수행하고, 그렇지 않으면 기존의 copy_() 방식을 사용하는 폴백(fallback) 로직을 포함합니다.

# Before (기존 방식)
combined_output = self.moe_alltoall.combine(...)
output.copy_(combined_output)

# After (개선된 방식)
def combine_into(self, payload, runtime_max_tokens_per_rank, output):
    if self._combine_supports_output:
        self.moe_alltoall.combine(..., output=output)
    else:
        combined_output = self.moe_alltoall.combine(...)
        output.copy_(combined_output)

2. vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py

finalize 메서드에서는 이제 복잡한 조건문 없이 추상화된 combine_into를 호출하기만 하면 됩니다. 이는 코드의 가독성을 높이고 유지보수성을 개선합니다.

-        combined_output = self.all2all_manager.moe_alltoall.combine(...)
-        output.copy_(combined_output)
+        self.all2all_manager.combine_into(
+            payload=fused_expert_output,
+            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
+            output=output,
+        )

왜 이게 좋은가

이번 최적화는 메모리 복사(Memory Copy)를 제거함으로써 GPU 메모리 대역폭 사용량을 줄였습니다. 실제 벤치마크 결과, Nemotron Ultra 550B 모델 환경에서 Output throughput이 약 2.31% 향상되었고, Mean TPOT(Time Per Output Token)이 약 3.10% 감소하는 성과를 거두었습니다.

일반적 교훈

1. API 설계의 중요성: 라이브러리(FlashInfer)가 외부 버퍼를 직접 입력받을 수 있도록 설계되면, 상위 프레임워크(vLLM)는 불필요한 메모리 할당과 복사를 피할 수 있습니다.
2. Capability-gated 최적화: 라이브러리 버전에 따라 기능 지원 여부가 다를 때, supports_kw와 같은 유틸리티를 사용하여 런타임에 기능을 감지하고 폴백을 제공하는 방식은 안정적인 배포를 가능하게 합니다.
3. 책임의 분리: 리뷰어의 피드백을 통해 기능 감지 로직을 개별 인스턴스가 아닌 Manager 클래스로 이동시킨 것은 객체 지향적 설계 관점에서 매우 적절한 개선이었습니다.

리뷰어 피드백 반영

리뷰어 shaharmor98은 기능 감지 로직을 finalize 메서드 내부가 아닌 Manager 클래스 초기화 시점으로 옮길 것을 제안했습니다. 이를 통해 매번 호출 시점에 검사할 필요 없이 효율적으로 관리할 수 있게 되었으며, supports_kw를 사용하여 더욱 견고한 API 검증이 가능해졌습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html
• https://github.com/flashinfer-ai/flashinfer

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