[vLLM] Compressed Tensors: 양자화+희소성 통합 프레임워크

들어가며

Compressed Tensors는 Neural Magic이 개발한 양자화 및 희소성 통합 프레임워크다. 단일 설정 포맷으로 W8A8 Int8, W8A8 FP8, W4A16, W4A8, 2:4 구조적 희소성 등 다양한 압축 조합을 표현할 수 있다. vLLM은 이를 통해 SparseGPT, GPTQ, SmoothQuant 등으로 압축된 모델을 통합적으로 로딩한다.

공식 문서

vLLM 공식 문서: LLM Compressor

핵심 구조/코드 분석

CompressedTensorsConfig

설정 클래스는 타겟-스킴 매핑과 희소성 설정을 모두 포함한다:

class CompressedTensorsConfig(QuantizationConfig):
    def __init__(
        self,
        target_scheme_map: dict[str, Any],
        ignore: list[str],
        quant_format: str,
        sparsity_scheme_map: dict[str, SparsityCompressionConfig],
        sparsity_ignore_list: list[str],
        kv_cache_scheme: dict[str, Any] | None = None,
        transform_config: dict[str, Any] | None = None,
    ):

target_scheme_map은 모듈 이름 패턴을 키로, 해당 모듈의 양자화 스킴을 값으로 가진다. 이를 통해 레이어별로 서로 다른 양자화 전략을 적용할 수 있다.

다양한 Scheme 지원

vLLM은 Compressed Tensors의 스킴별로 전용 구현 클래스를 제공한다:

from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
    CompressedTensors24,           # 2:4 구조적 희소성
    CompressedTensorsW8A8Fp8,      # FP8 가중치+활성화
    CompressedTensorsW8A8Int8,     # INT8 가중치+활성화
    CompressedTensorsW8A16Fp8,     # FP8 가중치, FP16 활성화
    CompressedTensorsW4A8Fp8,      # FP4 가중치, FP8 활성화
    CompressedTensorsW4A8Int,      # INT4 가중치, INT8 활성화
    CompressedTensorsW4A16Fp4,     # FP4 가중치, FP16 활성화
    CompressedTensorsW4A16Mxfp4,   # MXFP4 가중치, FP16 활성화
    CompressedTensorsWNA16,        # WNA16 (Marlin 커널)
)

이 구조의 장점은 새로운 양자화 조합이 추가되어도 스킴 클래스만 하나 추가하면 된다는 것이다.

레이어별 스킴 매칭

get_quant_method에서 레이어 prefix를 기반으로 적절한 스킴을 찾는다:

def get_quant_method(self, layer, prefix):
    if isinstance(layer, LinearBase):
        scheme = self._get_scheme_from_parts(layer, prefix)
        if scheme is None:
            return UnquantizedLinearMethod()
        return CompressedTensorsLinearMethod(scheme)
    elif isinstance(layer, FusedMoE):
        return CompressedTensorsMoEMethod(self, prefix)
    elif isinstance(layer, Attention):
        if self.kv_cache_scheme is not None:
            return BaseKVCacheMethod(self)

find_matched_target 함수가 prefix와 target_scheme_map의 패턴을 매칭하여, 해당 레이어에 적용할 양자화 전략을 결정한다.

Transform 지원

최근 추가된 transform_config는 가중치 변환(예: Hadamard rotation)을 지원한다:

from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import (
    CompressedTensorsLinearTransformMethod,
    get_linear_transform_schemes,
)

QuaRot, SpinQuant 등의 회전 기반 양자화 기법에서 사용되며, 양자화 전후로 가중치에 직교 변환을 적용한다.

왜 이 설계인가

1. 통합 추상화: 양자화와 희소성을 별도 시스템이 아닌 하나의 설정 포맷으로 통합한다. 이는 2:4 희소성 + INT8 양자화 같은 복합 압축을 자연스럽게 표현할 수 있게 한다.

2. Neural Magic 생태계 연동: SparseML로 생성한 모델을 별도 변환 없이 바로 로딩할 수 있다. CompressionFormat 열거형이 다양한 압축 포맷(pack-quantized, marlin-24, float-quantized 등)을 지원한다.

3. 확장 가능한 스킴 구조: CompressedTensorsScheme 기반 클래스를 상속하여 새로운 양자화 조합을 쉽게 추가할 수 있다. W4A4 NvFP4, W4A16 MXFP4 등 최신 포맷도 이 구조로 추가되었다.

참고 자료

• vLLM 공식 문서 - Quantization
• 소스 코드: vllm/model_executor/layers/quantization/compressed_tensors/