본문으로 건너뛰기

[vllm] vLLM, CUDA 네이티브 SwiGLU 커널 도입으로 Humming MoE 성능 1.4% 향상

PR 링크: vllm-project/vllm#53685 상태: Merged | 변경: +17 / -0

들어가며

최근 대규모 언어 모델(LLM) 분야에서는 Mixture-of-Experts (MoE) 아키텍처가 큰 주목을 받고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 계산 효율성을 높이는 방식입니다. vLLM은 이러한 MoE 모델을 효율적으로 서빙하기 위한 다양한 최적화를 진행해왔습니다. 이번 PR은 vLLM에서 Hummingbird MoE (Humming MoE)를 사용할 때, 활성화 함수인 SwiGLU의 특정 연산을 네이티브 CUDA 커널로 대체하여 성능을 향상시킨 내용을 다룹니다.

이 PR은 특히 deepseek-ai/DeepSeek-V4-Flash-0731 모델과 같이 Humming MoE를 사용하는 경우, SwiGLU 활성화 함수의 clamp 연산이 비효율적인 파이썬/PyTorch 구현으로 처리되는 문제를 해결합니다. 이를 통해 처리량(throughput)을 약 1.4% 개선하는 성과를 거두었습니다.

코드 변경 분석

이번 PR의 핵심은 vllm/model_executor/layers/fused_moe/utils.py 파일에서 SwiGLU 연산의 일부를 최적화하는 것입니다. 특히, topk_ids가 제공되지 않는 경우(즉, 라우팅이 없는 경우)에 대한 처리가 개선되었습니다.

vllm/model_executor/layers/fused_moe/utils.py 변경 사항

기존 코드에서는 topk_ids가 None일 경우, SwiGLU 연산의 clamp 부분을 PyTorch의 기본 구현인 _swiglu_limit_torch 함수를 호출하여 처리했습니다. 이 함수는 Python 레벨에서 동작하거나 일반적인 PyTorch 연산을 사용하기 때문에 GPU 커널의 성능을 최대한 활용하지 못하는 경우가 많습니다.

Before:

def swiglu_limit_func(
    output: torch.Tensor,
    input: torch.Tensor,
    topk_ids: Optional[torch.Tensor] = None,
    swiglu_limit: float = 1.0,
    expert_map: Optional[torch.Tensor] = None,
) -> None:
    # ... (중략) ...
    else:
        _swiglu_limit_torch(output, input, swiglu_limit)

이번 PR에서는 topk_ids가 None이고 현재 플랫폼이 CUDA인 경우, PyTorch의 내부 CUDA 연산인 torch.ops._C.silu_and_mul_with_clamp를 직접 호출하도록 변경했습니다.

After:

def swiglu_limit_func(
    output: torch.Tensor,
    input: torch.Tensor,
    topk_ids: Optional[torch.Tensor] = None,
    swiglu_limit: float = 1.0,
    expert_map: Optional[torch.Tensor] = None,
) -> None:
    # ... (중략) ...
    elif current_platform.is_cuda():
        torch.ops._C.silu_and_mul_with_clamp(output, input, swiglu_limit, 1.0, 0.0)
    else:
        _swiglu_limit_torch(output, input, swiglu_limit)

torch.ops._C.silu_and_mul_with_clamp는 PyTorch에서 제공하는 최적화된 CUDA 커널로, silu (Swish) 함수와 곱셈, 그리고 clamp 연산을 한 번에 처리합니다. 이 연산은 GPU에서 매우 효율적으로 실행되도록 설계되었습니다.

tests/kernels/core/test_activation.py 변경 사항

또한, 새로운 swiglu_limit_func의 동작을 검증하기 위한 테스트 케이스가 추가되었습니다.

Before: (해당 테스트 함수 없음)

After:

@torch.inference_mode()
def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
    x = torch.randn(7, 1024, dtype=torch.bfloat16, device="cuda")
    output = torch.empty(7, 512, dtype=x.dtype, device=x.device)

    swiglu_limit_func(output, x, swiglu_limit=7.0)
    gate, up = x.chunk(2, dim=-1)
    expected = torch.nn.functional.silu(gate.clamp(max=7.0)) * up.clamp(
        min=-7.0, max=7.0
    )

    torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)

이 테스트는 swiglu_limit_func가 topk_ids 없이 호출될 때, 예상되는 결과와 동일한 출력을 생성하는지 검증합니다. 이는 새로운 CUDA 커널이 정확하게 동작함을 보장하는 중요한 단계입니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 테스트 결과는 이 변경의 효과를 명확하게 보여줍니다.

Variant Output token/s (three trials) Mean Total token/s mean TTFT mean TPOT mean
Baseline 667.07, 653.91, 663.19 661.39 5952.52 1361.48 ms 13.648 ms
Optimized 672.95, 668.88, 670.09 670.64 6035.75 1325.21 ms 13.598 ms

최적화된 버전은 평균 처리량(Output token/s)에서 약 1.40% (670.64 / 661.39)의 향상을 보였습니다. 이는 초당 처리할 수 있는 토큰 수가 늘어났음을 의미하며, 특히 대규모 MoE 모델에서 이러한 미세한 성능 향상이 모이면 전체적인 추론 속도에 상당한 영향을 미칠 수 있습니다.

또한, 평균 TTFT(Time To First Token)도 1361.48ms에서 1325.21ms로 감소했으며, TPOT(Time Per Output Token)도 약간 개선되었습니다. 이는 응답 지연 시간을 줄이고 사용자 경험을 향상시키는 데 기여합니다.

일반적 교훈

  1. 네이티브 커널 활용의 중요성: PyTorch와 같은 딥러닝 프레임워크는 내부적으로 고도로 최적화된 CUDA 커널을 제공합니다. 이러한 커널을 직접 활용하면 Python 오버헤드를 줄이고 GPU 하드웨어의 성능을 최대한 끌어낼 수 있습니다. 특히, torch.ops._C와 같이 내부적으로 노출되는 연산들은 성능이 검증된 경우가 많습니다.
  2. 조건부 최적화: 모든 경우에 동일한 최적화가 적용될 수는 없습니다. 이 PR처럼 topk_ids의 존재 여부와 같은 특정 조건에 따라 다른 구현(네이티브 CUDA 커널 vs. PyTorch 구현)을 사용하는 것은 효율적인 최적화 전략입니다.
  3. 테스트의 중요성: 새로운 최적화 로직이 추가될 때는 반드시 정확성을 검증하는 테스트 케이스를 함께 작성해야 합니다. 이를 통해 변경 사항이 의도한 대로 동작하는지, 그리고 예상치 못한 부작용은 없는지를 확인할 수 있습니다.
  4. MoE 아키텍처의 특성 이해: MoE 모델은 라우팅 메커니즘과 각 전문가 내에서의 연산이 성능에 큰 영향을 미칩니다. 활성화 함수와 같은 기본 연산의 최적화는 MoE 모델의 전체 성능에 직접적인 영향을 줄 수 있습니다.

리뷰 댓글 분석

리뷰어 chaunceyjiang의 댓글은 이 PR의 핵심 변경 사항을 정확히 지적하고 있습니다. apply_activation 함수에서 topk_ids가 전달되지 않아 기존에는 _swiglu_limit_torch로 폴백(fallback)되던 경로를, 이 PR이 ops._C.silu_and_mul_with_clamp로 대체했다는 점을 명확히 했습니다. 이는 이 변경이 기존 코드의 특정 분기(branch)를 개선하는 것임을 확인시켜 줍니다.

결론

이번 vLLM의 PR은 Humming MoE 아키텍처에서 SwiGLU 활성화 함수의 특정 연산을 네이티브 CUDA 커널로 대체함으로써, 실제 모델(DeepSeek-V4-Flash-0731)의 처리량을 1.4% 이상 향상시키는 실질적인 성능 개선을 이루었습니다. 이는 LLM 서빙 프레임워크에서 저수준 최적화가 얼마나 중요한지를 다시 한번 보여주는 사례입니다. 앞으로도 vLLM 팀은 이러한 지속적인 최적화를 통해 더 빠르고 효율적인 LLM 추론을 지원할 것으로 기대됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글