| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm import LLM, SamplingParams |
| 9 | from vllm.config import CompilationConfig, CompilationMode |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | from ....utils import check_answers, fork_new_process_for_each_test, prep_prompts |
| 13 | |
| 14 | # global seed |
| 15 | SEED = 42 |
| 16 | |
| 17 | |
| 18 | @pytest.fixture |
| 19 | def test_prompts(): |
| 20 | """Adapted from tests/v1/e2e/spec_decode/utils.py.""" |
| 21 | prompt_types = ["repeat", "sentence"] |
| 22 | # Setting higher num prompts increases the chance of numerics mismatch |
| 23 | # due to matrix multiplication numerics depending on batch dimension |
| 24 | num_prompts = 10 |
| 25 | prompts = [] |
| 26 | |
| 27 | random.seed(0) |
| 28 | random_prompt_type_choices = random.choices(prompt_types, k=num_prompts) |
| 29 | |
| 30 | for kind in random_prompt_type_choices: |
| 31 | word_choices = ["test", "temp", "hello", "where"] |
| 32 | word = random.choice(word_choices) |
| 33 | if kind == "repeat": |
| 34 | prompt = f"""please repeat the word '{word}' 10 times.""" |
| 35 | elif kind == "sentence": |
| 36 | prompt = f"""please give a ten-word sentence that |
| 37 | uses the word {word} at least once.""" |
| 38 | else: |
| 39 | raise ValueError(f"Unknown prompt type: {kind}") |
| 40 | prompts.append(prompt) |
| 41 | |
| 42 | return prompts |
| 43 | |
| 44 | |
| 45 | use_fork_for_test = ( |
| 46 | fork_new_process_for_each_test |
| 47 | if not (current_platform.is_rocm() or current_platform.is_xpu()) |
| 48 | else lambda x: x |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | @use_fork_for_test |
| 53 | @pytest.mark.parametrize( |
| 54 | "kv_sharing_fast_prefill,use_v2_model_runner", |
| 55 | [ |
| 56 | (False, False), |
| 57 | (True, False), |
| 58 | (True, True), |
| 59 | # (False, True) omitted: fast-prefill-off behavior is runner-generic, |
| 60 | # and skipping it saves ~9 min of mostly torch.compile time in CI. The |
| 61 | # no-silent-fallback assertion below is exercised by the V2 cases. |
| 62 | ], |
| 63 | ) |
| 64 | @pytest.mark.parametrize("enforce_eager", [True, False]) |
| 65 | def test_kv_sharing_fast_prefill( |
| 66 | monkeypatch: pytest.MonkeyPatch, |
| 67 | kv_sharing_fast_prefill: bool, |
| 68 | enforce_eager: bool, |
| 69 | use_v2_model_runner: bool, |
| 70 | ): |
| 71 | if not enforce_eager and current_platform.is_rocm(): |
| 72 | # Relevant context: https://github.com/vllm-project/vllm/pull/29244 |
| 73 | pytest.skip( |
| 74 | "ROCm: torch.compile produces incorrect output for gemma-3n's GELU " |
| 75 | "with tanh approximation. Use enforce_eager=True instead." |
| 76 | ) |
| 77 | |
| 78 | sampling_params = SamplingParams(temperature=0.0, max_tokens=100) |
| 79 | compilation_config = CompilationConfig( |
| 80 | # This allows vLLM compilation backend to handle allocating and |
| 81 | # managing buffers for cudagraph |
| 82 | cudagraph_copy_inputs=True, |
| 83 | mode=CompilationMode.VLLM_COMPILE |
| 84 | if not enforce_eager |
| 85 | else CompilationMode.NONE, |
| 86 | ) |
| 87 | batch_size = 10 |
| 88 | |
| 89 | with monkeypatch.context() as m: |
| 90 | # Make scheduling deterministic for reproducibility |
| 91 | if current_platform.is_rocm(): |
| 92 | # Use spawn to prevent cuda re-initialization error |
| 93 | m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 94 | else: |
| 95 | m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 96 | m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0") |
| 97 | |
| 98 | prompts, answer, indices = prep_prompts(batch_size) |
| 99 | |
| 100 | llm = LLM( |
| 101 | model="google/gemma-3n-E2B-it", |
| 102 | enforce_eager=enforce_eager, |
| 103 | compilation_config=compilation_config, |
| 104 | seed=SEED, |
| 105 | kv_sharing_fast_prefill=kv_sharing_fast_prefill, |
| 106 | attention_backend="TRITON_ATTN", |
| 107 | ) |
| 108 | # Guard against a silent fallback to the other model runner. |
| 109 | assert llm.llm_engine.vllm_config.use_v2_model_runner == use_v2_model_runner |
| 110 | responses = llm.generate(prompts, sampling_params) |
| 111 | check_answers( |
| 112 | indices, |
| 113 | answer, |
| 114 | [response.outputs[0].text for response in responses], |
| 115 | accept_rate=1.0, |
| 116 | ) |