diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index aeecaee62..e57a7a3b5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -64,6 +64,7 @@ from sglang.srt.layers.quantization.compressed_tensors.utils import ( should_ignore_layer, ) from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod from sglang.srt.layers.quantization.unquant import ( UnquantizedFusedMoEMethod, UnquantizedLinearMethod, @@ -131,6 +132,17 @@ class CompressedTensorsConfig(QuantizationConfig): self.packed_modules_mapping = packed_modules_mapping or {} self.linear_fp8_config = linear_fp8_config + @property + def kv_cache_quant_algo(self) -> Optional[str]: + """Duck-typed by configure_kv_cache_dtype to resolve --kv-cache-dtype + auto: loaded scales need the fp8 pool they calibrate, never bf16.""" + if ( + self.kv_cache_scheme is not None + and CompressedTensorsKVCacheMethod.is_supported_scheme(self.kv_cache_scheme) + ): + return "FP8" + return None + def get_linear_method(self) -> CompressedTensorsLinearMethod: return CompressedTensorsLinearMethod(self) @@ -156,8 +168,9 @@ class CompressedTensorsConfig(QuantizationConfig): self.sparsity_ignore_list = hf_to_sglang_mapper.apply_list( self.sparsity_ignore_list ) - if self.kv_cache_scheme is not None: - self.kv_cache_scheme = hf_to_sglang_mapper.apply_dict(self.kv_cache_scheme) + # kv_cache_scheme is deliberately not remapped: it holds schema fields + # (type/num_bits/strategy), never module names, and apply_dict drops + # keys a mapper deletion rule happens to match. def get_quant_method( self, @@ -187,6 +200,23 @@ class CompressedTensorsConfig(QuantizationConfig): layer.scheme = scheme return CompressedTensorsLinearMethod(self) + from sglang.srt.layers.radix_attention import RadixAttention + + if isinstance(layer, RadixAttention): + if self.kv_cache_scheme is None: + return None + if not CompressedTensorsKVCacheMethod.is_supported_scheme( + self.kv_cache_scheme + ): + # Degrade, don't refuse to boot: unquantized-scale KV serves fine. + logger.warning_once( + f"Ignoring compressed-tensors kv_cache_scheme " + f"{self.kv_cache_scheme}: only static symmetric " + f"per-tensor FP8 scales are supported." + ) + return None + return CompressedTensorsKVCacheMethod(self) + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if isinstance(layer, FusedMoE): @@ -271,6 +301,7 @@ class CompressedTensorsConfig(QuantizationConfig): quant_format=quant_format, sparsity_scheme_map=sparsity_scheme_map, sparsity_ignore_list=sparsity_ignore_list, + kv_cache_scheme=config.get("kv_cache_scheme"), config=config, packed_modules_mapping=packed_modules_mapping, linear_fp8_config=linear_fp8_config, @@ -1085,6 +1116,27 @@ class CompressedTensorsConfig(QuantizationConfig): return weight_quant.num_bits == input_quant.num_bits == 8 +class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): + """Load calibrated k_scale / v_scale from a compressed-tensors checkpoint + that declares a ``kv_cache_scheme`` (static per-tensor FP8).""" + + def __init__(self, quant_config: CompressedTensorsConfig): + assert self.is_supported_scheme(quant_config.kv_cache_scheme) + super().__init__(quant_config) + + @staticmethod + def is_supported_scheme(kv_cache_scheme: Dict[str, Any]) -> bool: + """Static symmetric per-tensor FP8 — all BaseKVCacheMethod can + represent. Dynamic schemes serialize no k_scale/v_scale tensors.""" + return ( + kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + and kv_cache_scheme.get("strategy") == "tensor" + and kv_cache_scheme.get("symmetric", True) + and not kv_cache_scheme.get("dynamic", False) + ) + + class CompressedTensorsLinearMethod(LinearMethodBase): def __init__(self, quantization_config: CompressedTensorsConfig): diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 291bcbb95..eb61b5fd9 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1293,6 +1293,9 @@ QWEN3_5_KV_SCALE_MAPPER = WeightsMapper( orig_to_new_substr={ ".self_attn.k_proj.k_scale": ".attn.k_scale", ".self_attn.v_proj.v_scale": ".attn.v_scale", + # compressed-tensors stores kv_cache_scheme scales on the attention module. + ".self_attn.k_scale": ".attn.k_scale", + ".self_attn.v_scale": ".attn.v_scale", }, ) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index feac0250b..130f9aad4 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -33,7 +33,7 @@ from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM +from sglang.srt.models.qwen3_5 import QWEN3_5_KV_SCALE_MAPPER, Qwen3_5ForCausalLM from sglang.srt.runtime_context import ( get_model, get_parallel, @@ -232,6 +232,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False ): + weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_kv_cache.py b/test/registered/unit/layers/quantization/test_compressed_tensors_kv_cache.py new file mode 100644 index 000000000..1333e221b --- /dev/null +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_kv_cache.py @@ -0,0 +1,98 @@ +"""Unit tests for compressed-tensors KV cache scale loading — CPU-only.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import unittest + +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + CompressedTensorsKVCacheMethod, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.test.test_utils import CustomTestCase + +_FP8_TENSOR_KV_SCHEME = { + "type": "float", + "num_bits": 8, + "strategy": "tensor", + "symmetric": True, + "dynamic": False, +} + + +def _config(kv_cache_scheme): + cfg = { + "format": "float-quantized", + "quant_method": "compressed-tensors", + "ignore": [], + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 8, + "type": "float", + "strategy": "channel", + "symmetric": True, + "dynamic": False, + }, + "input_activations": { + "num_bits": 8, + "type": "float", + "strategy": "token", + "symmetric": True, + "dynamic": True, + }, + } + }, + } + if kv_cache_scheme is not None: + cfg["kv_cache_scheme"] = kv_cache_scheme + return CompressedTensorsConfig.from_config(cfg) + + +def _attn(): + # __new__ is enough: get_quant_method only isinstance-checks the layer. + return RadixAttention.__new__(RadixAttention) + + +class TestCompressedTensorsKVCacheMethod(CustomTestCase): + def test_declared_scheme_gets_kv_cache_method(self): + """A declared supported scheme must produce the KV cache method; + without it the calibrated k_scale/v_scale have no parameters to + load into and fp8 KV runs unscaled.""" + config = _config(_FP8_TENSOR_KV_SCHEME) + method = config.get_quant_method(_attn(), "model.layers.0.attn") + self.assertIsInstance(method, CompressedTensorsKVCacheMethod) + + def test_no_scheme_returns_none(self): + config = _config(None) + self.assertIsNone(config.get_quant_method(_attn(), "model.layers.0.attn")) + + def test_kv_cache_quant_algo_resolves_auto_dtype(self): + """configure_kv_cache_dtype duck-types this field for --kv-cache-dtype + auto: supported schemes must report FP8, everything else None, so + loaded scales always meet an fp8 pool.""" + self.assertEqual(_config(_FP8_TENSOR_KV_SCHEME).kv_cache_quant_algo, "FP8") + self.assertIsNone(_config(None).kv_cache_quant_algo) + self.assertIsNone( + _config(dict(_FP8_TENSOR_KV_SCHEME, dynamic=True)).kv_cache_quant_algo + ) + + def test_unsupported_scheme_degrades_to_none(self): + """Unsupported declared schemes must skip the method, not fail + the boot: such checkpoints serve with an unquantized-scale cache.""" + for bad in ( + dict(_FP8_TENSOR_KV_SCHEME, type="int"), + dict(_FP8_TENSOR_KV_SCHEME, strategy="channel"), + dict(_FP8_TENSOR_KV_SCHEME, symmetric=False), + dict(_FP8_TENSOR_KV_SCHEME, dynamic=True), + ): + self.assertIsNone( + _config(bad).get_quant_method(_attn(), "model.layers.0.attn") + ) + + +if __name__ == "__main__": + unittest.main()