[vllm] vLLM의 FP8 Scaled MM 최적화: Padding 제거를 통한 20% 성능 향상

PR 링크: vllm-project/vllm#43706 상태: Merged | 변경: +57 / -9

들어가며

vLLM은 대규모 언어 모델(LLM) 추론을 위한 고성능 엔진입니다. 최근 vLLM의 scaled_mm 커널에서 FP8 연산 시 불필요한 Padding을 제거하여 성능을 최적화하는 PR이 병합되었습니다. 기존에는 정렬되지 않은 입력에 대해 일괄적으로 Padding을 적용했으나, 이번 개선을 통해 정렬된 입력(4-aligned)에 대해서는 직접 연산을 수행함으로써 약 20%의 성능 향상을 달성했습니다.

코드 분석

1. vllm/model_executor/kernels/linear/scaled_mm/cutlass.py 수정

핵심 변경 사항은 apply_block_scaled_mm 메서드에서 호출하는 연산자를 padded_cutlass에서 dynamic_padded_cutlass로 변경한 것입니다. 또한, _dynamic_padded_cutlass 함수를 도입하여 입력 데이터의 정렬 상태에 따라 분기 처리를 수행합니다.

Before

if self.is_hopper:
    return torch.ops.vllm.padded_cutlass(
        A, B, As, Bs, list(self.weight_group_shape), out_dtype,
    )

After

if self.is_hopper:
    return torch.ops.vllm.dynamic_padded_cutlass(
        A, B, As, Bs, list(self.weight_group_shape), out_dtype,
    )

새로 추가된 _dynamic_padded_cutlass는 torch.compiler.is_compiling() 환경과 런타임 환경 모두에서 qx.shape[0] % 4 != 0 조건을 확인합니다. 4로 나누어떨어지는 경우(정렬된 경우)에는 run_direct를 호출하여 Padding 오버헤드를 완전히 제거합니다.

if qx.shape[0] % 4 != 0:
    return run_padded(qx, weight, x_scale, weight_scale)
return run_direct(qx, weight, x_scale, weight_scale)

왜 이게 좋은가

이 최적화의 핵심은 '불필요한 메모리 복사 및 연산 방지'입니다. GPU 커널에서 Padding은 데이터 정렬을 맞추기 위해 필수적일 때가 많지만, 데이터가 이미 정렬된 상태라면 Padding을 추가하는 것은 메모리 대역폭과 연산 자원의 낭비입니다.

성능 테스트 결과, 다양한 M(Batch size) 값에서 일관되게 약 1.2배(20%)의 속도 향상을 보였습니다. 특히 m=4, n=16384, k=1024와 같은 조건에서 1.22x의 속도 향상을 기록했습니다.

교훈

1. 조건부 실행의 힘: 모든 상황에 일괄적인 로직을 적용하기보다, 데이터의 특성(Alignment)을 파악하여 분기 처리하는 것이 성능 최적화의 기본입니다.
2. Torch Compile 고려: torch.cond를 사용하여 컴파일 타임과 런타임 모두에서 최적화된 경로를 선택하도록 설계한 점은 최신 PyTorch 환경에서의 모범 사례입니다.
3. 리뷰어 피드백의 중요성: 초기 제안보다 더 깔끔한 분기 처리를 위해 리뷰어(lgeiger)의 제안을 수용하여 코드를 개선한 점이 돋보입니다.

결론

이번 최적화는 vLLM의 FP8 추론 성능을 한 단계 끌어올렸습니다. 하드웨어의 특성을 깊이 이해하고, 불필요한 오버헤드를 제거하는 것만으로도 상당한 성능 이득을 얻을 수 있음을 보여주는 좋은 사례입니다.

참고 자료

• https://pytorch.org/docs/stable/torch.compiler.html
• https://docs.nvidia.com/cuda/cutlass/index.html

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