[vllm] vLLM 컴파일 최적화: Transformers 모델을 위한 FusedAddRMSNorm 도입

PR 링크: vllm-project/vllm#48757 상태: Merged | 변경: +298 / -2

들어가며

최신 대규모 언어 모델(LLM)의 추론 성능을 극대화하기 위해 vLLM은 torch.compile 기반의 컴파일 파이프라인을 적극적으로 활용하고 있습니다. 하지만 Transformers 기반 모델들은 구조적으로 residual add와 RMSNorm을 별도의 연산으로 수행하는 경우가 많습니다. 이로 인해 메모리 대역폭을 낭비하게 되고, 이후 이어지는 연산들과의 융합(fusion) 기회를 놓치는 문제가 발생합니다. 본 PR은 이러한 연산들을 fused_add_rms_norm으로 정규화(canonicalize)하고, 중간에 끼어든 reshape 연산을 최적화하여 전체적인 추론 성능을 개선합니다.

코드 분석

1. vllm/compilation/passes/fusion/add_rms_fusion.py

핵심은 VllmPatternReplacement를 사용하여 기존의 분리된 연산 패턴을 하나의 융합된 연산으로 교체하는 것입니다.

Before (분리된 연산):

residual_out = residual + branch if self.residual_first else branch + residual
rms = vllm.ir.ops.rms_norm(residual_out, weight, self.epsilon)

After (융합된 연산):

return vllm.ir.ops.fused_add_rms_norm(branch, residual, weight, self.epsilon)

또한, RMSNormReshapeFusionPass를 통해 RMSNorm 이후의 reshape 연산을 앞으로 이동시켜, fused_add_rms_norm이 더 넓은 범위의 패턴과 매칭될 수 있도록 했습니다. 이는 특히 AllReduce + RMSNorm과 같은 복합 패턴을 발견하는 데 결정적인 역할을 합니다.

2. vllm/compilation/passes/pass_manager.py

새롭게 정의된 패스들을 PassManager에 등록하여 컴파일 시점에 자동으로 적용되도록 했습니다. 특히 model_config.using_transformers_backend()를 통해 Transformers 기반 모델에서만 이 최적화가 활성화되도록 게이트를 설정했습니다.

if enable_transformers_norm_canonicalization:
    self.passes += [AddRMSNormFusionPass(config)]
    # ...
    self.passes += [RMSNormReshapeFusionPass(config)]

왜 이게 좋은가

이 최적화는 메모리 접근 횟수를 획기적으로 줄여줍니다. Add와 RMSNorm을 별도로 수행하면 중간 결과값을 메모리에 쓰고 다시 읽어야 하지만, 융합된 커널은 레지스터 수준에서 이를 처리합니다.

• 성능 수치: Qwen3-32B-FP8 모델(8xH100, TP8) 환경에서 출력 토큰 처리량이 +18.3% 증가하고, TPOT(Time Per Output Token)가 -15.7% 감소하는 놀라운 성과를 보였습니다.
• 교훈: 컴파일러 최적화에서 '패턴 매칭'의 순서는 매우 중요합니다. AR+RMS 융합이 RMS+Quant보다 우선순위를 가짐으로써 더 큰 병목을 해결할 수 있었으며, 중간의 reshape와 같은 사소한 연산을 정규화하는 것만으로도 전체 그래프의 융합 가능성을 크게 높일 수 있음을 보여줍니다.

리뷰어 피드백 반영

리뷰 과정에서 branch + residual과 residual + branch 순서에 상관없이 융합이 가능하도록 일반화하는 논의가 있었으며, 이는 AddRMSNormPattern 클래스에서 residual_first 파라미터를 통해 구현되었습니다. 또한, 컴파일 시간 증가(약 70%)라는 트레이드오프가 있지만, 추론 성능 향상폭이 이를 충분히 상쇄한다는 점이 확인되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.vllm.ai/en/latest/models/engine_args.html

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