[misc] Unify MLA scaling init and remove dead buffer / scaling code (#33363)

test/registered/unit/configs/test_model_config_scaling.py

@@ -1,13 +1,22 @@
 import math
 import unittest
 
-from sglang.srt.configs.model_config import compute_mla_mscale_scaling
+from sglang.srt.configs.model_config import ModelConfig, 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")
 
 
+def _mla_scaling(rope_scaling) -> tuple[float, float]:
+    """Run ModelConfig._init_mla_scaling on fixed head dims -> (base, result)."""
+    config = ModelConfig.__new__(ModelConfig)
+    config.qk_nope_head_dim = 128
+    config.qk_rope_head_dim = 64
+    config._init_mla_scaling(rope_scaling)
+    return 1 / math.sqrt(192), config.scaling
+
+
 class TestMlaMscaleScaling(CustomTestCase):
     def test_ignores_transformers_v5_default_rope_parameters(self):
         base_scaling = 1 / math.sqrt(72)
@@ -65,5 +74,24 @@ class TestMlaMscaleScaling(CustomTestCase):
         )
 
 
+class TestInitMlaScaling(CustomTestCase):
+    """ModelConfig._init_mla_scaling must not re-add a "default" fallback of its
+    own: DeepseekV2AttentionMLA stamps rope_type="deepseek_yarn" on any non-empty
+    rope_scaling, so a dict without rope_type/type still gets a yarn rope and its
+    mscale belongs in the scale FlashInferMLA reads as sm_scale."""
+
+    def test_applies_mscale_without_rope_type(self):
+        base, scaling = _mla_scaling({"factor": 40, "mscale_all_dim": 1})
+        self.assertGreater(scaling, base)
+
+    def test_ignores_default_rope_type(self):
+        base, scaling = _mla_scaling({"rope_type": "default", "factor": 40})
+        self.assertEqual(scaling, base)
+
+    def test_no_rope_scaling_keeps_base(self):
+        base, scaling = _mla_scaling(None)
+        self.assertEqual(scaling, base)
+
+
 if __name__ == "__main__":
     unittest.main()