Fix Ministral3 accuracy issue by aligning YaRN RoPE scaling with Transformers implementation (#31232)

Co-authored-by: Elizaveta Martirosian <elizaveta.martirosian@gmail.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/srt/layers/rotary_embedding/factory.py

@@ -243,7 +243,14 @@ def get_rope(
                 k: v
                 for k, v in rope_scaling.items()
                 if k
-                in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow")
+                in (
+                    "extrapolation_factor",
+                    "attn_factor",
+                    "beta_fast",
+                    "beta_slow",
+                    "mscale",
+                    "mscale_all_dim",
+                )
             }
             extra_kwargs["truncate"] = rope_scaling.get("truncate", True)
             if "mrope_section" in rope_scaling:

python/sglang/srt/layers/rotary_embedding/yarn.py

@@ -83,6 +83,8 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
         beta_fast: int = 32,
         beta_slow: int = 1,
         truncate: bool = True,
+        mscale: float = None,
+        mscale_all_dim: float = None,
     ) -> None:
         self.scaling_factor = scaling_factor
         self.extrapolation_factor = extrapolation_factor
@@ -90,8 +92,18 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
         self.truncate = truncate
-        # Get n-d magnitude scaling corrected for interpolation
+
-        self.mscale = float(yarn_get_mscale_simple(self.scaling_factor) * attn_factor)
+        if mscale is not None and mscale_all_dim is not None:
+            # Match Hugging Face's YaRN RoPE scaling (supports mscale/mscale_all_dim)
+            self.mscale = float(
+                yarn_get_mscale(self.scaling_factor, mscale)
+                / yarn_get_mscale(self.scaling_factor, mscale_all_dim)
+            )
+        else:
+            # Get n-d magnitude scaling corrected for interpolation
+            self.mscale = float(yarn_get_mscale_simple(self.scaling_factor))
+        self.mscale *= attn_factor
+
         super().__init__(
             head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
         )