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 13, 22:00 UTC
[Bugfix] Padded Eagle Specdec with Chunked Prefill (#26263)
Signed-off-by: Rémi Delacourt <remi@mistral.ai> Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Signed-off-by: remi <remi@mistral.ai> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Rémi Delacourt committed
Nov 3, 2025 at 08:22 UTC
cec7c288333339028f6fe8e0ac3222e3924da90b
1 file changed
+21
-7
tests/v1/e2e/test_spec_decode.py
+21
-7
@@ -202,9 +202,9 @@ def test_speculators_model_integration(
202
203
204
@pytest.mark.parametrize(
205
- ["model_setup", "mm_enabled"],
205
+ ["model_setup", "mm_enabled", "chunked_prefill_enabled"],
206
[
207
- (("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), False),
207
+ (("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), False, False),
208
pytest.param(
209
(
210
"eagle3",
@@ -213,11 +213,12 @@ def test_speculators_model_integration(
213
1,
214
),
215
False,
216
+ False,
217
marks=pytest.mark.skip(
218
reason="Skipping due to its head_dim not being a a multiple of 32"
219
),
220
),
220
- (
221
+ pytest.param(
222
(
223
"eagle",
224
"meta-llama/Llama-3.1-8B-Instruct",
@@ -225,7 +226,9 @@ def test_speculators_model_integration(
226
1,
227
),
228
False,
228
- ),
229
+ True,
230
+ marks=large_gpu_mark(min_gb=40),
231
+ ), # works on 4x H100
232
(
233
(
234
"eagle3",
@@ -234,6 +237,7 @@ def test_speculators_model_integration(
237
1,
238
),
239
False,
240
+ False,
241
),
242
pytest.param(
243
(
@@ -243,6 +247,7 @@ def test_speculators_model_integration(
247
4,
248
),
249
False,
250
+ False,
251
marks=large_gpu_mark(min_gb=80),
252
), # works on 4x H100
253
pytest.param(
@@ -253,6 +258,7 @@ def test_speculators_model_integration(
258
4,
259
),
260
True,
261
+ True,
262
marks=large_gpu_mark(min_gb=80),
263
), # works on 4x H100
264
(
@@ -263,6 +269,7 @@ def test_speculators_model_integration(
269
1,
270
),
271
False,
272
+ False,
273
),
274
],
275
ids=[
@@ -281,6 +288,7 @@ def test_eagle_correctness(
288
sampling_config: SamplingParams,
289
model_setup: tuple[str, str, str, int],
290
mm_enabled: bool,
291
+ chunked_prefill_enabled: bool,
292
attn_backend: str,
293
):
294
if attn_backend == "TREE_ATTN":
@@ -317,9 +325,13 @@ def test_eagle_correctness(
325
m.setenv("VLLM_ROCM_USE_AITER", "1")
326
327
method, model_name, spec_model_name, tp_size = model_setup
328
+ max_model_len = 2048
329
+ max_num_batched_tokens = max_model_len
330
+ if chunked_prefill_enabled:
331
+ max_num_batched_tokens = 128
332
333
ref_llm = LLM(
322
- model=model_name, max_model_len=2048, tensor_parallel_size=tp_size
334
+ model=model_name, max_model_len=max_model_len, tensor_parallel_size=tp_size
335
)
336
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
337
del ref_llm
@@ -334,9 +346,11 @@ def test_eagle_correctness(
346
"method": method,
347
"model": spec_model_name,
348
"num_speculative_tokens": 3,
337
- "max_model_len": 2048,
349
+ "max_model_len": max_model_len,
350
},
339
- max_model_len=2048,
351
+ max_model_len=max_model_len,
352
+ max_num_batched_tokens=max_num_batched_tokens,
353
+ enable_chunked_prefill=chunked_prefill_enabled,
354
)
355
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
356
matches = 0