diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index db705eae9..15dd56361 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1679,6 +1679,11 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models. Warns if 'factor' is missing from rope_scaling (common in v5 configs). """ + if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get( + "apply_scale", True + ): + return base_scaling + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) if "factor" not in rope_scaling: logger.warning( diff --git a/test/registered/unit/configs/test_model_config_scaling.py b/test/registered/unit/configs/test_model_config_scaling.py new file mode 100644 index 000000000..f04e1681f --- /dev/null +++ b/test/registered/unit/configs/test_model_config_scaling.py @@ -0,0 +1,52 @@ +import math +import unittest + +from sglang.srt.configs.model_config import compute_mla_mscale_scaling +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestMlaMscaleScaling(CustomTestCase): + def test_respects_disabled_yarn_scaling(self): + base_scaling = 1 / math.sqrt(128) + rope_scaling = { + "rope_type": "deepseek_yarn", + "factor": 128, + "mscale_all_dim": 1, + "apply_yarn_scaling": False, + } + + self.assertEqual( + compute_mla_mscale_scaling(rope_scaling, base_scaling), base_scaling + ) + + def test_applies_yarn_scaling_by_default(self): + base_scaling = 1 / math.sqrt(128) + rope_scaling = { + "rope_type": "deepseek_yarn", + "factor": 128, + "mscale_all_dim": 1, + } + + self.assertGreater( + compute_mla_mscale_scaling(rope_scaling, base_scaling), base_scaling + ) + + def test_respects_disabled_native_apply_scale(self): + base_scaling = 1 / math.sqrt(128) + rope_scaling = { + "rope_type": "mistral", + "factor": 128, + "mscale_all_dim": 1, + "apply_scale": False, + } + + self.assertEqual( + compute_mla_mscale_scaling(rope_scaling, base_scaling), base_scaling + ) + + +if __name__ == "__main__": + unittest.main()