[vllm] group_topk 커널 최적화 - 1.9% Throughput, 2.1% TPOT 개선

2025년 12월 9일수정: 2025년 12월 9일

PR 링크: vllm-project/vllm#30159 상태: Merged | 변경: +128/-47

들어가며

MoE 모델에서 각 토큰이 어떤 expert로 라우팅될지 결정하는 group_topk 커널은 매 forward pass마다 호출된다. 이 PR은 런타임 분기를 컴파일 타임 템플릿으로 전환하고, 루프 언롤링을 적용하여 1.9%의 throughput 향상과 2.1%의 TPOT(Time Per Output Token) 개선을 달성한다.

핵심 코드 분석

런타임 분기를 컴파일 타임 템플릿으로 전환

// Before: 매번 런타임 분기
if (scoring_func == SCORING_SIGMOID) {
    value = apply_sigmoid(value);
}

// After: 컴파일 타임에 결정
template <ScoringFunc SF, typename T>
__device__ inline T apply_scoring(T val) {
    if constexpr (SF == SCORING_SIGMOID) {
        return apply_sigmoid(val);
    } else {
        return val;
    }
}

if constexpr는 컴파일 타임에 분기가 제거되어 불필요한 명령어가 생성되지 않는다. GPU에서 분기는 warp divergence를 유발할 수 있어, 이 최적화의 효과가 크다.

정적 N-group 루프 언롤링

// Before: 동적 루프
for (int i_group = 0; i_group < n_group; i_group++) {
    process_group(i_group);
}

// After: 정적 크기면 언롤링
template <typename T, typename IdxT, ScoringFunc SF, int NGroup = -1>
if constexpr (kUseStaticNGroup) {
    #pragma unroll
    for (int i_group = 0; i_group < NGroup; ++i_group) {
        process_group(i_group);
    }
} else {
    for (int i_group = 0; i_group < n_group_i32; ++i_group) {
        process_group(i_group);
    }
}

NGroup이 컴파일 타임에 알려진 경우 #pragma unroll로 루프가 완전히 펼쳐진다.

불필요한 연산 스킵

// Before: 항상 renormalize 계산
topk_sum += cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());

// After: renormalize가 필요할 때만 계산
if (renormalize) {
    topk_sum += cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
}