[Quant] Load compressed-tensors kv_cache_scheme scales (#35455)

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:
+        # kv_cache_scheme is deliberately not remapped: it holds schema fields
-            self.kv_cache_scheme = hf_to_sglang_mapper.apply_dict(self.kv_cache_scheme)
+        # (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):

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",
     },
 )
 

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"),

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()