From a47c58c3f9333341fc3cbfc3ee52dc49fbc37f53 Mon Sep 17 00:00:00 2001 From: DEV TYAGI Date: Tue, 20 Oct 2026 20:47:41 +0530 Subject: [PATCH 1/2] fix(kv_cache): scope checkpoint scale filtering to mismatch and handle bf16 runtime (#42419) Signed-off-by: DEV TYAGI --- python/sglang/srt/model_loader/loader.py | 56 ++++++++++++++++++ .../model_loader/test_default_model_loader.py | 58 +++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index fdb318756740..6a397d0f77c7 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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": diff --git a/test/registered/unit/model_loader/test_default_model_loader.py b/test/registered/unit/model_loader/test_default_model_loader.py index c70c0da94812..21c72582e815 100644 --- a/test/registered/unit/model_loader/test_default_model_loader.py +++ b/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() From 6e1a8d7ca28bacb852c06a216bedf80a8308672b Mon Sep 17 00:00:00 2001 From: DEV TYAGI Date: Wed, 21 Oct 2026 15:06:31 +0530 Subject: [PATCH 2/2] fix(kv_cache): resolve auto kv_cache_dtype to fp8 pool on compressed-tensors checkpoints Signed-off-by: DEV TYAGI --- python/sglang/srt/model_loader/loader.py | 7 ++++++- .../model_loader/test_default_model_loader.py | 17 ++++++++++------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 6a397d0f77c7..0ba6010129a6 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -1074,7 +1074,12 @@ def load_weights_only(model, weights, target_device): if isinstance(kv_scheme, dict): checkpoint_bits = kv_scheme.get("num_bits") expected_bits = None - if runtime_kv_dtype: + if runtime_kv_dtype == "auto": + # configure_kv_cache_dtype resolves auto to an FP8 pool here. + algo = getattr(quant_config, "kv_cache_quant_algo", None) + if isinstance(algo, str) and algo.upper() == "FP8": + expected_bits = 8 + elif runtime_kv_dtype: if "fp8" in runtime_kv_dtype or "int8" in runtime_kv_dtype: expected_bits = 8 elif "fp4" in runtime_kv_dtype: diff --git a/test/registered/unit/model_loader/test_default_model_loader.py b/test/registered/unit/model_loader/test_default_model_loader.py index 21c72582e815..5875230992eb 100644 --- a/test/registered/unit/model_loader/test_default_model_loader.py +++ b/test/registered/unit/model_loader/test_default_model_loader.py @@ -131,18 +131,21 @@ def load_weights_and_postprocess(self, model, weights, target_device): 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), + # (checkpoint_bits, runtime_kv_dtype, expect_scales_dropped, kv_cache_quant_algo) + (8, "nvfp4", True, None), + (8, "bfloat16", True, None), + (8, "fp8", False, None), + (8, "auto", False, "FP8"), + (4, "nvfp4", False, None), + (4, "fp8", True, None), ] - for checkpoint_bits, runtime_dtype, expect_dropped in cases: + for checkpoint_bits, runtime_dtype, expect_dropped, *extra in cases: + quant_algo = extra[0] if extra else None with self.subTest(ckpt=checkpoint_bits, runtime=runtime_dtype): model = Mock() model.quant_config = SimpleNamespace( kv_cache_scheme={"num_bits": checkpoint_bits}, + kv_cache_quant_algo=quant_algo, is_nvfp4_online=False, get_name=lambda: "dummy", )