본문으로 건너뛰기

[vllm] vLLM, Qwen 3.8B 모델의 성능을 3배까지 향상시킨 최적화 분석

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

들어가며

최근 vLLM 프로젝트에서는 Qwen 3.8B 모델의 추론 성능을 획기적으로 개선하기 위한 Pull Request (PR)가 제출되었습니다. 이 PR은 특히 Tensor Parallelism (TP)이 1일 때, NVIDIA B300 (SM103) GPU 아키텍처에 최적화된 Skinny GEMM (General Matrix Multiply) 커널 계획을 추가하는 데 중점을 두고 있습니다. 결과적으로 cuBLAS 대비 최대 3.06배의 속도 향상을 달성했으며, 이는 초당 디코드 단계에서 약 0.52ms의 지연 시간 감소로 이어집니다. 본 글에서는 이 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 최적화가 성능 향상에 크게 기여하는지 기술적인 관점에서 설명하고자 합니다.

코드 분석

이번 PR의 핵심 변경 사항은 vllm/models/qwen4_exp/nvidia/low_latency_gemm.py 파일에 새로운 Skinny GEMM 설정을 추가하는 것입니다. Skinny GEMM은 행렬 곱셈에서 한쪽 차원이 매우 작을 때 발생하는 비효율성을 해결하기 위해 설계된 최적화 기법입니다. vLLM은 다양한 모델과 하드웨어 아키텍처에 맞춰 이러한 Skinny GEMM 설정을 세밀하게 조정하여 사용합니다.

vllm/models/qwen4_exp/nvidia/low_latency_gemm.py 변경 분석

이 PR은 기존 코드에 Qwen 3.8B 모델의 특정 레이어(GDN B/A, HC final down, QSA QKV, shared gate/up, shared down 등)에 대한 SM103 아키텍처(NVIDIA B300) 및 TP=1 환경에 최적화된 SkinnyGemmConfig들을 추가했습니다. 이전에는 SM103 아키텍처의 경우 TP=4 형태의 Shape만 지원되었으나, 이번 PR을 통해 TP=1에서도 최적화된 커널을 사용할 수 있게 되었습니다.

주요 추가 설정:

  • GDN and QSA output projections, TP=1 (DGX Station GB300):
    • 기존 코드에는 해당 설정이 없었습니다.
    • 새로 추가된 설정은 다음과 같습니다: 
      +    # GDN and QSA output projections, TP=1 (DGX Station GB300).
  • (2560, 6144): {
  •    1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=6144),

  •    2: SkinnyGemmConfig(2, 128, 2, static_k=6144),
  • }, 
*   `(2560, 6144)`은 (N, K) 행렬 차원을 나타내며, `SkinnyGemmConfig`는 `(num_stages, num_tiled_k, num_groups, k_unroll, vector_width, static_k)` 등의 파라미터를 포함합니다. 이 파라미터들은 CUDA 커널의 내부 동작을 최적화하기 위한 힌트 역할을 합니다. 예를 들어, `k_unroll`은 K 차원을 얼마나 많이 병렬 처리할지, `vector_width`는 벡터 연산의 크기를 지정합니다.
  • GDN fused B/A projection, TP=1:

  • GDN fused B/A projection, TP=1.

  • (96, 2560): {
  •    1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4, static_k=2560),

  •    2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),

  •    4: SkinnyGemmConfig(4, 160, 1, k_unroll=2, static_k=2560),
  • }, 
*   이 설정은 특히 `k_unroll=4`와 `vector_width=4`를 사용하여 K 차원의 병렬 처리 효율을 높였습니다.
  • QSA fused QKV/gate + replicated indexer Q/K, TP=1: 
*   `k_unroll=6`은 이 특정 행렬 곱셈에서 K 차원을 6개의 청크로 나누어 처리함으로써 성능을 향상시킵니다.
  • Shared-expert fused gate/up projection, TP=1:

  • Shared-expert fused gate/up projection, TP=1.

  • (1280, 2560): {
  •    1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=2560),

  •    2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=2560),

  •    4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
  • },
  • Shared-expert down projection, TP=1:
  • HC final down projection, TP=1:

  • HC final down projection, TP=1.

  • (320, 10240): {
  •    1: SkinnyGemmConfig(1, 256, 1, k_unroll=2, static_k=10240),

  •    2: SkinnyGemmConfig(2, 256, 1, k_unroll=2, static_k=10240),

  •    4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
  • }, 
*   이 설정은 `static_k=10240`으로 K 차원이 고정되어 있음을 명시하며, `k_unroll=2`를 통해 최적화합니다.

이러한 설정들은 특정 행렬 차원(N, K)과 Tensor Parallelism(TP) 값에 맞춰 가장 효율적인 CUDA 커널 실행 방식을 정의합니다. SkinnyGemmConfig의 각 파라미터는 GPU 하드웨어의 특성(예: SIMT 코어 수, 레지스터 파일 크기, 캐시 구조)을 고려하여 튜닝됩니다. 특히, K 차원이 작은 경우 발생하는 메모리 대역폭 병목 현상을 완화하기 위해 k_unroll, vector_width, num_groups 등의 값을 조절하여 연산 밀도를 높이는 것이 중요합니다.

왜 이게 좋은가

이번 PR은 다음과 같은 이유로 매우 훌륭한 최적화라고 할 수 있습니다.

  1. 구체적인 하드웨어 및 모델 최적화: Qwen 3.8B 모델의 특정 레이어와 NVIDIA B300 (SM103) 아키텍처, 그리고 TP=1이라는 구체적인 조건에 맞춰 최적화된 커널 설정을 추가했습니다. 이는 범용적인 최적화보다 훨씬 높은 성능 향상을 기대할 수 있게 합니다.
  2. 성능 향상 수치: PR 설명에 따르면, cuBLAS 대비 최대 3.06배의 속도 향상을 달성했으며, 이는 디코드 단계에서 약 0.52ms의 지연 시간 감소로 이어집니다. 이러한 정량적인 성능 개선은 실제 서비스 환경에서 사용자 경험을 크게 향상시킬 수 있습니다.
    • Cold L2 vs Hot L2: PR의 테이블에서 볼 수 있듯이, L2 캐시가 차가운 상태(cold L2)와 따뜻한 상태(hot L2) 모두에서 성능 향상이 관찰되었습니다. 이는 최적화된 커널이 캐시 효율성 측면에서도 이점을 가지며, 반복적인 연산에서 더 큰 성능 향상을 보일 수 있음을 시사합니다.
  3. Skinny GEMM의 효과적인 활용: LLM에서 행렬 곱셈은 연산의 상당 부분을 차지합니다. 특히 트랜스포머 아키텍처의 Feed-Forward Network (FFN) 레이어 등에서 발생하는 'Skinny GEMM' (한쪽 차원이 매우 작은 행렬 곱셈)은 최적화하기 까다롭지만, 성능 병목의 주요 원인이 됩니다. vLLM은 이러한 Skinny GEMM에 특화된 커널 전략을 적용하여 성능을 극대화하고 있습니다.
  4. 테스트 및 검증: PR에서는 pytest를 이용한 커널 테스트와 실제 성능 측정 결과를 제시하며 변경 사항의 유효성을 입증했습니다. 이는 코드의 신뢰성을 높이는 중요한 요소입니다.

일반적인 교훈

  • 하드웨어 특성 이해의 중요성: 최신 GPU 아키텍처(예: NVIDIA B300)의 세부적인 성능 특성(코어 수, 메모리 계층 구조, 명령어 집합 등)을 깊이 이해하고 이를 커널 설계에 반영하는 것이 고성능 컴퓨팅의 핵심입니다.
  • 모델별 최적화: 특정 모델 아키텍처(예: Qwen)와 레이어의 수학적 특성(예: 행렬 차원 비율)을 분석하여 맞춤형 최적화를 수행하면 범용 최적화보다 훨씬 큰 성능 향상을 얻을 수 있습니다.
  • 정량적 검증의 필요성: 성능 개선 PR은 반드시 실제 측정 데이터를 통해 그 효과를 입증해야 합니다. 다양한 워크로드와 환경(cold/hot cache)에서의 성능 측정이 중요합니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 /ci run 명령어로 CI(Continuous Integration)를 실행하라는 요청이었습니다. 이는 코드 변경 사항이 CI 파이프라인을 통과하는지, 특히 새로운 커널 설정이 기존 테스트를 통과하는지 확인하기 위한 일반적인 절차입니다. 특별한 기술적 논의나 반박은 없었으며, 이는 제안된 변경 사항이 합리적이고 잘 검증되었다는 것을 간접적으로 시사합니다.

결론

이번 vLLM PR은 Qwen 3.8B 모델의 추론 성능을 NVIDIA B300 GPU에서 크게 향상시키는 중요한 작업을 수행했습니다. TP=1 환경에서 Skinny GEMM 커널 설정을 세밀하게 조정함으로써 cuBLAS 대비 최대 3배 이상의 성능 향상을 달성한 것은 vLLM의 최적화 역량을 잘 보여주는 사례입니다. 이러한 모델 및 하드웨어 특화 최적화는 대규모 언어 모델의 효율적인 서빙을 위해 필수적이며, 앞으로도 vLLM의 지속적인 발전에 기여할 것으로 기대됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글