본문으로 건너뛰기

[vLLM] MXFP8/MXFP4: 마이크로스케일링 포맷 양자화

들어가며

Microscaling(MX) 포맷은 OCP(Open Compute Project)에서 제안한 차세대 양자화 표준이다. 기존 per-tensor/per-channel 스케일링의 한계를 극복하기 위해 블록 단위 스케일링을 도입한 것이 핵심이다. vLLM은 MXFP8(block-32 FP8)과 MXFP4(block-32 FP4) 두 가지를 모두 지원하며, 온라인 양자화 방식으로 bf16/fp16 체크포인트를 런타임에 변환한다.

관련 논문: Microscaling Data Formats for Deep Learning (arxiv 2310.10537)

공식 문서

vLLM 공식 문서: Quantization

핵심 구조/코드 분석

MXFP8: Block-32 스케일링

Mxfp8Config는 Fp8Config를 상속하면서 동적 활성화 스킴만 허용한다:

class Mxfp8Config(Fp8Config):
    def __init__(
        self,
        activation_scheme: str = "dynamic",
        ignored_layers: list[str] | None = None,
    ) -> None:
        if activation_scheme != "dynamic":
            raise ValueError("mxfp8 only supports dynamic activation scheme.")
        super().__init__(
            is_checkpoint_fp8_serialized=False,
            activation_scheme=activation_scheme,
            ignored_layers=ignored_layers,
            weight_block_size=None,
        )

핵심은 MXFP8_BLOCK_SIZE = 32라는 상수다. 가중치의 입력 차원이 반드시 32의 배수여야 하며, 가중치 로딩 후 mxfp8_e4m3_quantize로 FP8 값과 uint8 스케일을 생성한다:

def process_weights_after_loading(self, layer: Module) -> None:
    weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous())
    layer.input_scale = None
    replace_parameter(layer, "weight", weight_fp8.data)
    replace_parameter(layer, "weight_scale", weight_scale.data)

32개 원소마다 하나의 스케일 팩터가 할당되므로, per-tensor 대비 훨씬 정밀한 동적 범위 표현이 가능하다.

MXFP4: 4비트 MoE 특화

MXFP4는 현재 MoE 레이어 전용으로 구현되어 있다. Linear 레이어는 아직 UnquantizedLinearMethod로 폴백된다:

def get_quant_method(self, layer, prefix):
    if isinstance(layer, LinearBase):
        logger.debug_once(
            "MXFP4 linear layer is not implemented - falling back to "
            "UnquantizedLinearMethod.")
        return UnquantizedLinearMethod()
    elif isinstance(layer, FusedMoE):
        return Mxfp4MoEMethod(layer.moe_config)

MXFP4의 가중치는 torch.uint8로 저장되며, 2개의 FP4 값이 1바이트에 패킹된다. 따라서 가중치 shape에서 마지막 차원이 hidden_size // 2가 된다:

w13_weight = torch.nn.Parameter(
    torch.zeros(num_experts, 2 * intermediate_size_per_partition,
                hidden_size // 2, dtype=weight_dtype),
    requires_grad=False,
)

MoE 백엔드 선택

MXFP8 MoE는 select_mxfp8_moe_backend로, MXFP4 MoE는 select_mxfp4_moe_backend로 최적의 커널 백엔드를 자동 선택한다. MXFP4의 경우 FlashInfer+TRT-LLM, Triton 등 다양한 백엔드를 지원하며 SM100(Blackwell) 전용 커널도 포함되어 있다.

왜 이 설계인가

  1. Block-32 스케일링의 이점: 32개 원소 단위로 스케일을 적용하면 outlier가 다른 원소에 미치는 영향을 최소화할 수 있다. Per-tensor 스케일링은 하나의 극단값이 전체 텐서의 정밀도를 저하시키지만, block-32는 해당 블록 내에서만 영향을 받는다.

  2. 온라인 양자화: 별도의 양자화 체크포인트 없이 bf16 모델을 로딩하면서 실시간으로 변환한다. 이는 사용자 진입 장벽을 낮추고, HuggingFace에 공개된 대부분의 모델을 바로 사용할 수 있게 한다.

  3. MoE 우선 지원: MXFP4는 MoE 모델에서 가장 큰 메모리 절감 효과를 얻을 수 있다. 전문가 수가 수백 개인 모델(DeepSeek-V3 등)에서 4비트 양자화의 이점이 극대화된다.

논문 핵심 내용

Microscaling Data Formats for Deep Learning (2310.10537) 논문은 OCP(Open Compute Project)에서 제안한 MX 포맷의 설계와 검증 결과를 다룬다.

핵심 아이디어: 기존 per-tensor 스케일링은 하나의 이상치가 전체 텐서의 정밀도를 떨어뜨리는 문제가 있었다. MX 포맷은 **32개 원소 단위(block-32)**로 독립적인 스케일 팩터를 적용해서 이 문제를 해결했다. 하드웨어 효율성, 모델 정확도, 사용 편의성이라는 세 가지 상충하는 요구를 균형 잡는 것이 설계 목표였다.

논문의 가장 중요한 성과는 sub-8-bit(8비트 미만) 가중치, 활성화, 그래디언트로 생성형 언어 모델을 학습한 최초의 사례라는 점이다. 학습 레시피를 전혀 수정하지 않고 FP32의 드롭인 대체(drop-in replacement)로 사용할 수 있었다. 20개 이상의 벤치마크에서 검증했고, 최소한의 정확도 손실만 발생했다.

MX 포맷 원소 타입 블록 크기 스케일 타입 주 용도
MXFP8 E4M3/E5M2 32 E8M0 (uint8) 추론/학습 범용
MXFP6 E3M2/E2M3 32 E8M0 메모리 절약 학습
MXFP4 E2M1 32 E8M0 극한 메모리 절약
MXINT8 INT8 32 E8M0 정수 연산 선호 시

Block-32 스케일링의 장점은 per-tensor 대비 이상치 영향을 블록 내로 국한시키면서도, per-element 스케일링보다 메타데이터 오버헤드가 훨씬 적다는 거다. 32개 원소당 1바이트의 스케일만 추가되므로 약 3% 정도의 오버헤드밖에 없다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글