본문으로 건너뛰기

[vLLM] Marlin Kernels: 양자화 고속 GEMM 커널

들어가며

Marlin은 GPTQ 양자화된 모델을 거의 FP16 수준의 속도로 추론할 수 있게 해주는 고성능 GEMM(General Matrix Multiply) 커널이다. 일반적인 GPTQ 추론은 역양자화 오버헤드로 인해 FP16 대비 느려지는 경우가 있지만, Marlin은 메모리 대역폭에 최적화된 커널 설계로 이를 해결한다. vLLM은 호환되는 GPTQ 체크포인트를 자동으로 Marlin 포맷으로 변환하여 사용한다.

관련 논문: Marlin: Mixed-Precision Auto-Regressive Parallel INference on GPUs (arxiv 2408.11743)

공식 문서

vLLM 공식 문서: Quantization

핵심 구조/코드 분석

GPTQMarlinConfig: 자동 변환 감지

GPTQMarlinConfig는 GPTQ 체크포인트가 Marlin 호환인지 자동으로 감지한다:

TYPE_MAP = {
    (4, True): scalar_types.uint4b8,
    (8, True): scalar_types.uint8b128,
}

4비트 대칭, 8비트 대칭 양자화만 지원하며, override_quantization_method에서 런타임 변환을 결정한다:

@classmethod
def override_quantization_method(cls, hf_quant_cfg, user_quant):
    can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg)
    is_valid_user_quant = (
        user_quant is None or user_quant == "marlin"
        or user_quant == "gptq_marlin"
    )
    if can_convert and is_valid_user_quant:
        return cls.get_name()

사용자가 --quantization gptq를 명시하지 않는 한, 호환되는 GPTQ 모델은 자동으로 gptq_marlin으로 전환된다.

Dynamic Per-Module 양자화 (GPTQModel)

GPTQModel의 dynamic 설정으로 레이어별 다른 양자화 비트를 적용할 수 있다:

# 레이어 10-15는 8비트, 레이어 16-21은 8비트+group_size 64
dynamic = {
    r"+:.*\.(?:1[0-5])\..*": {"bits": 8},
    r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64},
    r"-:.*\.moe\..*": {},  # MoE 레이어는 양자화 스킵
}

정규식 기반의 매칭으로, +: 접두사는 양자화 적용, -: 접두사는 양자화 스킵을 의미한다.

MoE 지원: fused_marlin_moe

MoE 모델에서도 Marlin 커널을 사용할 수 있다. 단, 레이어의 intermediate_size가 Marlin의 정렬 요구사항을 만족해야 한다:

def get_quant_method(self, layer, prefix):
    if isinstance(layer, FusedMoE):
        if not check_moe_marlin_supports_layer(layer, self.group_size):
            logger.warning_once(
                f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
                "Falling back to Moe WNA16 kernels."
            )
            return MoeWNA16Config.from_config(self.full_config) \
                .get_quant_method(layer, prefix)

Marlin이 지원하지 않는 shape의 MoE 레이어는 WNA16 커널로 자동 폴백된다.

Marlin 호환성 검증

@classmethod
def is_gptq_marlin_compatible(cls, quant_config):
    quant_method = quant_config.get("quant_method", "").lower()
    num_bits = quant_config.get("bits")
    group_size = quant_config.get("group_size")
    sym = quant_config.get("sym")
    desc_act = quant_config.get("desc_act")
    if quant_method != "gptq":
        return False
    if (num_bits, sym) not in cls.TYPE_MAP:
        return False
    return check_marlin_supported(
        quant_type=cls.TYPE_MAP[(num_bits, sym)],
        group_size=group_size
    )

GPTQ 메서드, 지원되는 비트/대칭 조합, group_size 호환성을 모두 검증한다.

왜 이 설계인가

  1. 투명한 변환: 사용자는 GPTQ 모델을 그대로 로딩하면 되고, vLLM이 Marlin 호환성을 자동 감지하여 최적의 커널로 전환한다. 별도의 사전 변환 과정이 불필요하다.

  2. 폴백 전략: Marlin이 지원하지 않는 경우(비대칭 양자화, 특정 shape 등)에는 WNA16 등 대체 커널로 자동 폴백하여, 어떤 GPTQ 모델이든 동작하도록 보장한다.

  3. SM75+ 최소 요구: Turing 아키텍처(RTX 20 시리즈) 이상에서 동작하며, 이는 INT4 텐서 코어를 활용하기 위한 최소 요구사항이다.

논문 핵심 내용

Marlin: Mixed-Precision Auto-Regressive Parallel INference on GPUs (2408.11743) 논문은 양자화된 가중치의 배치 추론에서도 높은 속도 향상을 유지하는 GEMM 커널을 제안했다.

핵심 아이디어: 기존 양자화 추론은 단일 사용자(배치 1)에서는 메모리 대역폭 절감으로 빠르지만, 배치 크기가 커지면 연산 병목으로 전환되어 속도 이점이 사라졌다. Marlin은 비동기 메모리 접근, 복잡한 태스크 스케줄링과 파이프라이닝, 맞춤형 양자화 지원을 결합해서 배치 16-32에서도 거의 최대(4x) 양자화 속도 향상을 달성했다.

Llama-2-7B 엔드투엔드 추론 속도 향상 (NVIDIA A10)

배치 크기 Marlin 속도 향상
1 2.93x
16 2.74x
32 2.26x
128 1.20x

커널 수준 성능 (72k x 18k 행렬, 4비트)

배치 크기 속도 향상 병목 유형
16-32 ~3.87x 메모리 대역폭 병목
64 ~2.5x 전환 구간
128 ~1.5x 연산 병목

vLLM 서빙 엔진과 통합했을 때 최대 2.8x 엔드투엔드 속도 향상을 달성했다. Sparse-Marlin 변형(NVIDIA 2:4 sparsity 추가 적용)은 3.2x까지 도달했다. 흥미로운 점은 RTX 3090 같은 소비자용 GPU에서 A100보다 더 높은 속도 향상을 보였다는 거다. 이는 소비자용 GPU의 메모리 대역폭 대비 연산 비율이 낮아서 메모리 병목이 더 심하기 때문이다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글