[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
설계 근거
- 스킴 패턴 기반 자동 선택: 비트 수, 전략, 동적 여부의 조합으로 스킴을 결정하여, 사용자가 구체적인 커널을 알 필요가 없다.
- target_scheme_map: 레이어 이름 패턴 → 양자화 설정 매핑으로, 모델의 다른 부분에 다른 양자화를 적용할 수 있다.
- 혼합 양자화:
linear_fp8_config로 MoE 전문가와 Dense 레이어의 양자화를 독립적으로 설정한다. - 희소성 통합: 양자화와 희소성을 하나의 설정에서 관리하여, 두 기법을 결합한 모델도 지원한다.
관련 포스트
참고
- SGLang 소스:
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py - Compressed Tensors 라이브러리: neuralmagic/compressed-tensors
- vLLM 원본 구현: vllm-project/vllm
관련 포스트
- [sglang] DeepSeek-V4의 Latency 최적화: Fused mHC Post/Pre Kernel 도입
- [sglang] sglang ROCm MXFP4 어텐션에서 불필요한 contiguous copy 제거를 통한 성능 최적화
- [sglang] sglang의 torch.compile 활용: Advanced Indexing Gather 최적화로 LLM 추론 가속화
- [sglang] sglang diffusion 모델 성능 향상: Cache-DiT와 torch.compile의 최적화된 적용 순서
- [sglang] NixlKVManager 성능 향상: 비동기 및 멀티스레드 KV 전송 도입
SGLang 의 다른글
- 이전글 [SGLang] AutoRound: 자동 라운딩 최적화 양자화
- 현재글 : [SGLang] Compressed Tensors: 통합 양자화 프레임워크
- 다음글 [SGLang] W4A8, W8A8, W4A4: 혼합 정밀도 양자화 스킴
댓글