[fix] Mistral-Large-3 accuracy regression on `main` (#57563)
Signed-off-by: jdebache <jdebache@nvidia.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Julien Debache committed
Sep 18, 2026 at 16:19 UTC
2bdae2a5a83fe8e5640078a3d4bfe86cce323d0b
8 files changed
+79
-27
tests/transformers_utils/test_config.py
+52
@@ -28,6 +28,7 @@ from vllm.transformers_utils.configs.glm5_next import (
28
Glm5NextTextConfig,
29
Glm5NextVisionConfig,
30
)
31
+from vllm.transformers_utils.configs.mistral import adapt_config_dict
32
33
34
def test_patch_legacy_rope_type_preserves_nope_layers():
@@ -77,6 +78,57 @@ def test_patch_legacy_rope_type_normalizes_telechat3_yarn():
78
}
79
80
81
+def test_mistral_yarn_apply_scale_false_disables_yarn_magnitude_scaling():
82
+ """`yarn.apply_scale: false` must reach the DeepSeek-style attentions.
83
+
84
+ Transformers spells it `attention_factor = 1.0`, which DeepseekV2Attention
85
+ and its siblings read to select `deepseek_llama_scaling` over
86
+ `deepseek_yarn`; without it Mistral-Large-3 runs with a spurious
87
+ yarn_get_mscale(factor)^2 attention scaling.
88
+ """
89
+ params = {
90
+ "dim": 7168,
91
+ "n_layers": 61,
92
+ "head_dim": 192,
93
+ "hidden_dim": 16384,
94
+ "n_heads": 128,
95
+ "n_kv_heads": 128,
96
+ "norm_eps": 1e-5,
97
+ "vocab_size": 131072,
98
+ "rope_theta": 10000.0,
99
+ "max_position_embeddings": 294912,
100
+ "q_lora_rank": 1536,
101
+ "kv_lora_rank": 512,
102
+ "qk_nope_head_dim": 128,
103
+ "qk_rope_head_dim": 64,
104
+ "v_head_dim": 128,
105
+ "moe": {
106
+ "num_experts": 128,
107
+ "num_experts_per_tok": 4,
108
+ "num_shared_experts": 1,
109
+ "expert_hidden_dim": 4096,
110
+ "first_k_dense_replace": 3,
111
+ "route_every_n": 1,
112
+ "routed_scale": 1.0,
113
+ "num_expert_groups": 1,
114
+ "num_expert_groups_per_tok": 1,
115
+ },
116
+ "llama_4_scaling": {"beta": 0.1, "original_max_position_embeddings": 8192},
117
+ "yarn": {
118
+ "alpha": 1,
119
+ "apply_scale": False,
120
+ "beta": 32,
121
+ "factor": 36,
122
+ "original_max_position_embeddings": 8192,
123
+ },
124
+ }
125
+
126
+ config = adapt_config_dict(params, defaults={})
127
+
128
+ assert config.architectures == ["MistralLarge3ForCausalLM"]
129
+ assert config.rope_parameters["attention_factor"] == 1.0
130
+
131
+
132
def test_glm5_next_accepts_deepseek_sparse_attention_layers():
133
layer_types = ["linear_attention", "deepseek_sparse_attention"]
134
vllm/model_executor/models/AXK1.py
+6
-6
@@ -324,9 +324,9 @@ class AXK1Attention(nn.Module):
324
assert config.rope_parameters is not None
325
if config.rope_parameters["rope_type"] != "default":
326
config.rope_parameters["rope_type"] = (
327
- "deepseek_yarn"
328
- if config.rope_parameters.get("apply_yarn_scaling", True)
329
- else "deepseek_llama_scaling"
327
+ "deepseek_llama_scaling"
328
+ if config.rope_parameters.get("attention_factor") == 1.0
329
+ else "deepseek_yarn"
330
)
331
332
self.rotary_emb = get_rope(
@@ -496,9 +496,9 @@ class AXK1MLAAttention(nn.Module):
496
assert config.rope_parameters is not None
497
if config.rope_parameters["rope_type"] != "default":
498
config.rope_parameters["rope_type"] = (
499
- "deepseek_yarn"
500
- if config.rope_parameters.get("apply_yarn_scaling", True)
501
- else "deepseek_llama_scaling"
499
+ "deepseek_llama_scaling"
500
+ if config.rope_parameters.get("attention_factor") == 1.0
501
+ else "deepseek_yarn"
502
)
503
504
self.rotary_emb = get_rope(
vllm/model_executor/models/deepseek_v2.py
+6
-6
@@ -540,9 +540,9 @@ class DeepseekV2Attention(nn.Module):
540
)
541
if config.rope_parameters["rope_type"] != "default":
542
config.rope_parameters["rope_type"] = (
543
- "deepseek_yarn"
544
- if config.rope_parameters.get("apply_yarn_scaling", True)
545
- else "deepseek_llama_scaling"
543
+ "deepseek_llama_scaling"
544
+ if config.rope_parameters.get("attention_factor") == 1.0
545
+ else "deepseek_yarn"
546
)
547
548
self.rotary_emb = get_rope(
@@ -1107,9 +1107,9 @@ class DeepseekV2MLAAttention(nn.Module):
1107
1108
if config.rope_parameters["rope_type"] != "default":
1109
config.rope_parameters["rope_type"] = (
1110
- "deepseek_yarn"
1111
- if config.rope_parameters.get("apply_yarn_scaling", True)
1112
- else "deepseek_llama_scaling"
1110
+ "deepseek_llama_scaling"
1111
+ if config.rope_parameters.get("attention_factor") == 1.0
1112
+ else "deepseek_yarn"
1113
)
1114
1115
self.rotary_emb = get_rope(
vllm/models/deepseek_v32/attention.py
+3
-3
@@ -153,9 +153,9 @@ class DeepseekV32Attention(MLAAttention):
153
# DSA checkpoints may use plain ("default") or yarn-scaled RoPE.
154
if config.rope_parameters["rope_type"] != "default":
155
config.rope_parameters["rope_type"] = (
156
- "deepseek_yarn"
157
- if config.rope_parameters.get("apply_yarn_scaling", True)
158
- else "deepseek_llama_scaling"
156
+ "deepseek_llama_scaling"
157
+ if config.rope_parameters.get("attention_factor") == 1.0
158
+ else "deepseek_yarn"
159
)
160
if config.rope_parameters["rope_type"] == "deepseek_yarn":
161
mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
vllm/models/deepseek_v4/common/rope.py
+3
-3
@@ -31,9 +31,9 @@ def build_deepseek_v4_rope(
31
if compress_ratio > 1 and rope_parameters["rope_type"] != "default":
32
# YaRN applies only to compressor (CSA/HCA) layers.
33
rope_parameters["rope_type"] = (
34
- "deepseek_yarn"
35
- if rope_parameters.get("apply_yarn_scaling", True)
36
- else "deepseek_llama_scaling"
34
+ "deepseek_llama_scaling"
35
+ if rope_parameters.get("attention_factor") == 1.0
36
+ else "deepseek_yarn"
37
)
38
else:
39
# Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
vllm/models/deepseek_v41/common/rope.py
+3
-3
@@ -32,9 +32,9 @@ def build_deepseek_v4_rope(
32
# YaRN applies to all compressed-KV layers (compress_ratio > 0); only
33
# pure sliding-window layers (compress_ratio == 0) use plain RoPE.
34
rope_parameters["rope_type"] = (
35
- "deepseek_yarn"
36
- if rope_parameters.get("apply_yarn_scaling", True)
37
- else "deepseek_llama_scaling"
35
+ "deepseek_llama_scaling"
36
+ if rope_parameters.get("attention_factor") == 1.0
37
+ else "deepseek_yarn"
38
)
39
else:
40
# Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
vllm/models/glm5next/common/attention.py
+3
-3
@@ -496,9 +496,9 @@ class Glm5NextMLAAttention(nn.Module):
496
assert config.rope_parameters is not None
497
if config.rope_parameters["rope_type"] != "default":
498
config.rope_parameters["rope_type"] = (
499
- "deepseek_yarn"
500
- if config.rope_parameters.get("apply_yarn_scaling", True)
501
- else "deepseek_llama_scaling"
499
+ "deepseek_llama_scaling"
500
+ if config.rope_parameters.get("attention_factor") == 1.0
501
+ else "deepseek_yarn"
502
)
503
504
self.rotary_emb: RotaryEmbedding | None = get_rope(
vllm/models/kimi_k3/nvidia/mla.py
+3
-3
@@ -165,9 +165,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
165
rope_parameters = dict(config.rope_parameters)
166
if rope_parameters["rope_type"] != "default":
167
rope_parameters["rope_type"] = (
168
- "deepseek_yarn"
169
- if rope_parameters.get("apply_yarn_scaling", True)
170
- else "deepseek_llama_scaling"
168
+ "deepseek_llama_scaling"
169
+ if rope_parameters.get("attention_factor") == 1.0
170
+ else "deepseek_yarn"
171
)
172
self.rotary_emb = get_rope(
173
qk_rope_head_dim,