| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm import LLM |
| 7 | |
| 8 | MODEL = "hmellor/tiny-random-LlamaForCausalLM" |
| 9 | PROMPT = "Hello my name is Robert and I" |
| 10 | |
| 11 | |
| 12 | @pytest.fixture(scope="module") |
| 13 | def llm() -> LLM: |
| 14 | return LLM( |
| 15 | MODEL, |
| 16 | enforce_eager=True, |
| 17 | enable_prefix_caching=True, |
| 18 | long_prefill_token_threshold=2, |
| 19 | max_num_batched_tokens=6, |
| 20 | max_num_seqs=3, |
| 21 | block_size=16, |
| 22 | ) |
| 23 | |
| 24 | |
| 25 | def test_concurrent_partial_prefill(llm): |
| 26 | outputs = llm.generate([PROMPT] * 3) |
| 27 | assert len(outputs) == 3 |
| 28 | for output in outputs: |
| 29 | assert len(output.outputs) == 1 |
| 30 | |
| 31 | |
| 32 | def test_prefix_cache_stats_is_recorded(llm): |
| 33 | # 17 tokens will make sure first 16 tokens are cached in a block |
| 34 | input_tokens = {"prompt_token_ids": [101] * 17} |
| 35 | _ = llm.generate([input_tokens]) |
| 36 | outputs = llm.generate([input_tokens]) |
| 37 | assert outputs[0].num_cached_tokens == 16 |