[vllm] vLLM XPU 백엔드 최적화: W8A8 및 W8A16 FP8 Linear 커널 도입

PR 링크: vllm-project/vllm#43645 상태: Merged | 변경: +188 / -12

들어가며

vLLM은 다양한 하드웨어 가속기를 지원하기 위해 노력하고 있습니다. 특히 Intel XPU 환경에서 FP8 연산은 모델 추론 성능을 극대화하는 핵심 요소입니다. 기존의 XPUFP8ScaledMMLinearKernel은 제한적인 양자화 세분성(granularity) 지원으로 인해 최적의 성능을 내기 어려웠습니다. 이번 PR에서는 XPUW8A8FP8LinearKernel과 XPUW8A16FP8LinearKernel을 도입하여, 가중치와 활성화 값의 양자화 방식에 따른 유연한 커널 선택을 가능하게 했습니다.

코드 분석

1. 커널 구현 및 분기 (vllm/model_executor/kernels/linear/scaled_mm/xpu.py)

기존의 단일 커널을 W8A8(가중치 8비트, 활성화 8비트)과 W8A16(가중치 8비트, 활성화 16비트)으로 분리했습니다. 이를 통해 활성화 양자화 유무에 따라 최적의 커널을 선택할 수 있게 되었습니다.

Before:

class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
    # ... 기존의 단일 커널 구현

After:

class XPUW8A8FP8LinearKernel(FP8ScaledMMLinearKernel):
    # W8A8 전용 구현, 동적/정적 양자화 지원
    def apply_scaled_mm(self, ...):
        return torch.ops._xpu_C.fp8_gemm(A, B, out_dtype, As, Bs, bias)

class XPUW8A16FP8LinearKernel(FP8ScaledMMLinearKernel):
    # W8A16 전용 구현

2. 백엔드 설정 확장 (vllm/config/kernel.py)

사용자가 명시적으로 커널 백엔드를 선택할 수 있도록 --linear-backend 옵션을 확장했습니다.

# vllm/config/kernel.py
_LINEAR_BACKENDS = [
    "xpu",
    "xpu_woq",
]

xpu는 W8A8을 포함한 기본 커널을, xpu_woq는 Weight-Only Quantization(W8A16)을 타겟팅합니다.

왜 이게 좋은가

이번 최적화의 핵심은 '세분성(Granularity)에 따른 유연한 커널 선택'입니다. 리뷰 과정에서 확인된 성능 지표에 따르면, Per-token 양자화는 높은 동시성 환경에서 유리하고, Per-tensor 양자화는 대규모 배치 프리필(prefill) 단계에서 유리합니다.

• 성능 수치: meta-llama/Llama-3.2-1B-Instruct 모델 기준, Per-token 방식은 고부하 환경에서 out-tok/s가 약 24% 향상되었으며, Per-tensor 방식은 대규모 프리필에서 약 25%의 성능 이득을 보였습니다.
• 교훈: 하드웨어 가속기(XPU) 환경에서는 단일 커널로 모든 상황을 커버하기보다, 데이터의 양자화 세분성에 맞춰 커널을 분리하는 것이 추론 지연 시간(latency)과 처리량(throughput)을 모두 잡는 핵심 전략임을 알 수 있습니다.

리뷰어 피드백 반영

리뷰 과정에서 chaojun-zhang은 성능 비교를 통해 Per-token과 Per-tensor가 워크로드에 따라 장단점이 있음을 증명했습니다. 이에 따라 기본값을 유지하면서도 사용자가 --linear-backend를 통해 명시적으로 최적의 커널을 선택할 수 있도록 인터페이스를 개선했습니다.

참고 자료

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

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