본문으로 건너뛰기

[SGLang] Compressed Tensors: 통합 양자화 프레임워크

들어가며

양자화 방식은 FP8, INT8, INT4, FP4 등 다양하며, 각각 가중치/활성화 정밀도 조합이 다르다. SGLang의 Compressed Tensors는 이런 다양한 양자화 스킴을 하나의 통합 프레임워크로 관리한다. python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py에서 설정 파싱, 스킴 자동 선택, 디스패치 로직을 구현한다.

구조도

CompressedTensorsConfig
├── target_scheme_map: Dict[target_name → {weights, input_activations}]
├── ignore: List[str]          (양자화 제외 레이어)
├── quant_format: str          (압축 포맷)
├── sparsity_scheme_map        (희소성 설정)
├── linear_fp8_config          (혼합 양자화용 FP8 설정)
└── packed_modules_mapping     (fused 모듈 매핑)

스킴 자동 선택 흐름:
get_quant_method(layer, prefix)
├── LinearBase
│   ├── linear_fp8_config 있음? → Fp8LinearMethod
│   └── get_linear_scheme()
│       ├── W8A8 FP8  → CompressedTensorsW8A8Fp8
│       ├── W8A8 INT8 → CompressedTensorsW8A8Int8
│       ├── W4A4 FP4  → CompressedTensorsW4A4Fp4
│       ├── W8A16 FP8 → CompressedTensorsW8A16Fp8
│       └── WNA16     → CompressedTensorsWNA16
└── FusedMoE
    └── get_moe_scheme()
        ├── W8A8 FP8 MoE  → CompressedTensorsW8A8Fp8MoE
        ├── W4A4 FP4 MoE  → CompressedTensorsW4A4Nvfp4MoE
        ├── MxINT4 MoE    → CompressedTensorsMxInt4MoE
        └── WNA16 MoE     → CompressedTensorsWNA16MoE

핵심 코드 분석

1. CompressedTensorsConfig: 중앙 설정 관리

설정은 HuggingFace quantization_config에서 파싱된다.

class CompressedTensorsConfig(QuantizationConfig):
    def __init__(self, target_scheme_map, ignore, quant_format,
                 sparsity_scheme_map, sparsity_ignore_list,
                 kv_cache_scheme=None, config=None,
                 packed_modules_mapping=None,
                 linear_fp8_config=None):
        self.target_scheme_map = target_scheme_map
        self.ignore = ignore
        self.quant_format = quant_format
        self.linear_fp8_config = linear_fp8_config

linear_fp8_config는 혼합 양자화 시나리오를 위한 것이다. 예를 들어 MoE 전문가는 INT4, Linear 레이어는 FP8로 양자화하는 경우에 사용한다.

2. 설정 파싱: config_groups

HuggingFace 양자화 설정의 config_groups를 파싱하여 target_scheme_map을 구성한다.

@classmethod
def _quantization_scheme_map_from_config(cls, config):
    target_scheme_map = dict()
    config_groups = config.get("config_groups", dict())
    for _, quant_config in config_groups.items():
        targets = quant_config.get("targets")
        for target in targets:
            target_scheme_map[target] = {}
            target_scheme_map[target]["weights"] = \
                QuantizationArgs.model_validate(
                    quant_config.get("weights"))
            if is_activation_quantization_format(quant_format):
                input_activations = quant_config.get("input_activations")
                if input_activations:
                    target_scheme_map[target]["input_activations"] = \
                        QuantizationArgs.model_validate(input_activations)

QuantizationArgs는 pydantic 모델로, 양자화 전략(TENSOR/CHANNEL/BLOCK), 비트 수, 대칭 여부 등을 검증한다.

3. 레이어 타입별 디스패치

get_quant_method는 레이어 타입에 따라 LinearMethod 또는 MoEMethod를 반환한다.

def get_quant_method(self, layer, prefix):
    if isinstance(layer, LinearBase):
        if self.linear_fp8_config is not None:
            return Fp8LinearMethod(self.linear_fp8_config)
        scheme = self.get_linear_scheme(layer=layer, layer_name=prefix)
        if scheme is None:
            return UnquantizedLinearMethod()
        layer.scheme = scheme
        return CompressedTensorsLinearMethod(self)

    if isinstance(layer, FusedMoE):
        layer.scheme = self.get_moe_scheme(layer=layer, layer_name=prefix)
        if layer.scheme is None:
            return UnquantizedFusedMoEMethod(...)
        return CompressedTensorsFusedMoEMethod(self)

스킴 객체가 layer.scheme에 저장되어, 이후 CompressedTensorsLinearMethod가 실제 연산 시 참조한다.

4. FusedMoE 타겟 자동 확장

Linear 타겟 설정이 있으면 FusedMoE에도 자동 적용한다.

def _add_fused_moe_to_target_scheme_map(self):
    if ("Linear" not in self.target_scheme_map
            or "FusedMoE" in self.target_scheme_map):
        return
    self.target_scheme_map["FusedMoE"] = self.target_scheme_map["Linear"]
    self.target_scheme_map["DeepEPMoE"] = self.target_scheme_map["Linear"]

MoE 전문가의 Linear 레이어가 FusedMoE로 합쳐지므로, Linear 설정을 FusedMoE에도 전파한다.

5. 혼합 양자화: linear_fp8_config

MoE 전문가는 INT4, 나머지 Linear는 FP8로 양자화하는 혼합 구성을 지원한다.

@classmethod
def from_config(cls, config):
    linear_fp8_config = None
    if "linear_fp8_config" in config:
        from sglang.srt.layers.quantization.fp8 import Fp8Config
        fp8_cfg = config["linear_fp8_config"]
        is_fp8 = fp8_cfg.get("quant_method") == "fp8"
        linear_fp8_config = Fp8Config(
            is_checkpoint_fp8_serialized=is_fp8,
            activation_scheme=fp8_cfg.get("activation_scheme", "dynamic"),
            weight_block_size=fp8_cfg.get("weight_block_size"),
        )

6. 스킴 판별 로직

가중치/활성화의 비트 수, 전략, 동적 여부를 조합하여 적절한 스킴을 판별한다.

def _is_dynamic_token_w8a8(self, weight_quant, input_quant):
    is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8
    weight_strategy = (
        weight_quant.strategy == QuantizationStrategy.TENSOR.value
        or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
    )
    is_token = (
        weight_strategy
        and input_quant.strategy == QuantizationStrategy.TOKEN.value
    )
    is_dynamic = not weight_quant.dynamic and input_quant.dynamic
    return is_8_bits and is_token and weight_quant.symmetric and is_dynamic

이 함수는 "가중치 8비트 정적 + 활성화 8비트 동적 토큰별" 조합을 감지한다.

7. 희소성(Sparsity) 통합

양자화와 희소성 압축을 함께 지원한다.

@classmethod
def _parse_sparsity_config(cls, config):
    if not (sparsity_config := config.get(SPARSITY_CONFIG_NAME)):
        return dict(), []
    sparsity_config = SparsityCompressionConfig.model_validate(sparsity_config)
    sparse_scheme_map = {
        target: sparsity_config
        for target in sparsity_config.targets or list()
    }
    return sparse_scheme_map, sparsity_config.ignore or list()

2:4 구조적 희소성(structured sparsity) 등을 양자화와 결합하여 추가 압축을 달성할 수 있다.

8. GPU 능력 검사

스킴별 최소 GPU 능력을 검증한다.

def _check_scheme_supported(self, min_capability, error=True):
    capability_tuple = DeviceCapability(
        *torch.cuda.get_device_capability()
    )
    capability = capability_tuple.to_int()
    supported = capability >= min_capability
    if error and not supported:
        raise RuntimeError(
            f"Min capability: {min_capability}. "
            f"Current capability: {capability}."
        )
    return supported

설계 근거

  1. 스킴 패턴 기반 자동 선택: 비트 수, 전략, 동적 여부의 조합으로 스킴을 결정하여, 사용자가 구체적인 커널을 알 필요가 없다.
  2. target_scheme_map: 레이어 이름 패턴 → 양자화 설정 매핑으로, 모델의 다른 부분에 다른 양자화를 적용할 수 있다.
  3. 혼합 양자화: linear_fp8_config로 MoE 전문가와 Dense 레이어의 양자화를 독립적으로 설정한다.
  4. 희소성 통합: 양자화와 희소성을 하나의 설정에서 관리하여, 두 기법을 결합한 모델도 지원한다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글