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

2026년 5월 13일수정: 2026년 5월 13일

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

들어가며

대규모 언어 모델(LLM)의 추론 효율성을 높이기 위한 양자화 기술은 이제 선택이 아닌 필수입니다. 그중에서도 MXFP4(Microscaling Formats)는 높은 압축률과 정밀도를 동시에 잡으려는 시도로 주목받고 있습니다.

최근 vLLM 프로젝트에 병합된 이 PR은 Intel GPU(XPU) 환경에서 MXFP4 W4A4(Weight 4-bit, Activation 4-bit) 경로를 활성화하고, 전용 GEMM(General Matrix Multiply) 커널을 추가하는 내용을 담고 있습니다. 이번 글에서는 시니어 엔지니어의 시각으로 이 변경사항이 왜 중요한지, 그리고 코드 레벨에서 어떤 최적화가 이루어졌는지 분석해 보겠습니다.

코드 분석: XPU를 위한 MXFP4 커널 구현

1. 커널 레지스트리 확장 (`vllm/model_executor/kernels/linear/init.py`)

먼저 vLLM의 선형 레이어 커널 관리 시스템에 XPU 전용 MXFP4 커널을 등록하는 작업이 수행되었습니다.

[Before]

# 기존에는 MarlinMxFp4 등 특정 벤더/플랫폼 위주의 커널만 존재
from vllm.model_executor.kernels.linear.mxfp4.marlin import (
    MarlinMxFp4LinearKernel,
)

[After]

from vllm.model_executor.kernels.linear.mxfp4.xpu import (
    XPUMxFp4LinearKernel,
)

# ... 중략 ...

    PlatformEnum.XPU: [
        XPUMxFp4LinearKernel,
    ],

이 변경을 통해 vLLM은 실행 플랫폼이 XPU일 경우 자동으로 XPUMxFp4LinearKernel을 선택하게 됩니다. 이는 전략 패턴(Strategy Pattern)을 활용하여 하드웨어별 최적화된 커널을 유연하게 교체할 수 있도록 설계된 vLLM의 아키텍처를 잘 따르고 있습니다.

2. XPUMxFp4LinearKernel 구현 (`vllm/model_executor/kernels/linear/mxfp4/xpu.py`)

가장 핵심이 되는 파일입니다. 여기서는 가중치 로딩 후의 전처리(Preprocessing)와 실제 추론 시의 연산(Inference) 로직이 포함되어 있습니다.

가중치 레이아웃 최적화

def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # Before: 일반적인 텐서 형태
    # After: XPU 하드웨어 가속에 최적화된 데이터 타입 및 레이아웃 변경
    weight = layer.weight.view(torch.float4_e2m1fn_x2)
    replace_parameter(layer, "weight", weight.data.t())

    weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
    weight_scale = weight_scale.t().contiguous()
    replace_parameter(layer, "weight_scale", weight_scale.data)

이 부분에서 주목할 점은 torch.float4_e2m1fn_x2와 torch.float8_e8m0fnu라는 특수한 데이터 타입을 사용한다는 것입니다.

Transposition (.t()): 행렬 곱셈 연산 시 메모리 접근 효율성(Memory Coalescing)을 높이기 위해 가중치를 미리 전치(Transpose)하여 저장합니다.
View Casting: 데이터를 물리적으로 복사하지 않고 MXFP4 규격에 맞는 뷰로 재해석하여 메모리 오버헤드를 최소화합니다.

런타임 양자화 및 GEMM 호출

def apply_weights(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    out_dtype = x.dtype
    # 입력 활성화 함수(Activation)를 실시간으로 MXFP4로 양자화
    x_fp4, x_blockscale = quant_mxfp4(x)
    
    # XPU 전용 컴파일된 연산자 호출
    return torch.ops._xpu_C.fp4_gemm(
        x_fp4,
        layer.weight,
        x_blockscale,
        layer.weight_scale,
        out_dtype,
        bias,
    )

추론 시점(apply_weights)에서는 입력값 x를 quant_mxfp4를 통해 실시간 양자화합니다. 이후 torch.ops._xpu_C.fp4_gemm이라는 하위 레벨의 C++ 커널을 호출합니다. 이는 Python 오버헤드를 피하고 Intel GPU의 하드웨어 가속기(XMX 등)를 최대한 활용하기 위한 설계입니다.

왜 이게 좋은 최적화인가?

메모리 대역폭 절감: 4-bit 양자화를 통해 모델 가중치 크기를 FP16 대비 1/4로 줄였습니다. 이는 LLM 추론의 병목 지점인 메모리 대역폭 제한(Memory Bandwidth Bound) 문제를 직접적으로 해결합니다.
하드웨어 특화 가속: 단순히 수치적으로 줄이는 것에 그치지 않고, XPU의 특수 데이터 타입(float4_e2m1fn_x2)을 활용하여 하드웨어 레벨의 병렬 연산을 이끌어냈습니다.
정밀도 유지: PR 설명에 언급되었듯 Yi30, Llama-3.2-1B 모델에서 정확도 테스트를 통과했습니다. MXFP4는 블록 단위 스케일링을 사용하므로 일반적인 4-bit 정수 양자화보다 정밀도 손실이 적습니다.

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer Mamba SSU 커널 최적화: Async State Prefetching과 Vectorized Load를 통한 성능 혁신
현재글 : [vllm] vLLM XPU 가속을 위한 MXFP4 W4A4 GEMM 커널 도입 분석
다음글 [vllm] vLLM의 Triton 통합 어텐션 커널에 Tensor Descriptor 최적화 도입