[vllm] [vLLM] CPU MoE 성능을 극대화하는 전략: 고정 임계값을 넘어선 동적 GEMM 디스패치 최적화

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

들어가며

Mixture of Experts(MoE) 모델은 추론 시 모든 파라미터를 사용하지 않고 입력 토큰에 따라 필요한 '전문가(Expert)' 레이어만 활성화하여 효율성을 극대화합니다. 하지만 CPU 환경에서 MoE를 구현할 때는 각 전문가에게 할당되는 토큰의 양(M)이 매번 달라지기 때문에, 어떤 연산 커널을 사용할지 결정하는 디스패치(Dispatch) 전략이 성능의 핵심이 됩니다.

기존 vLLM의 CPU FP8/MXFP4 MoE 구현에서는 모든 전문가의 평균적인 워크로드를 계산하여 일괄적으로 커널을 선택하는 'Fixed Global Crossover' 방식을 사용했습니다. 이 방식은 특정 전문가에게 토큰이 몰리는 'Concentrated Routing' 상황에서 최적의 성능을 내지 못한다는 단점이 있었습니다. 본 포스트에서는 이를 해결하기 위해 개별 전문가 블록 단위로 커널을 동적으로 선택하도록 개선한 PR을 분석합니다.

코드 분석: 전역 평균에서 개별 최적화로

1. gemm.h: 커널 선택 조건의 세분화

먼저, BRGEMM(Batch Reduced GEMM)을 사용할 수 있는지 판단하는 헬퍼 함수가 확장되었습니다. 기존에는 행(M)의 크기만 고려했다면, 이제는 열(N)의 크기까지 고려하여 더 유연한 판단이 가능해졌습니다.

Before:

inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
  return brgemm_supported() && M > 4;
}

After:

template <typename T>
inline bool can_use_brgemm(int M, int N) {
  return brgemm_supported() && (can_use_brgemm<T>(M) || N != block_size_n());
}

리뷰어 bigPYJ1151의 제안에 따라 M과 N을 모두 인자로 받는 오버로딩이 추가되었습니다. 이는 특정 블록 사이즈(block_size_n)를 벗어나는 경우에도 BRGEMM의 이점을 취할 수 있도록 설계된 것입니다.

2. moe_fp8.cpp: 루프 내부의 동적 디스패치

가장 핵심적인 변화는 fused_experts_fp_kernel_impl 함수에 있습니다. 기존에는 커널 실행 전 avg_M을 계산하여 모든 전문가에게 동일한 use_brgemm 플래그를 적용했습니다.

Before:

// 전역 평균을 기반으로 한 번만 결정
int64_t avg_M = std::max(int64_t(1), M * topk / E);
const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);

// ... 루프 시작 ...
    /* brg */ use_brgemm,

After: 이 PR은 전역 변수인 avg_M을 제거하고, 실제 각 전문가 블록의 크기인 m_size를 기반으로 루프 내부에서 결정을 내립니다.

// 루프 내부에서 개별 전문가의 m_size를 확인
int64_t m_size = offsets[mb + 1] - offsets[mb];
if (m_size == 0) {
  return;
}
// 현재 블록의 m_size와 n_size에 최적화된 커널 선택
const bool brg = can_use_brgemm<packed_t>(m_size, n_size);

// ... 중략 ...

    /*   brg          */ brg,

이 변경을 통해 토큰이 많이 할당된 전문가는 고성능 BRGEMM 커널을 사용하고, 토큰이 적게 할당된(Small expert blocks) 전문가는 오버헤드가 적은 TinyGEMM 커널을 사용하도록 최적화되었습니다.

왜 이게 좋은가?

1. Concentrated Routing 대응 능력 향상

실제 LLM 추론 시 토큰들은 모든 전문가에게 균등하게 분배되지 않습니다. 특정 전문가가 '인기'가 많을 경우, 해당 전문가의 m_size는 매우 커집니다. 기존 방식은 전체 평균에 묶여 이 전문가조차 비효율적인 커널을 사용할 수 있었으나, 이제는 각 전문가가 자신의 워크로드에 맞는 최적의 무기를 선택할 수 있습니다.

2. 구체적인 성능 지표

벤치마크 결과는 이 최적화의 위력을 증명합니다.

• Qwen3-Coder-30B-A3B (FP8): 처리량(Throughput)이 약 1.10배 향상되었습니다.
• GPT-OSS-20B (TP=4): 특정 동시성(Concurrency 32) 환경에서 최대 1.59배라는 놀라운 성능 향상을 보였습니다 (611.7 -> 975.9 output tok/s).
• 전체 평균: 27개 테스트 포인트에서 기하평균 1.0866배의 성능 향상을 기록했습니다.

3. 유연한 아키텍처

리뷰어 bigPYJ1151은 SGLang과의 코드 정렬(Alignment)을 위해 bool 타입을 사용하고 일관된 오버로딩을 유지할 것을 권장했습니다. 이는 오픈소스 프로젝트에서 코드의 유지보수성과 타 프로젝트와의 호환성을 고려한 시니어 엔지니어의 통찰을 보여줍니다.

일반적인 교훈

이 PR은 "평균의 함정"을 경계해야 함을 시사합니다. 시스템 최적화 시 전체 평균 데이터에 기반한 결정은 구현이 간단하지만, 데이터 분포가 불균형할 때는 성능 저하의 원인이 됩니다. 워크로드의 특성을 파악하고 결정 지점(Decision Point)을 런타임의 가장 세밀한 단위(Granularity)로 옮기는 것만으로도 추가적인 하드웨어 자원 없이 상당한 성능 이득을 얻을 수 있습니다.

마치며

vLLM의 이번 업데이트는 CPU 환경에서 FP8/MXFP4 양자화 모델을 구동하는 사용자들에게 큰 혜택을 제공합니다. 특히 MoE 구조를 가진 최신 모델(Qwen, GPT-OSS 등)에서 그 효과가 두드러집니다. CPU 추론 성능을 한 단계 끌어올린 이 정교한 디스패치 최적화는 고성능 서빙 엔진이 지향해야 할 방향을 잘 보여주고 있습니다.

참고 자료

• https://oneapi-src.github.io/oneDNN/dev_guide_brgemm.html
• https://github.com/vllm-project/vllm/blob/main/csrc/cpu/sgl-kernels/gemm.h

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