| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from tests.utils import multi_gpu_test |
| 6 | from tests.v1.e2e.spec_decode.utils import compute_acceptance_len |
| 7 | from vllm import LLM, SamplingParams |
| 8 | from vllm.distributed import cleanup_dist_env_and_memory |
| 9 | |
| 10 | MODEL = "meta-llama/Llama-3.2-1B-Instruct" |
| 11 | DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3" |
| 12 | PROMPTS = [ |
| 13 | "The capital of France is", |
| 14 | "2 + 2 equals", |
| 15 | "In one word, the color of the sky is", |
| 16 | "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:", |
| 17 | ] |
| 18 | |
| 19 | ACCEPTANCE_TOLERANCE = 0.95 |
| 20 | |
| 21 | |
| 22 | def _run(pp_size: int) -> float: |
| 23 | llm = LLM( |
| 24 | model=MODEL, |
| 25 | tensor_parallel_size=1, |
| 26 | pipeline_parallel_size=pp_size, |
| 27 | max_model_len=512, |
| 28 | gpu_memory_utilization=0.45, |
| 29 | disable_log_stats=False, |
| 30 | compilation_config={"cudagraph_mode": "FULL_AND_PIECEWISE"}, |
| 31 | speculative_config={ |
| 32 | "method": "eagle3", |
| 33 | "model": DRAFT, |
| 34 | "num_speculative_tokens": 3, |
| 35 | }, |
| 36 | ) |
| 37 | try: |
| 38 | llm.generate( |
| 39 | PROMPTS, |
| 40 | SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True), |
| 41 | ) |
| 42 | acceptance = compute_acceptance_len(llm.get_metrics()) |
| 43 | assert acceptance > 1 |
| 44 | return acceptance |
| 45 | finally: |
| 46 | del llm |
| 47 | torch.accelerator.empty_cache() |
| 48 | cleanup_dist_env_and_memory() |
| 49 | |
| 50 | |
| 51 | @multi_gpu_test(num_gpus=4) |
| 52 | def test_eagle3_pipeline_parallel_acceptance(): |
| 53 | baseline = _run(1) |
| 54 | for pp_size in (2, 4): |
| 55 | parallel = _run(pp_size) |
| 56 | assert parallel >= baseline * ACCEPTANCE_TOLERANCE, ( |
| 57 | f"PP={pp_size} acceptance regressed: {parallel:.3f} < " |
| 58 | f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}" |
| 59 | ) |