A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 06:00 UTC
[Bugfix][2/n] Fix speculative decoding CI - Fix test_ngram_e2e_greedy_correctness (#19644)
Lu Fang committed
Jun 15, 2025 at 12:15 UTC
ee1531bc3874c10b233c8825aef976c796bf1ecc
5 files changed
+50
-3
tests/spec_decode/e2e/test_integration.py
+10
-1
@@ -14,10 +14,13 @@ MAIN_MODEL = "JackFram/llama-68m"
14
@pytest.mark.parametrize(
15
"common_llm_kwargs",
16
[{
17
+ "model_name": "JackFram/llama-68m",
18
19
# Verify equality when cuda graphs allowed.
20
"enforce_eager": False,
20
- "model_name": "JackFram/llama-68m",
21
+
22
+ # The original model is float32, keep it for numerical stability.
23
+ "dtype": "float32",
24
}])
25
@pytest.mark.parametrize(
26
"per_test_common_llm_kwargs",
@@ -59,6 +62,9 @@ def test_spec_decode_cuda_graph(vllm_runner, common_llm_kwargs,
62
63
# Skip cuda graph recording for fast test.
64
"enforce_eager": True,
65
+
66
+ # The original model is float32, keep it for numerical stability.
67
+ "dtype": "float32",
68
}])
69
@pytest.mark.parametrize("per_test_common_llm_kwargs", [])
70
@pytest.mark.parametrize(
@@ -117,6 +123,9 @@ def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
123
124
# Skip cuda graph recording for fast test.
125
"enforce_eager": True,
126
+
127
+ # The original model is float32, keep it for numerical stability.
128
+ "dtype": "float32",
129
}])
130
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
131
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
tests/spec_decode/e2e/test_logprobs.py
+17
-1
@@ -17,7 +17,10 @@ from .conftest import run_equality_correctness_test
17
"model_name": "JackFram/llama-160m",
18
19
# Skip cuda graph recording for fast test.
20
- "enforce_eager": True
20
+ "enforce_eager": True,
21
+
22
+ # The original model is float32, keep it for numerical stability.
23
+ "dtype": "float32",
24
}])
25
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
26
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -75,6 +78,9 @@ def test_logprobs_equality(vllm_runner, common_llm_kwargs,
78
79
# Skip cuda graph recording for fast test.
80
"enforce_eager": True,
81
+
82
+ # The original model is float32, keep it for numerical stability.
83
+ "dtype": "float32",
84
}])
85
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
86
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -128,6 +134,9 @@ def test_logprobs_different_k(vllm_runner, common_llm_kwargs,
134
135
# Skip cuda graph recording for fast test.
136
"enforce_eager": True,
137
+
138
+ # The original model is float32, keep it for numerical stability.
139
+ "dtype": "float32",
140
}])
141
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
142
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -182,6 +191,9 @@ def test_logprobs_when_skip_speculation(vllm_runner, common_llm_kwargs,
191
192
# Skip cuda graph recording for fast test.
193
"enforce_eager": True,
194
+
195
+ # The original model is float32, keep it for numerical stability.
196
+ "dtype": "float32",
197
}])
198
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
199
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -256,8 +268,12 @@ def test_logprobs_temp_1(vllm_runner, common_llm_kwargs,
268
"common_llm_kwargs",
269
[{
270
"model_name": "JackFram/llama-160m",
271
+
272
# Skip cuda graph recording for fast test.
273
"enforce_eager": True,
274
+
275
+ # The original model is float32, keep it for numerical stability.
276
+ "dtype": "float32",
277
}])
278
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
279
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
tests/spec_decode/e2e/test_mlp_correctness.py
+3
@@ -494,6 +494,9 @@ def test_mlp_disable_queue(vllm_runner, common_llm_kwargs,
494
495
# Skip cuda graph recording for fast test.
496
"enforce_eager": True,
497
+
498
+ # Precision
499
+ "dtype": PRECISION,
500
}])
501
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
502
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
tests/spec_decode/e2e/test_ngram_correctness.py
+18
@@ -40,6 +40,9 @@ from .conftest import run_equality_correctness_test
40
41
# Print spec metrics.
42
"disable_log_stats": False,
43
+
44
+ # The original model is float32, keep it for numerical stability.
45
+ "dtype": "float32",
46
}])
47
@pytest.mark.parametrize("per_test_common_llm_kwargs", [
48
{
@@ -97,6 +100,9 @@ def test_ngram_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
100
101
# Print spec metrics.
102
"disable_log_stats": False,
103
+
104
+ # The original model is float32, keep it for numerical stability.
105
+ "dtype": "float32",
106
}])
107
@pytest.mark.parametrize("per_test_common_llm_kwargs", [
108
{
@@ -160,6 +166,9 @@ def test_ngram_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
166
167
# Skip cuda graph recording for fast test.
168
"enforce_eager": True,
169
+
170
+ # The original model is float32, keep it for numerical stability.
171
+ "dtype": "float32",
172
}])
173
@pytest.mark.parametrize("per_test_common_llm_kwargs", [
174
{
@@ -221,6 +230,9 @@ def test_ngram_e2e_greedy_correctness_with_preemption(
230
231
# Skip cuda graph recording for fast test.
232
"enforce_eager": True,
233
+
234
+ # The original model is float32, keep it for numerical stability.
235
+ "dtype": "float32",
236
}])
237
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
238
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -281,6 +293,9 @@ def test_ngram_different_k(vllm_runner, common_llm_kwargs,
293
294
# Skip cuda graph recording for fast test.
295
"enforce_eager": True,
296
+
297
+ # The original model is float32, keep it for numerical stability.
298
+ "dtype": "float32",
299
}])
300
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
301
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -337,6 +352,9 @@ def test_ngram_disable_queue(vllm_runner, common_llm_kwargs,
352
353
# Skip cuda graph recording for fast test.
354
"enforce_eager": True,
355
+
356
+ # The original model is float32, keep it for numerical stability.
357
+ "dtype": "float32",
358
}])
359
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
360
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
vllm/model_executor/models/eagle.py
+2
-1
@@ -74,6 +74,7 @@ class EAGLE(nn.Module):
74
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
75
super().__init__()
76
config = vllm_config.model_config.hf_config
77
+ self.dtype = vllm_config.model_config.dtype
78
self.config = config
79
80
architectures = getattr(self.config.model, "architectures", [])
@@ -250,7 +251,7 @@ class EAGLE(nn.Module):
251
lm_head_weight = torch.zeros(
252
self.lm_head.org_vocab_size,
253
self.lm_head.embedding_dim,
253
- dtype=self.config.torch_dtype,
254
+ dtype=self.dtype,
255
)
256
257
weight_loader = getattr(self.lm_head.weight, "weight_loader",