[Fix] kv_cache: Drop fp8 scales when using nvfp4 KV cache

python/sglang/srt/model_loader/loader.py

@@ -1055,6 +1055,62 @@ def load_weights_only(model, weights, target_device):
                 for name, loaded_weight in weights
             )
 
+        try:
+            from sglang.srt.runtime_context import get_model
+
+            model_bag = get_model()
+            runtime_kv_dtype = getattr(model_bag, "kv_cache_dtype", None)
+
+            kv_scheme = None
+            if quant_config is not None and hasattr(quant_config, "kv_cache_scheme"):
+                kv_scheme = quant_config.kv_cache_scheme
+            elif hasattr(model, "config") and hasattr(model.config, "hf_config"):
+                quant_cfg = getattr(model.config.hf_config, "quantization_config", None)
+                if isinstance(quant_cfg, dict):
+                    kv_scheme = quant_cfg.get("kv_cache_scheme")
+                elif hasattr(quant_cfg, "kv_cache_scheme"):
+                    kv_scheme = getattr(quant_cfg, "kv_cache_scheme", None)
+
+            if isinstance(kv_scheme, dict):
+                checkpoint_bits = kv_scheme.get("num_bits")
+                expected_bits = None
+                if runtime_kv_dtype:
+                    if "fp8" in runtime_kv_dtype or "int8" in runtime_kv_dtype:
+                        expected_bits = 8
+                    elif "fp4" in runtime_kv_dtype:
+                        expected_bits = 4
+                    elif (
+                        "bf16" in runtime_kv_dtype
+                        or "bfloat16" in runtime_kv_dtype
+                        or "fp16" in runtime_kv_dtype
+                        or "float16" in runtime_kv_dtype
+                    ):
+                        expected_bits = 16
+
+                if checkpoint_bits and (
+                    expected_bits is None or checkpoint_bits != expected_bits
+                ):
+                    logger.info(
+                        "ignoring kv cache scales from checkpoint (checkpoint_bits=%s, runtime_dtype=%s). "
+                        "mismatch prevents scale corruption; falling back to dynamic scales.",
+                        checkpoint_bits,
+                        runtime_kv_dtype,
+                    )
+
+                    def _filter_kv_scales(ws):
+                        for n, w in ws:
+                            if (
+                                n.endswith("k_scale")
+                                or n.endswith("v_scale")
+                                or n.endswith("kv_scale")
+                            ):
+                                continue
+                            yield n, w
+
+                    weights = _filter_kv_scales(weights)
+        except (ValueError, AttributeError):
+            pass
+
         weights = monokernel_weights_iterator(model=model, weights=weights)
         model.load_weights(weights)
         if (is_nvfp4_online or is_modelopt_fp4_online) and target_device.type == "cuda":

test/registered/unit/model_loader/test_default_model_loader.py

@@ -129,6 +129,64 @@ def load_weights_and_postprocess(self, model, weights, target_device):
                     startup_prefetch_active=True,
                 )
 
+    def test_kv_cache_scale_dropped_on_mismatch(self):
+        cases = [
+            # (checkpoint_bits, runtime_kv_dtype, expect_scales_dropped)
+            (8, "nvfp4", True),
+            (8, "bfloat16", True),
+            (8, "fp8", False),
+            (4, "nvfp4", False),
+            (4, "fp8", True),
+        ]
+        for checkpoint_bits, runtime_dtype, expect_dropped in cases:
+            with self.subTest(ckpt=checkpoint_bits, runtime=runtime_dtype):
+                model = Mock()
+                model.quant_config = SimpleNamespace(
+                    kv_cache_scheme={"num_bits": checkpoint_bits},
+                    is_nvfp4_online=False,
+                    get_name=lambda: "dummy",
+                )
+                loaded = []
+                model.load_weights.side_effect = lambda ws: loaded.extend(list(ws))
+
+                dummy_weights = [
+                    ("model.layers.0.self_attn.q_proj.weight", torch.zeros(1)),
+                    ("model.layers.0.self_attn.k_scale", torch.ones(1)),
+                    ("model.layers.0.self_attn.v_scale", torch.ones(1)),
+                ]
+
+                model_bag = SimpleNamespace(kv_cache_dtype=runtime_dtype)
+                with (
+                    patch.object(loader_mod, "is_cuda_alike", return_value=False),
+                    patch(
+                        "sglang.srt.runtime_context.get_model",
+                        return_value=model_bag,
+                    ),
+                    patch.object(
+                        loader_mod,
+                        "monokernel_weights_iterator",
+                        side_effect=lambda model, weights: weights,
+                    ),
+                ):
+                    DefaultModelLoader.load_weights_only(
+                        model, iter(dummy_weights), torch.device("cpu")
+                    )
+
+                loaded_names = [name for name, _ in loaded]
+                if expect_dropped:
+                    self.assertEqual(
+                        loaded_names, ["model.layers.0.self_attn.q_proj.weight"]
+                    )
+                else:
+                    self.assertEqual(
+                        loaded_names,
+                        [
+                            "model.layers.0.self_attn.q_proj.weight",
+                            "model.layers.0.self_attn.k_scale",
+                            "model.layers.0.self_attn.v_scale",
+                        ],
+                    )
+
 
 if __name__ == "__main__":
     unittest.main()