본문으로 건너뛰기

[vllm] vLLM의 XPU 가속을 위한 MXFP8 GEMM 커널 도입 분석

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

들어가며

대규모 언어 모델(LLM)의 추론 성능을 높이기 위해 양자화(Quantization) 기술은 이제 필수적입니다. 그중에서도 **MXFP8(Microscaling Formats)**은 높은 정밀도를 유지하면서도 메모리 대역폭과 연산 속도를 획기적으로 개선할 수 있는 포맷으로 주목받고 있습니다.

최근 vLLM 프로젝트에 병합된 이 PR은 Intel의 GPU 라인업인 XPU 환경에서 MXFP8 GEMM(General Matrix Multiply)을 지원하기 위한 전용 커널을 도입했습니다. 기존에는 XPU에서 MXFP8 연산을 수행할 때 범용적인 에뮬레이션 커널에 의존해야 했으나, 이번 업데이트를 통해 하드웨어 가속을 직접 활용할 수 있게 되었습니다.

코드 분석: XPU 전용 MXFP8 커널 구현

이번 변경사항의 핵심은 XPUMxFp8LinearKernel 클래스의 구현과 이를 vLLM의 커널 레지스트리에 등록하는 과정에 있습니다.

1. 커널 레지스트리 등록

먼저, vLLM이 XPU 플랫폼에서 선형 연산을 수행할 때 MXFP8 커널을 선택할 수 있도록 시스템에 등록하는 과정이 포함되었습니다.

# vllm/model_executor/kernels/linear/__init__.py

# Before: XPU 플랫폼에 대한 MXFP8 커널 정의가 없었음

# After: XPUMxFp8LinearKernel 추가 및 우선순위 설정
    PlatformEnum.XPU: [
        XPUMxFp8LinearKernel,
        EmulationMxfp8LinearKernel,
    ],

여기서 주목할 점은 XPUMxFp8LinearKernel을 EmulationMxfp8LinearKernel보다 앞에 배치했다는 것입니다. 이는 시스템이 XPU 환경임을 감지하면 성능이 낮은 에뮬레이션 방식 대신 최적화된 전용 커널을 우선적으로 사용하도록 유도합니다.

2. 가중치 전처리 최적화

MXFP8 연산은 데이터 레이아웃에 민감합니다. 효율적인 연산을 위해 모델 로드 후 가중치(Weight)와 스케일(Scale) 텐서를 미리 변환하는 과정이 구현되었습니다.

# vllm/model_executor/kernels/linear/mxfp8/xpu.py

def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # weight_scale을 MXFP8 규격에 맞는 e8m0fnu 타입으로 뷰 변경 후 전치(Transpose)
    weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
    weight_scale = weight_scale.t().contiguous()
    
    # 가중치 본체도 전치하여 메모리 레이아웃 최적화
    replace_parameter(layer, "weight", layer.weight.t())
    replace_parameter(layer, "weight_scale", weight_scale.data)

process_weights_after_loading 메서드는 추론 루프가 시작되기 전 단 한 번 실행됩니다. 여기서 t().contiguous()를 통해 가중치를 미리 전치해 두면, 실제 추론 시(Runtime) 매번 수행해야 하는 연산 오버헤드를 제거할 수 있습니다.

3. 런타임 GEMM 실행

실제 연산 단계에서는 Intel XPU 전용 연산 라이브러리인 _xpu_C의 fp8_gemm을 호출합니다.

# vllm/model_executor/kernels/linear/mxfp8/xpu.py

def apply_weights(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    out_dtype = x.dtype
    # 입력 활성화 함수(Activation)를 MXFP8로 양자화
    x_fp8, x_scale = quant_mxfp8(x)
    
    # XPU 전용 가속 연산 호출
    return torch.ops._xpu_C.fp8_gemm(
        x_fp8,
        layer.weight,
        out_dtype,
        x_scale,
        layer.weight_scale,
        bias,
    )

이 부분은 W8A8(Weight 8-bit, Activation 8-bit) 연산을 수행합니다. 입력 텐서 x를 런타임에 양자화하고, 미리 준비된 가중치와 함께 하드웨어 가속기에서 GEMM을 수행하여 최종 결과를 산출합니다.

왜 이게 좋은 개선인가?

  1. 하드웨어 가속 활용: 기존의 에뮬레이션 방식은 FP8 연산을 지원하지 않는 하드웨어에서 소프트웨어적으로 흉내를 내는 방식이라 성능 저하가 심했습니다. 이번 PR은 Intel XPU의 전용 FP8 연산 유닛을 직접 활용함으로써 처리량(Throughput)을 극대화합니다.
  2. 메모리 효율성: MXFP8 포맷은 가중치를 8비트로 저장하므로, FP16 대비 모델 크기를 절반으로 줄일 수 있습니다. 이는 더 적은 GPU 메모리로 더 큰 모델을 서빙하거나, KV 캐시 공간을 더 많이 확보하는 데 유리합니다.
  3. 유연한 폴백(Fallback): VLLM_DISABLED_KERNELS 환경 변수를 통해 문제가 발생할 경우 수동으로 에뮬레이션 커널로 돌아갈 수 있는 안전장치를 마련했습니다. 이는 대규모 배포 환경에서 안정성을 보장하는 중요한 설계입니다.

결론

vLLM의 이번 업데이트는 NVIDIA 중심의 생태계를 넘어 Intel XPU 플랫폼에서도 고성능 양자화 추론이 가능함을 보여줍니다. 특히 transformer 5.4 및 compressed-tensors 0.14.0과의 호환성을 염두에 둔 설계는 최신 양자화 생태계의 흐름을 잘 반영하고 있습니다. Intel 가속기를 사용하는 엔지니어라면 이번 MXFP8 지원을 통해 상당한 성능 향상을 기대할 수 있을 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글