[vllm] vLLM ROCm 환경에서 Shared-Expert Fusion을 통한 MoE 추론 성능 최적화

PR 링크: vllm-project/vllm#46545 상태: Merged | 변경: +110 / -20

들어가며

최신 Mixture-of-Experts (MoE) 모델들은 효율성을 위해 여러 전문가(Expert)를 사용하지만, 일부 모델은 모든 토큰이 거쳐야 하는 '공유 전문가(Shared-expert)'를 별도의 dense MLP 레이어로 처리합니다. MiniMax-M3 모델의 경우, 60개 레이어마다 이 공유 전문가가 별도의 GEMM 연산으로 실행되는데, 이는 특히 낮은 동시성(concurrency) 환경에서 launch-bound 병목을 유발합니다. 본 PR은 이 공유 전문가를 라우팅된 그룹 GEMM(Grouped GEMM)에 통합(Fusion)하여 불필요한 커널 실행을 제거하고 추론 성능을 극대화합니다.

코드 분석

1. Router 로직 수정 (router/fused_topk_bias_router.py)

공유 전문가를 라우팅된 전문가 슬롯 뒤에 추가하여, 모든 토큰이 공유 전문가를 거치도록 강제합니다.

# Before: 라우팅된 top-k 전문가만 처리
# After: 공유 전문가 슬롯을 강제로 추가
if self.num_fused_shared_experts > 0:
    base = self.global_num_experts
    shared_ids = torch.arange(base, base + n, ...).expand(m, n)
    shared_w = torch.full((m, n), self.shared_expert_weight, ...)
    topk_ids = torch.cat([topk_ids, shared_ids], dim=-1)
    topk_weights = torch.cat([topk_weights, shared_w], dim=-1)

2. Expert 정렬 로직 수정 (experts/mxfp8_native_moe.py)

기존에는 global_num_experts를 기준으로 전문가를 정렬했으나, 공유 전문가가 추가되면서 인덱스 범위가 확장되었습니다. 이를 실제 가중치 행의 개수로 동적 계산하도록 수정했습니다.

- global_num_experts,
+ num_align_experts = w13.shape[0] if expert_map is None else global_num_experts
+ num_align_experts,

3. 모델 설정 통합 (models/minimax_m3/amd/model.py)

MiniMax-M3 모델에서 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS 환경 변수를 통해 이 기능을 선택적으로 활성화할 수 있도록 구현했습니다.

왜 이게 좋은가

이 최적화는 'Launch-bound' 문제를 해결하는 데 매우 효과적입니다. 특히 디코딩 단계에서 커널 실행 횟수를 줄임으로써 성능 향상을 이끌어냈습니다.

성능 수치 (Throughput)

• Concurrency 1: 59.2 -> 77.0 (+30.2%)
• Concurrency 32: 1115.5 -> 1246.8 (+11.8%)
• Concurrency 128: 2940.0 -> 3103.8 (+5.6%)

교훈

1. Kernel Fusion의 힘: 여러 번의 작은 GEMM 호출을 하나의 큰 그룹 GEMM으로 통합하는 것은 GPU 활용도를 높이는 가장 강력한 방법 중 하나입니다.
2. 수치적 등가성(Numerical Equivalence): 최적화 과정에서 모델의 정확도(GSM8K 결과)를 유지하는 것이 필수적입니다. 본 PR은 공유 전문가의 가중치를 1/routed_scaling_factor로 조정하여 기존의 별도 MLP 연산과 수학적으로 동일한 결과를 보장했습니다.
3. 옵트인(Opt-in) 설계: 특정 모델이나 하드웨어(ROCm)에 종속적인 최적화는 환경 변수를 통해 제어함으로써 코드 복잡도를 관리하고 안정성을 확보하는 것이 좋습니다.

리뷰어 피드백 반영

리뷰어들은 CI 리소스의 한계로 인해 모든 모델에 대한 자동 테스트가 어렵다는 점을 지적했습니다. 이에 따라, 특정 모델(MiniMax-M3)에 한정하여 기능을 활성화하고, 향후 다른 모델로 확장할 수 있는 유연한 구조를 유지하는 방향으로 합의되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cat.html
• https://docs.vllm.ai/en/latest/models/adding_model.html

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