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

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

대규모 Mixture of Experts(MoE) 모델을 여러 사용자에게 맞춤형으로 제공할 때 발생하는 GPU 자원 낭비 문제를 해결하기 위해 vLLM에 Multi-LoRA 서빙 기능을 도입했다. 베이스 모델 가중치를 공유하면서 요청별로 어댑터만 교체하는 fused_moe_lora 커널을 새롭게 설계했다. Triton 컴파일러 최적화, Split-K 기법, CTA swizzling 등 커널 수준의 고도화된 튜닝을 통해 GPT-OSS 20B 모델 기준 출력 속도를 454% 개선하고 첫 토큰 생성 시간을 87% 단축했다. 이 최적화 기술은 현재 Amazon SageMaker AI와 Bedrock에서 즉시 활용 가능하다.

배경

vLLM 추론 엔진에 대한 기본 지식, MoE(Mixture of Experts) 및 LoRA 아키텍처 이해, GPU 커널 연산 및 Triton 컴파일러 개념

대상 독자

프로덕션 환경에서 대규모 MoE 모델을 다중 사용자에게 서빙해야 하는 ML 엔지니어 및 인프라 아키텍트

의미 / 영향

이 기술은 MoE 모델의 높은 추론 비용 문제를 해결하여 기업들이 수십 개의 맞춤형 모델을 경제적으로 운영할 수 있게 한다. 특히 오픈소스 모델인 GPT-OSS, Qwen, DeepSeek 등을 활용한 서비스의 확장성을 크게 높일 것으로 기대된다.

섹션별 상세

MoE 모델은 각 전문가(Expert) 네트워크가 gate_up 확장과 down 압축의 2단계 투영 과정을 거치며, LoRA 적용 시 각 단계에 shrink와 expand 연산이 추가되어 총 4개의 커널 연산이 발생한다.

MoE 레이어 내에서 LoRA 어댑터가 통합되는 구조를 보여주는 아키텍처 다이어그램이다. — Diagram라우터가 토큰을 전문가에게 전달하는 과정과 각 전문가 내부의 gate_up/down 투영 단계에 LoRA의 shrink/expand 연산이 어떻게 추가되는지 시각화한다. 이를 통해 MoE와 LoRA가 결합될 때 발생하는 연산 복잡성을 이해할 수 있다.

초기 구현에서는 Triton 컴파일러가 입력 길이에 의존하는 변수를 상수로 취급하여 매번 커널을 재컴파일하는 오버헤드가 발생했으나, do_not_specialize 힌트를 적용하여 커널 재사용을 가능케 함으로써 TTFT를 대폭 개선했다.

LoRA 연산의 특성인 'Skinny Matrix' 문제를 해결하기 위해 Split-K 전략을 도입하여 내부 차원의 합산 작업을 여러 스레드 그룹에 분산 배치하고 병렬 처리 성능을 극대화했다.

최적화 전 fused_moe_lora 커널의 Nsys 프로파일링 결과로 재컴파일 지연 시간을 나타낸다. — ScreenshotcuModuleLoadData 호출로 인해 발생하는 약 0.5초의 실행 지연을 명확히 보여준다. 이는 Triton 컴파일러의 특수화 기능으로 인해 매번 새로운 커널이 로드되면서 발생하는 성능 병목을 증명한다.

CTA swizzling 기법을 적용하여 인접한 열을 처리하는 스레드 그룹들이 동시에 실행되도록 조정함으로써 L2 캐시 재사용률을 높이고 메모리 대역폭 효율을 최적화했다.

Amazon SageMaker AI 및 Bedrock 환경에 최적화된 블록 사이즈와 파라미터 설정을 자동으로 로드하는 기능을 추가하여 기본 설정 대비 OTPS를 19% 추가 향상시켰다.

GPT-OSS 20B 모델의 Multi-LoRA 추론 성능(OTPS 및 TTFT) 비교 차트이다. — Chart초기 구현 대비 vLLM 0.15.0 및 AWS 커스텀 튜닝 적용 시 OTPS는 약 5.5배 증가하고 TTFT는 8배 이상 감소했음을 수치로 나타낸다. 최적화 기법들이 실제 추론 효율성에 미치는 영향력을 입증한다.

Programmatic Dependent Launch(PDL)를 구현하여 shrink 커널이 실행되는 동안 expand 커널이 가중치를 미리 로드할 수 있도록 연산을 오버랩하여 실행 지연을 최소화했다.

실무 Takeaway

Multi-LoRA를 활용하면 개별 모델마다 전용 GPU를 할당할 필요 없이 하나의 GPU에서 수십 개의 미세 조정된 MoE 모델을 효율적으로 서빙하여 인프라 비용을 절감할 수 있다.
Triton 커널 작성 시 do_not_specialize 힌트를 사용하여 컨텍스트 길이에 따른 불필요한 재컴파일 오버헤드를 제거하고 첫 토큰 생성 시간(TTFT)을 10배 이상 단축할 수 있다.
LoRA와 같이 행렬의 한쪽 차원이 매우 작은 연산에는 Split-K 기법을 적용하여 GPU 스레드 그룹 간 부하 분산과 병렬성을 확보하는 것이 성능 최적화의 핵심이다.

언급된 리소스

GitHubvLLM GitHub Repository

문서vLLM LoRA Tuning Documentation