vLLM MoE 모델을 위한 Multi-LoRA 서빙

vLLM은 Multi-Low-Rank Adaptation (Multi-LoRA) 서빙을 위한 효율적인 솔루션을 도입하여, 조직이 단일 GPU에 수십 개의 파인튜닝된 모델을 호스팅할 수 있게 했습니다. 이는 각 커스텀 모델마다 전용 컴퓨팅 엔드포인트가 필요했던 상황을 없애며, 개별 모델이 충분한 트래픽을 받지 못해 하드웨어가 유휴 상태가 되는 문제를 해결합니다.

MoE 모델을 위한 Multi-LoRA 서빙

Multi-LoRA는 원본 베이스 모델 가중치를 고정한 채 작은 학습 가능한 어댑터를 요청마다 교체함으로써 여러 커스텀 모델이 동일한 GPU를 공유하도록 합니다. 이는 입력 토큰을 전문화된 신경망(전문가)으로 라우팅하는 라우터가 있는 Mixture of Experts (MoE) 모델에 특히 유용합니다.

MoE 아키텍처에서는 각 전문가가 "expand-then-compress" 패턴을 사용합니다: gate_up 프로젝션은 은닉 상태를 더 큰 중간 공간으로 확장하고, down 프로젝션은 이를 다시 압축합니다. LoRA가 적용될 때, 각 전문가마다 gate_up과 down 프로젝션 모두에 대해 축소와 확장 연산 두 개씩, 총 네 개의 추가 커널 연산이 필요합니다.

이를 처리하기 위해 vLLM은 fused_moe_lora 커널을 구현했습니다. 이 커널은 LoRA 연산을 fused_moe 커널에 직접 통합하여, 전문가 라우팅(다른 전문가에 할당된 토큰)과 어댑터 선택(다른 LoRA 어댑터를 사용하는 요청)으로 인해 발생하는 복합 희소성을 관리합니다.

성능을 위한 기술 최적화

MoE 모델에 대한 초기 multi-LoRA 구현은 높은 지연 시간을 겪었습니다. vLLM과 AWS는 NVIDIA Nsys와 NCU를 활용해 병목 현상을 식별하고 세 가지 주요 최적화 레이어를 통해 해결했습니다.

실행 최적화

프로파일링 결과 Triton 컴파일러가 입력 길이에 따라 변하는 변수를 컴파일 타임 상수로 처리하여, fused_moe_lora 커널이 새로운 컨텍스트 길이마다 재컴파일되는 문제가 발견되었습니다. 이는 기본 모델에 비해 Time To First Token (TTFT)이 10배 악화되는 결과를 초래했습니다. do_not_specialize 컴파일러 힌트를 추가해 커널이 한 번만 컴파일되고 모든 컨텍스트 길이에서 재사용되도록 함으로써 해결되었습니다.

커널 수준 최적화

LoRA 행렬은 "skinny"(즉, 랭크 r이 은닉 상태 차원보다 현저히 작음) 특성 때문에 표준 GEMM 커널의 성능이 저조합니다. 다음과 같은 최적화가 구현되었습니다:

  • Split-K Work Decomposition: 이 전략은 내부 차원 K에 대한 합산을 여러 스레드 그룹으로 나누어, 부분 합을 병렬로 계산함으로써 skinny 행렬에 대한 로드 밸런싱을 개선합니다. 원자적(add) 연산은 Triton 컴파일러에서 sem="relaxed"를 사용해 최적화되었습니다.
  • CTA Swizzling: Cooperative Thread Array(CTA) 스와이즐링은 GPU 스케줄을 재배열하여 인접한 열을 작업하는 스레드 그룹이 동시에 실행되도록 하여 L2 캐시 재사용을 증가시킵니다.
  • Masking Reduction: EVEN_K 파라미터를 도입해 행렬 차원이 블록 크기로 균등하게 나누어질 때 조건 마스킹 검사를 생략함으로써 각 로드 연산의 오버헤드를 감소시켰습니다.
  • Kernel Fusion: LoRA 가중치를 베이스 모델 가중치에 더하는 연산을 LoRA 확장 커널에 통합하여 커널 실행 오버헤드를 줄였습니다.

Amazon 전용 튜닝

표준 fused MoE에 최적화된 기본 Triton 커널 설정은 multi-LoRA의 복합 희소성을 고려하지 못했습니다. AWS는 네 가지 fused_moe_lora 연산(gate_up_shrink, gate_up_expand, down_shrink, down_expand)에 대해 맞춤형 튜닝 설정을 개발했습니다. 이러한 설정은 Amazon SageMaker AI와 Amazon Bedrock을 사용하는 고객에게 자동으로 로드됩니다.

성능 벤치마크 및 가용성

이러한 최적화는 GPT-OSS 20B, Qwen3-MoE, DeepSeek, Llama MoE와 같은 MoE 모델뿐만 아니라 Llama 3.3 70B, Qwen3 32B와 같은 밀집 모델의 성능을 크게 향상시켰습니다.

GPT-OSS 20B의 경우, vLLM 0.11.1rc3에서 vLLM 0.15.0으로 전환하면서 Output Tokens Per Second (OTPS)가 454% 향상되고 TTFT가 87% 감소했습니다. Amazon SageMaker AI와 Amazon Bedrock에서 제공되는 추가 최적화는 vLLM 0.15.0 대비 OTPS를 추가로 19% 증가시키고 TTFT를 8% 감소시킵니다.

GPT-OSS 20B 결과 요약:

Version/Platform OTPS TTFT
vLLM 0.11.1rc3 32 (Initial) ~1.2s (Estimated)
vLLM 0.15.0 144 135 ms
Amazon SageMaker/Bedrock 171 124 ms

이러한 개선 사항은 vLLM 버전 0.15.0 이상에서 사용할 수 있습니다.

Sources