[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:
-            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):

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