A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Initialize missing attributes in mistral eagle (#45217)
Signed-off-by: jpwang <jpwang@smail.nju.edu.cn> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
jpwang committed
Jun 12, 2026 at 08:21 UTC
6f573f486bc659adf51a8d4639e225097f2d8d39
2 files changed
+156
tests/model_executor/test_mistral_large_3_eagle.py
new
+146
@@ -0,0 +1,146 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+import torch.nn as nn
9
+
10
+from vllm.config.compilation import CompilationMode
11
+from vllm.model_executor.models import deepseek_v2 as deepseek_mod
12
+from vllm.model_executor.models import mistral_large_3_eagle as eagle_mod
13
+
14
+
15
+class DummyPPGroup:
16
+ world_size = 1
17
+ is_first_rank = True
18
+ is_last_rank = True
19
+
20
+
21
+class DummyEmbedding(nn.Module):
22
+ def __init__(self, vocab_size, hidden_size, *args, **kwargs):
23
+ super().__init__()
24
+ self.hidden_size = hidden_size
25
+
26
+ def forward(self, input_ids):
27
+ return torch.zeros(
28
+ (*input_ids.shape, self.hidden_size),
29
+ dtype=torch.float32,
30
+ device=input_ids.device,
31
+ )
32
+
33
+
34
+class DummyLinear(nn.Module):
35
+ def __init__(self, in_features, out_features, *args, **kwargs):
36
+ super().__init__()
37
+ self.out_features = out_features
38
+
39
+ def forward(self, x):
40
+ return torch.zeros(
41
+ (*x.shape[:-1], self.out_features),
42
+ dtype=x.dtype,
43
+ device=x.device,
44
+ )
45
+
46
+
47
+class DummyNorm(nn.Module):
48
+ def __init__(self, *args, **kwargs):
49
+ super().__init__()
50
+
51
+ def forward(self, hidden_states, residual=None):
52
+ return hidden_states, residual
53
+
54
+
55
+class DummyDecoderLayer(nn.Module):
56
+ def __init__(self, *args, **kwargs):
57
+ super().__init__()
58
+
59
+ def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
60
+ return hidden_states, residual
61
+
62
+
63
+def make_vllm_config(
64
+ *, model_type="mistral3", qk_nope_head_dim=128, qk_rope_head_dim=64
65
+):
66
+ hf_config = SimpleNamespace(
67
+ model_type=model_type,
68
+ first_k_dense_replace=0,
69
+ vocab_size=32000,
70
+ hidden_size=16,
71
+ num_hidden_layers=1,
72
+ rms_norm_eps=1e-5,
73
+ qk_nope_head_dim=qk_nope_head_dim,
74
+ qk_rope_head_dim=qk_rope_head_dim,
75
+ )
76
+
77
+ return SimpleNamespace(
78
+ model_config=SimpleNamespace(hf_config=hf_config),
79
+ quant_config=None,
80
+ parallel_config=SimpleNamespace(
81
+ eplb_config=SimpleNamespace(num_redundant_experts=0),
82
+ ),
83
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
84
+ cache_config=None,
85
+ compilation_config=SimpleNamespace(mode=CompilationMode.NONE),
86
+ )
87
+
88
+
89
+@pytest.fixture(autouse=True)
90
+def patch_heavy_modules(monkeypatch):
91
+ monkeypatch.setattr(eagle_mod, "get_pp_group", lambda: DummyPPGroup())
92
+ monkeypatch.setattr(deepseek_mod, "get_pp_group", lambda: DummyPPGroup())
93
+
94
+ monkeypatch.setattr(eagle_mod, "VocabParallelEmbedding", DummyEmbedding)
95
+ monkeypatch.setattr(eagle_mod, "RowParallelLinear", DummyLinear)
96
+ monkeypatch.setattr(eagle_mod, "RMSNorm", DummyNorm)
97
+ monkeypatch.setattr(eagle_mod, "DeepseekV2DecoderLayer", DummyDecoderLayer)
98
+
99
+
100
+@pytest.mark.cpu_test
101
+@pytest.mark.parametrize(
102
+ ("model_type", "qk_nope_head_dim", "qk_rope_head_dim", "expected_use_mha"),
103
+ [
104
+ # MLA-style config: should not use MHA.
105
+ ("mistral3", 128, 64, False),
106
+ # No MLA dims: should use MHA, matching DeepseekV2Model.__init__ logic.
107
+ ("mistral3", 0, 0, True),
108
+ # DeepSeek model type always uses MHA by the parent logic.
109
+ ("deepseek", 128, 64, True),
110
+ ],
111
+)
112
+def test_eagle_mistral_large3_initializes_deepseek_runtime_attrs(
113
+ model_type,
114
+ qk_nope_head_dim,
115
+ qk_rope_head_dim,
116
+ expected_use_mha,
117
+):
118
+ vllm_config = make_vllm_config(
119
+ model_type=model_type,
120
+ qk_nope_head_dim=qk_nope_head_dim,
121
+ qk_rope_head_dim=qk_rope_head_dim,
122
+ )
123
+
124
+ model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
125
+
126
+ assert model.aux_hidden_state_layers == ()
127
+ assert model.use_mha is expected_use_mha
128
+
129
+ # Add this if your fix also copies num_redundant_experts from
130
+ # DeepseekV2Model.__init__.
131
+ assert model.num_redundant_experts == 0
132
+
133
+
134
+@pytest.mark.cpu_test
135
+def test_eagle_mistral_large3_forward_reuses_deepseek_parent_forward():
136
+ vllm_config = make_vllm_config()
137
+ model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
138
+
139
+ input_ids = torch.tensor([[1, 2, 3]])
140
+ positions = torch.tensor([[0, 1, 2]])
141
+ hidden_states = torch.zeros((1, 3, 16))
142
+
143
+ output = model(input_ids, positions, hidden_states)
144
+
145
+ assert isinstance(output, torch.Tensor)
146
+ assert output.shape == hidden_states.shape
vllm/model_executor/models/mistral_large_3_eagle.py
+10
@@ -75,6 +75,16 @@ class EagleMistralLarge3Model(DeepseekV2Model):
75
)
76
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
77
self.aux_hidden_state_layers: tuple[int, ...] = ()
78
+
79
+ # Needed by load_weights
80
+ qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
81
+ qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
82
+ self.use_mha = config.model_type == "deepseek" or all(
83
+ dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
84
+ )
85
+ self.num_redundant_experts = (
86
+ vllm_config.parallel_config.eplb_config.num_redundant_experts
87
+ )
88
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
89
["hidden_states", "residual"], config.hidden_size
90
)