[CI] Stabilize speculator memory teardown (#50284)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Jul 30, 2026 at 05:42 UTC
61c1d098e5124f3566b22a3178ef46551a4bd8e3
1 file changed
+37
-29
tests/v1/e2e/spec_decode/test_spec_decode.py
+37
-29
@@ -379,6 +379,7 @@ def test_speculators_model_integration(
379
sampling_config: SamplingParams,
380
model_path: str,
381
expected_accuracy_threshold: float,
382
+ vllm_runner,
383
):
384
"""
385
Test that speculators models work with the simplified integration.
@@ -401,41 +402,48 @@ def test_speculators_model_integration(
402
test_prompts = get_test_prompts(mm_enabled=False)
403
404
# First run: Direct speculator model (simplified integration)
404
- spec_llm = LLM(model=model_path, max_model_len=4096, gpu_memory_utilization=0.92)
405
- evaluate_llm_for_gsm8k(
406
- spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
407
- )
408
- spec_outputs = spec_llm.chat(test_prompts, sampling_config)
409
-
410
- # Verify speculative config was auto-detected
411
- assert spec_llm.llm_engine.vllm_config.speculative_config is not None, (
412
- f"Speculative config should be auto-detected for {model_path}"
413
- )
405
+ with vllm_runner(
406
+ model_path,
407
+ trust_remote_code=False,
408
+ enable_chunked_prefill=None,
409
+ compilation_config=CompilationConfig(),
410
+ max_model_len=4096,
411
+ gpu_memory_utilization=0.92,
412
+ ) as spec_runner:
413
+ evaluate_llm_for_gsm8k(
414
+ spec_runner.llm, expected_accuracy_threshold=expected_accuracy_threshold
415
+ )
416
+ spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
417
415
- spec_config = spec_llm.llm_engine.vllm_config.speculative_config
416
- assert spec_config.num_speculative_tokens > 0, (
417
- f"Expected positive speculative tokens, "
418
- f"got {spec_config.num_speculative_tokens}"
419
- )
418
+ # Verify speculative config was auto-detected
419
+ assert spec_runner.llm.llm_engine.vllm_config.speculative_config is not None, (
420
+ f"Speculative config should be auto-detected for {model_path}"
421
+ )
422
421
- # Verify draft model is set to the speculator model
422
- assert spec_config.model == model_path, (
423
- f"Draft model should be {model_path}, got {spec_config.model}"
424
- )
423
+ spec_config = spec_runner.llm.llm_engine.vllm_config.speculative_config
424
+ assert spec_config.num_speculative_tokens > 0, (
425
+ f"Expected positive speculative tokens, "
426
+ f"got {spec_config.num_speculative_tokens}"
427
+ )
428
426
- # Extract verifier model for reference run
427
- verifier_model = spec_llm.llm_engine.vllm_config.model_config.model
429
+ # Verify draft model is set to the speculator model
430
+ assert spec_config.model == model_path, (
431
+ f"Draft model should be {model_path}, got {spec_config.model}"
432
+ )
433
429
- del spec_llm
430
- torch.accelerator.empty_cache()
431
- cleanup_dist_env_and_memory()
434
+ # Extract verifier model for reference run
435
+ verifier_model = spec_runner.llm.llm_engine.vllm_config.model_config.model
436
437
# Second run: Reference without speculative decoding
434
- ref_llm = LLM(model=verifier_model, max_model_len=4096, gpu_memory_utilization=0.92)
435
- ref_outputs = ref_llm.chat(test_prompts, sampling_config)
436
- del ref_llm
437
- torch.accelerator.empty_cache()
438
- cleanup_dist_env_and_memory()
438
+ with vllm_runner(
439
+ verifier_model,
440
+ trust_remote_code=False,
441
+ enable_chunked_prefill=None,
442
+ compilation_config=CompilationConfig(),
443
+ max_model_len=4096,
444
+ gpu_memory_utilization=0.92,
445
+ ) as ref_runner:
446
+ ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
447
448
# Compare outputs
449
matches = sum(