vLLM에서 MoE 모델을 위한 효율적인 Multi-LoRA 추론 구현 및 최적화

핵심 요약

여러 커스텀 AI 모델을 운영할 때 발생하는 GPU 자원 낭비를 해결하기 위해 vLLM과 AWS가 협력하여 MoE 모델용 Multi-LoRA 추론 솔루션을 개발했다. 기존 vLLM은 MoE 레이어에 대한 LoRA 커널이 부족하고 복합적인 희소성 처리 문제가 있었으나, 새로운 fused_moe_lora 커널을 통해 이를 해결했다. Triton 컴파일러 힌트, Split-K 전략, CTA 스위즐링 등 커널 레벨의 최적화를 적용한 결과, GPT-OSS 20B 모델 기준 출력 토큰 속도(OTPS)가 초기 대비 454% 향상되고 첫 토큰 생성 시간(TTFT)은 87% 단축되었다. 이 기술은 Amazon SageMaker AI 및 Bedrock에 적용되어 추가적인 성능 이점을 제공한다.

배경

GPU 아키텍처 기초, CUDA/Triton 프로그래밍 이해, MoE 및 LoRA 개념, vLLM 사용 경험

대상 독자

LLM 추론 엔진 개발자 및 MLOps 엔지니어

의미 / 영향

이번 최적화는 MoE 기반 오픈소스 모델의 상용 서비스 비용을 획기적으로 낮추고, 다중 어댑터 환경에서 vLLM의 경쟁력을 강화한다.

섹션별 상세

MoE 모델은 전문가(Expert) 라우팅을 통해 일부 파라미터만 활성화하며, LoRA는 원본 가중치를 고정한 채 작은 어댑터 행렬을 추가한다. MoE의 각 전문가는 gate_up 및 down 프로젝션 단계를 거치는데, LoRA 적용 시 각 단계마다 shrink와 expand 연산이 추가되어 총 4개의 LoRA 커널 연산이 필요하게 된다. 이러한 구조적 복잡성은 추론 시 성능 병목의 주요 원인이 된다.

기존 vLLM의 dense 모델용 LoRA 커널은 MoE의 전문가 라우팅을 지원하지 못하는 한계가 있었다. 이를 해결하기 위해 LoRA 연산을 fused_moe 커널에 통합한 새로운 fused_moe_lora 커널을 설계했다. 또한 Triton 컴파일러가 입력 길이에 따라 매번 커널을 재컴파일하는 오버헤드를 방지하기 위해 do_not_specialize 힌트를 추가하여 TTFT 지연 문제를 해결했다.

얇은 행렬(Skinny matrices) 연산 효율을 높이기 위해 Split-K 전략을 도입하여 연산을 병렬화하고 원자적 덧셈(Atomic add)을 최적화했다. CTA 스위즐링(Swizzling)을 통해 L2 캐시 재사용률을 높였으며, 불필요한 마스킹 및 도트 프로덕트 연산을 제거하여 연산 오버헤드를 최소화했다. 이러한 커널 레벨의 개선은 데이터 로드 효율을 극대화한다.

블록 크기(BLOCK_SIZE)와 같은 커널 설정 파라미터를 MoE LoRA의 복합 희소성 특성에 맞게 최적화했다. Amazon SageMaker AI 및 Bedrock 환경에서는 이러한 튜닝된 설정이 자동으로 로드되어 vLLM 0.15.0 순정 버전 대비 OTPS는 19% 향상되고 TTFT는 8% 추가로 개선되는 성과를 거두었다. 사용자는 커스텀 튜닝 설정을 직접 로드하여 성능을 최적화할 수도 있다.