| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()` |
| 4 | handling of maximum context length for decoder models. |
| 5 | |
| 6 | This test ensures: |
| 7 | - A prompt that is one token shorter than the model's maximum context length |
| 8 | can be processed successfully when requesting one additional token. |
| 9 | - A prompt that reaches the model's maximum context length throws a |
| 10 | `ValueError` when requesting at least one additional token. |
| 11 | """ |
| 12 | |
| 13 | import pytest |
| 14 | |
| 15 | from tests.conftest import VllmRunner |
| 16 | from tests.utils import create_new_process_for_each_test |
| 17 | from vllm.exceptions import VLLMValidationError |
| 18 | |
| 19 | |
| 20 | @create_new_process_for_each_test() |
| 21 | @pytest.mark.parametrize("model, max_model_len", [("JackFram/llama-160m", 2048)]) |
| 22 | @pytest.mark.parametrize( |
| 23 | "prompt_len, max_tokens", |
| 24 | [ |
| 25 | (2047, 1), # prompt_len = max_model_len - 1 -> allowed |
| 26 | (2048, 1), # prompt_len = max_model_len -> not allowed |
| 27 | ], |
| 28 | ) |
| 29 | def test_decoder_max_context_length_validation( |
| 30 | model: str, |
| 31 | max_model_len: int, |
| 32 | vllm_runner: type[VllmRunner], |
| 33 | prompt_len: int, |
| 34 | max_tokens: int, |
| 35 | ) -> None: |
| 36 | """Check vLLM decoder model input validation for edge cases where |
| 37 | the prompt length is (almost) equal to the max model length.""" |
| 38 | prompt_ids = [[43] * prompt_len] |
| 39 | |
| 40 | with vllm_runner( |
| 41 | model_name=model, |
| 42 | tokenizer_name=model, |
| 43 | max_model_len=max_model_len, |
| 44 | max_num_seqs=1, |
| 45 | tensor_parallel_size=1, |
| 46 | ) as vllm_model: |
| 47 | if prompt_len + max_tokens <= max_model_len: |
| 48 | # Should succeed as constraints are met |
| 49 | vllm_model.generate_greedy(prompt_ids, max_tokens) |
| 50 | else: |
| 51 | # Should raise the ValueError defined in |
| 52 | # vllm/v1/engine/processor.Processor_validate_model_input() |
| 53 | expected_msg = ( |
| 54 | f"The decoder prompt (length {prompt_len}) plus the number of " |
| 55 | f"requested output tokens (at least 1) is longer than " |
| 56 | f"the maximum model length of {max_model_len}. " |
| 57 | "Make sure that `max_model_len` is no smaller than the number of " |
| 58 | "text tokens (prompt + requested output tokens)." |
| 59 | ) |
| 60 | with pytest.raises(VLLMValidationError) as excinfo: |
| 61 | vllm_model.generate_greedy(prompt_ids, max_tokens) |
| 62 | assert expected_msg in str(excinfo.value) |
| 63 | |
| 64 | |
| 65 | @create_new_process_for_each_test() |
| 66 | @pytest.mark.parametrize("model", ["JackFram/llama-160m"]) |
| 67 | def test_auto_fit_max_model_len_rejects_oversized_input( |
| 68 | model: str, |
| 69 | vllm_runner: type[VllmRunner], |
| 70 | ) -> None: |
| 71 | """When max_model_len='auto' and KV cache memory is very limited, |
| 72 | the engine auto-fits max_model_len to a small value. The frontend |
| 73 | must see this reduced value and reject prompts that exceed it, |
| 74 | rather than accepting them and hanging.""" |
| 75 | # Use a small KV cache budget to force auto-fit to a small |
| 76 | # max_model_len. Pin block_size=16 so the budget is independent |
| 77 | # of the platform's default block size. One block for this model is |
| 78 | # 2 (K/V) * 16 (block) * 12 (heads) * 64 (head_dim) * 2 (fp16) |
| 79 | # * 12 (layers) = 589,824 bytes, and the pool must cover at least |
| 80 | # two: one is reserved as the null block, so a 1 MB single-block |
| 81 | # pool cannot serve any tokens at all. |
| 82 | kv_cache_bytes = 2_000_000 # 2 MB = 3 blocks -> 2 usable |
| 83 | |
| 84 | with vllm_runner( |
| 85 | model_name=model, |
| 86 | max_model_len=-1, |
| 87 | max_num_seqs=1, |
| 88 | enforce_eager=True, |
| 89 | block_size=16, |
| 90 | kv_cache_memory_bytes=kv_cache_bytes, |
| 91 | load_format="dummy", |
| 92 | ) as vllm_model: |
| 93 | auto_fitted_len = ( |
| 94 | vllm_model.llm.llm_engine.vllm_config.model_config.max_model_len |
| 95 | ) |
| 96 | # Sanity check: auto-fit should have reduced it well below the |
| 97 | # model's native context length. |
| 98 | assert auto_fitted_len < 2048, ( |
| 99 | f"Expected auto-fit to reduce max_model_len significantly, " |
| 100 | f"but got {auto_fitted_len}" |
| 101 | ) |
| 102 | |
| 103 | # A prompt longer than the auto-fitted length must be rejected. |
| 104 | oversized_prompt = [[43] * (auto_fitted_len + 10)] |
| 105 | with pytest.raises(VLLMValidationError, match="Please reduce the length"): |
| 106 | vllm_model.generate_greedy(oversized_prompt, max_tokens=4) |