[vllm] vLLM ROCM 최적화: GLM-4 MoE를 위한 Fused Shared Expert(FSE) 도입

PR 링크: vllm-project/vllm#44313 상태: Merged | 변경: +254 / -105

들어가며

최신 MoE(Mixture-of-Experts) 모델들은 성능 향상을 위해 '공유 전문가(Shared Expert)' 구조를 자주 사용합니다. 하지만 기존 vLLM 구현에서는 이 공유 전문가가 별도의 MLP 패스로 실행되어, 특히 낮은 동시성(Concurrency) 환경에서 커널 실행 오버헤드가 성능 병목으로 작용했습니다. 본 PR은 ROCm 환경에서 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 옵션을 통해 공유 전문가를 FusedMoE 커널 내부로 통합(FSE)하여, 별도의 연산 단계를 제거함으로써 추론 효율을 극대화합니다.

코드 분석

1. vllm/model_executor/models/glm4_moe.py

핵심 변경 사항은 Glm4MoE 초기화 시 공유 전문가를 별도의 MLP 모듈로 생성하지 않고, FusedMoE의 추가 슬롯으로 전달하는 것입니다.

Before:

if config.n_shared_experts is not None:
    self.shared_experts = Glm4MoeMLP(...)
else:
    self.shared_experts = None

After:

if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
    self.shared_experts = None
else:
    self.shared_experts = Glm4MoeMLP(...)

또한, FusedMoE 호출 시 n_shared_experts를 전달하여 커널이 이를 인지하게 하고, apply_routed_scale_to_output 설정을 조정하여 스케일링 오류를 방지했습니다.

2. vllm/model_executor/models/glm4_moe_mtp.py

MTP(Multi-Token Prediction) 모듈에서도 동일한 로직을 적용했습니다. load_weights 과정에서 공유 전문가 가중치를 슬라이싱하여 FusedMoE의 전문가 슬롯으로 매핑하는 로직이 추가되었습니다.

# FSE: split a widened mlp.shared_experts tensor into n_shared_experts chunks
chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
weight_to_load = loaded_weight[chunk_slice, :] # 예시: ColumnParallel

왜 이게 좋은가

이번 최적화는 특히 낮은 동시성(MC=4) 환경에서 TPOT(Time Per Output Token)을 약 18~19% 개선하고, 출력 처리량을 20% 이상 향상시켰습니다.

• 커널 오버헤드 감소: 별도의 MLP 연산을 제거함으로써 GPU 커널 실행 횟수를 줄였습니다.
• 정확도 유지: apply_routed_scale_to_output 설정을 not self.is_rocm_aiter_moe_enabled로 변경하여, 공유 전문가 슬롯에 부적절한 스케일링이 적용되는 문제를 해결했습니다. 이 수정 덕분에 GSM8K 벤치마크에서 정확도 회귀 없이 Baseline과 동일한 성능을 유지합니다.

교훈: MoE 모델에서 공유 전문가와 라우팅 전문가를 하나의 커널로 융합하는 것은 연산 밀도를 높이고 메모리 접근 패턴을 최적화하는 매우 효과적인 전략입니다. 특히 ROCm/AITER와 같은 가속기 환경에서는 커널 런치 오버헤드를 줄이는 것이 성능의 핵심입니다.

리뷰어 피드백 반영

리뷰어 dllehr-amd의 제안에 따라, 코드 내부에 장황하게 작성되었던 주석들을 간결하게 정리하고, 특정 하드웨어 아키텍처에 대한 직접적인 언급을 제거하여 유지보수성을 높였습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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