본문으로 건너뛰기

[vllm] vLLM Qwen3.5 GDN 최적화: `einops.rearrange`를 `torch.flatten`으로 교체하여 20배 성능 향상!

PR 링크: vllm-project/vllm#42311 상태: Merged | 변경: +3 / -3

들어가며

vLLM은 대규모 언어 모델(LLM) 추론을 위한 고성능 서빙 엔진으로, 최적화된 커널과 효율적인 메모리 관리를 통해 뛰어난 처리량을 제공합니다. 이러한 vLLM의 성능은 작은 부분의 최적화가 전체 시스템에 미치는 영향이 매우 크다는 것을 의미합니다. 이번에 분석할 PR은 Qwen3.5 모델의 GDN(Gated Diffusion Network) 레이어에서 발생하는 사소하지만 중요한 성능 병목을 해결하는 최적화에 관한 것입니다.

기존 Qwen3.5 GDN 출력 프로젝션(output projection)에서는 einops 라이브러리의 rearrange 함수를 사용하여 텐서의 차원을 재구성했습니다. einops는 텐서 조작을 위한 강력하고 가독성 높은 도구이지만, 내부적으로 문자열 파싱 및 유효성 검사 로직을 포함하고 있어 Python 레벨의 오버헤드를 발생시킵니다. 모델의 forward pass와 같이 빈번하게 호출되는 코드 경로에서 이러한 오버헤드는 누적되어 상당한 성능 저하를 초래할 수 있습니다. 이 PR은 이 부분을 PyTorch의 네이티브 flatten 연산으로 대체하여 불필요한 오버헤드를 제거하고 성능을 크게 향상시키는 것을 목표로 합니다.

코드 분석: einops.rearrange를 torch.flatten으로

핵심 변경사항은 vllm/model_executor/layers/mamba/gdn_linear_attn.py 파일에 있습니다. 이 파일은 GDN Linear Attention 레이어의 구현을 담당하며, 특히 _output_projection, forward_hip, forward_cpu 세 가지 메서드에서 텐서의 마지막 두 차원(h와 d)을 병합하는 로직이 수정되었습니다.

변경 전 (Before)

        core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")

변경 후 (After)

        core_attn_out = core_attn_out.flatten(-2)  # ... h d -> ... (h d)

이 변경은 텐서의 마지막 두 차원 h와 d를 하나의 차원으로 병합하는 동일한 논리적 작업을 수행합니다. einops.rearrange는 "... h d -> ... (h d)"와 같은 문자열 패턴을 통해 차원 재구성을 정의하는 반면, torch.flatten(-2)는 마지막에서 두 번째 차원(-2)부터 끝까지(-1)를 평탄화하라는 의미로, ... h d를 ... (h d)로 만드는 것과 정확히 일치합니다.

왜 이게 좋은 최적화인가?

이 최적화는 겉보기에는 간단한 변경처럼 보이지만, 딥러닝 모델의 성능에 지대한 영향을 미칩니다. 그 이유는 다음과 같습니다.

  1. Python 오버헤드 제거: einops.rearrange는 유연하고 가독성이 높지만, 내부적으로 문자열 파싱 및 유효성 검사 로직을 포함하여 Python 레벨의 오버헤드를 발생시킵니다. 이는 특히 모델의 forward pass와 같이 빈번하게 호출되는 코드 경로에서 누적되어 상당한 성능 저하를 초래할 수 있습니다. 반면 torch.flatten은 PyTorch의 네이티브 C/CUDA 구현으로, 이러한 Python 오버헤드 없이 직접적이고 최적화된 방식으로 텐서 차원 병합을 수행합니다.

  2. 프로파일링을 통한 병목 식별: PR 설명에 포함된 프로파일링 결과는 이 최적화의 효과를 명확히 보여줍니다.

    • CUDA (eager mode) 프로파일링 결과:

      Event Calls Avg Min Max
      einops/einops.py(561): rearrange 576 23.261 us 12.949 us 103.882 us
      aten::flatten 576 2.828 us 1.486 us 94.331 us

      rearrange 대비 flatten이 평균 8.23배 더 빠릅니다.

    • CPU (eager mode) 프로파일링 결과:

      Event Calls Avg Min Max
      einops/einops.py(561): rearrange 144 41.741 us 29.737 us 106.243 us
      aten::flatten 144 1.949 us 1.203 us 19.260 us

      rearrange 대비 flatten이 평균 21.41배 더 빠릅니다.

    CPU 환경에서 훨씬 더 큰 성능 향상을 보이는 것은 Python 인터프리터 오버헤드가 CPU 바운드 작업에서 더 두드러지게 나타나기 때문입니다.

  3. 종합 벤치마크 결과 (End-to-End Serving Benchmark): 리뷰어 vadiklyutiy와 ZJY0516은 CUDA Graph 환경에서의 성능과 eager mode에서의 성능 개선에 대한 질문을 제기했습니다. PR 작성자는 enforce-eager 모드로 실행된 vLLM 서빙 벤치마크 결과를 추가하여, 이 최적화가 실제 서비스 환경에서도 유의미한 영향을 미 미침을 입증했습니다. (RTX 4090, Qwen3.5-0.8B, batch size 600, input/output len 8)

    Metric Before (main) After (PR) Improvement
    Request throughput (req/s) 435.66 531.50 +22%
    Output token throughput (tok/s) 3485.32 4251.99 +22%
    Mean TPOT (ms) 88.38 67.59 -23.5%

    이 결과는 단일 연산의 최적화가 전체 시스템의 처리량과 지연 시간에 직접적인 영향을 미 미침을 보여줍니다. 특히 Mean TPOT(Time per Output Token)의 23.5% 감소는 LLM 서빙과 같이 지연 시간에 민감한 애플리케이션에서 매우 중요한 개선입니다.

  4. 정확성 유지: 이 최적화는 성능 향상과 더불어 모델의 정확성을 그대로 유지합니다. PR에 포함된 PPL(Perplexity) 테스트 결과는 einops.rearrange와 flatten 사용 시 거의 동일한 PPL 값을 보여주어, 변경이 모델의 출력에 영향을 주지 않음을 확인시켜 줍니다.

    Method our PPL HF expected PPL Relative difference
    einops.rearrange 19.5098 19.3885 0.6252%
    flatten 19.5100 19.3885 0.6263%

일반적인 교훈

이 PR은 다음과 같은 중요한 개발 원칙을 상기시켜 줍니다.

  • 네이티브 연산 우선: 성능에 민감한 딥러닝 코드에서는 가능한 한 PyTorch의 네이티브 연산을 활용하는 것이 좋습니다. 외부 라이브러리는 편의성과 가독성을 제공하지만, 잠재적인 성능 오버헤드를 항상 염두에 두어야 합니다.
  • 프로파일링의 중요성: 작은 코드 변경이 큰 성능 향상으로 이어질 수 있음을 보여줍니다. 프로파일링은 이러한 병목 지점을 정확히 식별하고, 최적화 노력을 집중할 수 있게 해주는 필수적인 도구입니다.
  • 미세 최적화의 누적 효과: 단일 연산의 최적화가 전체 시스템 성능에 미치는 영향은 미미해 보일 수 있지만, LLM과 같이 복잡하고 반복적인 연산이 많은 시스템에서는 이러한 미세 최적화가 누적되어 상당한 개선을 가져올 수 있습니다.

결론

Qwen3.5 GDN 레이어의 einops.rearrange를 torch.flatten으로 교체한 이 PR은 vLLM의 성능을 크게 향상시키는 모범적인 최적화 사례입니다. 이는 PyTorch의 네이티브 연산을 활용하고 불필요한 Python 오버헤드를 제거함으로써, 단일 연산에서 최대 21배의 속도 향상과 전체 시스템 처리량에서 20% 이상의 개선을 달성했습니다. 이러한 최적화는 vLLM이 LLM 서빙 분야에서 선두를 유지하는 데 기여하는 중요한 요소입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글