[vllm] vLLM, Qwen 모델의 성능을 극적으로 향상시킨 비결: SM121 TP=1 skinny-GEMM 최적화 분석

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

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 하지만 모델의 크기가 커질수록 추론(inference) 속도는 중요한 병목 현상이 됩니다. 특히 실시간 응답이 중요한 서비스에서는 이러한 성능 개선이 필수적입니다. vLLM은 LLM 추론을 위한 고성능 서빙 라이브러리로, 지속적인 최적화를 통해 LLM의 성능 한계를 넓혀가고 있습니다. 이번 글에서는 vLLM의 Qwen 모델 관련 PR([Perf][Qwen4Exp] Add SM121 TP=1 skinny-GEMM plans)을 통해 이루어진 성능 최적화에 대해 심층적으로 분석하고자 합니다. 이 PR은 특히 Qwen 모델의 특정 설정(TP=1, SM121 아키텍처)에서 발생하는 지연 시간을 크게 줄여, 전체적인 추론 성능을 9.5%까지 향상시키는 성과를 거두었습니다. 이 글에서는 어떤 코드가 어떻게 변경되었고, 왜 이러한 변경이 성능 향상으로 이어졌는지, 그리고 이 최적화가 주는 일반적인 교훈은 무엇인지 상세히 살펴보겠습니다.

코드 분석: SM121 TP=1 skinny-GEMM 계획 추가

이번 PR의 핵심은 NVIDIA의 최신 GPU 아키텍처인 SM121 (Ampere 아키텍처의 일부)에서 Qwen 모델의 특정 연산, 특히 TP=1 (Tensor Parallelism degree 1) 설정에서의 skinny-GEMM (행렬 곱셈) 연산에 대한 최적화된 커널 계획을 추가한 것입니다. 기존에는 이러한 특정 형태의 행렬 곱셈이 cuBLAS와 같은 범용 라이브러리에 의존하여 최적의 성능을 내지 못했습니다. 이 PR은 해당 연산에 특화된 skinny-GEMM 커널을 도입하여 성능을 개선했습니다.

1. vllm/models/qwen4_exp/nvidia/low_latency_gemm.py 파일 변경

이 파일은 Qwen 모델의 특정 레이어에 대한 저지연 행렬 곱셈 커널 계획을 정의합니다. 이번 PR에서는 QWEN4_EXP_SM121_GEMM_PLANS라는 새로운 딕셔너리를 추가했습니다.

Before:

이전에는 QWEN4_EXP_SM100_GEMM_PLANS와 QWEN4_EXP_SM90_GEMM_PLANS만 존재했습니다. SM121 아키텍처에 대한 최적화된 계획이 없었기 때문에, 해당 GPU에서는 cuBLAS와 같은 일반적인 구현에 의존해야 했습니다.

After:

# DGX Spark (GB10) plans, TP=1.
QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
    # Shared-expert gate.
    (1, 2560): {
        1: SkinnyGemmConfig(1, 64, 1, k_unroll=2, static_k=2560),
        2: SkinnyGemmConfig(2, 64, 1, k_unroll=2, vector_width=4, static_k=2560),
        4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
        8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
        16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
    },
    # GDN fused B/A projection.
    (96, 2560): {
        1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=2, static_k=2560),
        # ... (다른 레이어들에 대한 설정들) ...
    },
    # LM head.
    (248320, 2560): {
        1: SkinnyGemmConfig(1, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
        2: SkinnyGemmConfig(2, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
        4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560),
        8: SkinnyGemmConfig(8, 64, 1, k_unroll=2, vector_width=4, static_k=2560),
        16: SkinnyGemmConfig(16, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
    },
}

def _is_sm121() -> bool:
    return current_platform.is_device_capability((12, 1))

def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
    # ... (기존 로직) ...
    if _is_sm121():
        return QWEN4_EXP_SM121_GEMM_PLANS
    return {}

새롭게 추가된 QWEN4_EXP_SM121_GEMM_PLANS 딕셔너리는 (N, K) 형태의 행렬 차원과 num_tokens (배치 크기 또는 시퀀스 길이와 관련된 값)에 따라 최적화된 SkinnyGemmConfig를 매핑합니다. SkinnyGemmConfig는 m, n, k 블록 크기, k_unroll, vector_width, static_k 등 커널의 세부적인 실행 설정을 담고 있습니다. 특히 LM head 레이어 (248320, 2560)에 대한 설정이 포함된 것이 주목할 만합니다. 또한, _gemm_plans 함수에 _is_sm121() 체크 로직이 추가되어, SM121 GPU에서 실행될 때 이 새로운 계획이 선택되도록 합니다.

2. tests/kernels/test_bf16_skinny_gemm.py 파일 변경

테스트 파일에서는 새로운 SM121 계획이 올바르게 작동하는지 검증하는 코드가 추가되었습니다.

Before:

테스트는 주로 SM100, SM90 등에 대한 계획을 검증했습니다.

After:

QWEN4_EXP_SM121_CASES = [
    (n, k, num_tokens, config)
    for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS.items()
    for num_tokens, config in plans.items()
]

# ... (기존 테스트 함수들) ...

@pytest.mark.parametrize(

## 참고 자료
- https://github.com/vllm-project/vllm/blob/v0.4.2/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py#L207-L361
- https://github.com/vllm-project/vllm/blob/v0.4.2/tests/kernels/test_bf16_skinny_gemm.py#L91-L100
- https://github.com/vllm-project/vllm/blob/v0.4.2/vllm/model_executor/layers/linear.py#L100-L110
- https://github.com/vllm-project/vllm/blob/v0.4.2/vllm/model_executor/utils.py#L10-L20

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