Amazon SageMaker 및 Amazon Bedrock에서 vLLM을 활용한 MoE 모델 Multi-LoRA 서빙 최적화

핵심 요약

다수의 미세 조정된 모델을 운영할 때 발생하는 GPU 자원 낭비를 해결하기 위해 vLLM에 Mixture of Experts(MoE) 모델용 Multi-LoRA 서빙 기능이 도입되었다. AWS는 vLLM 커뮤니티와 협력하여 fused_moe_lora 커널을 개발하고, Triton 컴파일러 최적화, Split-K 전략, CTA 스위즐링 등 하드웨어 수준의 튜닝을 진행했다. 그 결과 GPT-OSS 20B 모델에서 초기 구현 대비 출력 토큰 속도(OTPS)는 454% 향상되었고, 첫 토큰 생성 시간(TTFT)은 87% 단축되었다. 이 최적화는 현재 vLLM 0.15.0 이상 버전과 Amazon SageMaker 및 Bedrock에서 즉시 활용 가능하다.

배경

vLLM 추론 엔진 기초, LoRA(Low-Rank Adaptation) 및 MoE 아키텍처 이해, CUDA 및 Triton 커널 연산 개념

대상 독자

프로덕션 환경에서 대규모 MoE 모델을 다수의 사용자에게 효율적으로 서빙하려는 ML 엔지니어 및 인프라 아키텍트

의미 / 영향

이번 최적화는 오픈소스 커뮤니티와 클라우드 제공업체의 협업이 실질적인 추론 효율성 향상으로 이어진 사례이다. MoE 모델의 경제성을 극대화하여 기업들이 더 적은 비용으로 고성능 맞춤형 AI 서비스를 운영할 수 있는 기술적 토대를 마련했다.

섹션별 상세

MoE 모델은 전문가(Expert) 라우팅을 통해 효율성을 높이지만, 여기에 다수의 LoRA 어댑터를 동시에 적용하면 라우팅과 어댑터 선택이라는 이중의 희소성(Sparsity) 문제가 발생한다. 기존 vLLM은 밀집(Dense) 모델용 LoRA 커널만 보유하고 있어 MoE 레이어의 전문가 구조를 처리하지 못했다. 이를 해결하기 위해 전문가별 투영(Projection) 연산에 LoRA의 shrink 및 expand 연산을 통합한 fused_moe_lora 커널이 새롭게 설계되었다.

초기 구현에서는 컨텍스트 길이에 따라 커널이 매번 재컴파일되어 TTFT가 기본 모델 대비 10배나 느려지는 문제가 발견되었다. 이를 해결하기 위해 Triton 컴파일러에 do_not_specialize 힌트를 추가하여 커널 재사용성을 높였으며, LoRA가 없는 레이어는 즉시 종료하는 로직을 도입했다. 또한 Programmatic Dependent Launch(PDL)를 통해 shrink와 expand 커널 실행을 중첩시키고, CudaGraph를 적용하여 CPU 오버헤드와 커널 런칭 지연을 최소화했다.

LoRA 연산의 특성상 발생하는 가느다란 행렬(Skinny Matrix) 연산 효율을 높이기 위해 Split-K 전략을 도입하여 연산을 병렬화했다. Cooperative Thread Array(CTA) 스위즐링을 적용해 L2 캐시 재사용률을 높였으며, 행렬 크기가 블록 크기에 딱 맞는 경우 마스킹 연산을 생략하는 EVEN_K 파라미터를 도입해 불필요한 연산을 제거했다. 마지막으로 LoRA 가중치 합산 과정을 expand 커널에 통합하여 커널 호출 횟수를 줄였다.

AWS는 SageMaker와 Bedrock 환경에 최적화된 블록 크기(BLOCK_SIZE)와 병렬화 파라미터(SPLIT_K)를 찾아내어 자동 로드되도록 구현했다. GPT-OSS 20B 모델 테스트 결과, vLLM 0.15.0 기본 버전보다도 OTPS는 19% 더 높고 TTFT는 8% 더 낮은 성능을 기록했다. 최종적으로 초기 구현체와 비교했을 때 OTPS는 약 5.5배 향상되는 성과를 거두었다.

이미지 분석

Diagram
MoE 레이어의 라우터가 토큰을 전문가에게 전달하는 과정과, 각 전문가의 투영 연산에 LoRA의 shrink 및 expand 단계가 어떻게 결합되는지 시각적으로 설명한다. 특히 gate_up과 down projection 각각에 어댑터 경로가 추가되는 구조를 명확히 보여준다.
MoE-LoRA 모델의 작동 원리를 보여주는 아키텍처 다이어그램이다.

Screenshot
커널 실행 전 cuModuleLoadData 호출로 인해 약 0.5초의 지연이 발생하는 것을 보여준다. 이는 Triton 컴파일러가 매번 새로운 커널을 컴파일하면서 발생하는 오버헤드로, 최적화의 필요성을 증명하는 핵심 데이터이다.
최적화 전 fused_moe_lora 커널의 프로파일링 결과 타임라인이다.

Chart
초기 구현, vLLM 0.15.0, 그리고 AWS 커스텀 튜닝이 적용된 버전의 성능을 비교한다. 최적화 후 OTPS가 비약적으로 상승하고 TTFT가 급격히 감소하는 수치를 통해 본문의 개선 효과를 정량적으로 뒷받침한다.
GPT-OSS 20B 모델의 OTPS 및 TTFT 성능 비교 차트이다.

실무 Takeaway

MoE 모델을 Multi-LoRA로 서빙하면 GPU 활용도를 극대화하여 유휴 자원 비용을 대폭 절감할 수 있다.
Triton 커널 작성 시 do_not_specialize와 같은 컴파일러 힌트를 활용해 불필요한 재컴파일 오버헤드를 방지하는 것이 성능 유지에 필수적이다.
가느다란 행렬 연산이 주를 이루는 LoRA 특성상 Split-K와 CTA 스위즐링 같은 하드웨어 친화적 최적화가 전체 추론 속도를 결정한다.

언급된 리소스

GitHubvLLM GitHub Repository

문서vLLM LoRA Tuning Documentation