| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU chunked-prefill / prefix-caching correctness for linear-attention models.""" |
| 4 | |
| 5 | import os |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from tests.models.utils import check_logprobs_close |
| 10 | from vllm import LLM, SamplingParams |
| 11 | from vllm.platforms import current_platform |
| 12 | |
| 13 | pytestmark = pytest.mark.cpu_model |
| 14 | |
| 15 | if not current_platform.is_cpu(): |
| 16 | pytest.skip("skipping CPU-only tests", allow_module_level=True) |
| 17 | |
| 18 | # Bound the KV cache so the run does not scale with host memory; these engines |
| 19 | # only need a few thousand tokens. |
| 20 | os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1") |
| 21 | |
| 22 | MODEL = "Qwen/Qwen3.5-0.8B" |
| 23 | CHUNK_TOKENS = 128 # max_num_batched_tokens for the chunked engine |
| 24 | NUM_LOGPROBS = 5 |
| 25 | SP = SamplingParams(max_tokens=32, temperature=0, logprobs=NUM_LOGPROBS) |
| 26 | |
| 27 | |
| 28 | def _long_prompt(repeat: int) -> str: |
| 29 | return "Solve the following arithmetic step by step. " * repeat + "What is 7*8?" |
| 30 | |
| 31 | |
| 32 | # Prompts long enough to span several CHUNK_TOKENS-sized chunks; a single-chunk |
| 33 | # prompt is bit-identical to full prefill regardless of the bug. |
| 34 | PROMPTS = [_long_prompt(r) for r in (40, 60, 80)] |
| 35 | # Spans several full cache blocks; prefix caching only reuses complete blocks. |
| 36 | PREFIX_PROMPT = "You are a helpful assistant. " * 230 + " Now answer: what is 2+2?" |
| 37 | |
| 38 | |
| 39 | def _make_llm(**overrides) -> LLM: |
| 40 | base = dict( |
| 41 | model=MODEL, |
| 42 | dtype="bfloat16", |
| 43 | max_model_len=2048, |
| 44 | enforce_eager=True, |
| 45 | trust_remote_code=True, |
| 46 | ) |
| 47 | base.update(overrides) |
| 48 | return LLM(**base) |
| 49 | |
| 50 | |
| 51 | def _tuples(outputs) -> list[tuple[list[int], str, object]]: |
| 52 | """(token_ids, text, sample_logprobs) per request, for check_logprobs_close.""" |
| 53 | return [ |
| 54 | (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs) |
| 55 | for o in outputs |
| 56 | ] |
| 57 | |
| 58 | |
| 59 | @pytest.fixture(scope="module") |
| 60 | def full_prefill_refs(): |
| 61 | """Reference (ids, text, logprobs) for PROMPTS and PREFIX_PROMPT, full prefill.""" |
| 62 | llm = _make_llm(enable_chunked_prefill=False, enable_prefix_caching=False) |
| 63 | refs = _tuples(llm.generate(PROMPTS, SP)) |
| 64 | prefix_ref = _tuples(llm.generate([PREFIX_PROMPT], SP))[0] |
| 65 | del llm |
| 66 | return refs, prefix_ref |
| 67 | |
| 68 | |
| 69 | def test_chunked_prefill_matches_full_prefill(full_prefill_refs): |
| 70 | """Batched multi-chunk prefill must stay close to per-prompt full prefill. |
| 71 | |
| 72 | Prompts are scheduled together so the scheduler interleaves prefill chunks |
| 73 | across requests (the cross-request path where the accuracy gap was strongest). |
| 74 | """ |
| 75 | refs, _ = full_prefill_refs |
| 76 | llm = _make_llm( |
| 77 | enable_chunked_prefill=True, |
| 78 | max_num_batched_tokens=CHUNK_TOKENS, |
| 79 | enable_prefix_caching=False, |
| 80 | ) |
| 81 | got = _tuples(llm.generate(PROMPTS, SP)) |
| 82 | del llm |
| 83 | |
| 84 | check_logprobs_close( |
| 85 | outputs_0_lst=refs, |
| 86 | outputs_1_lst=got, |
| 87 | name_0="full_prefill", |
| 88 | name_1="chunked_prefill", |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs): |
| 93 | """A prefix-cache hit must stay close to the cold-cache (reference) output. |
| 94 | |
| 95 | The warm run continues prefill from the restored GDN state; the |
| 96 | num_cached_tokens check guards against a vacuous (no-hit) pass. |
| 97 | """ |
| 98 | _, ref = full_prefill_refs |
| 99 | llm = _make_llm(enable_prefix_caching=True) |
| 100 | llm.generate([PREFIX_PROMPT], SP) # prime the cache |
| 101 | warm_out = llm.generate([PREFIX_PROMPT], SP)[0] |
| 102 | warm = _tuples([warm_out])[0] |
| 103 | del llm |
| 104 | |
| 105 | assert warm_out.num_cached_tokens > 0, ( |
| 106 | "expected a prefix-cache hit but num_cached_tokens=0; " |
| 107 | "PREFIX_PROMPT may be shorter than one cache block" |
| 108 | ) |
| 109 | check_logprobs_close( |
| 110 | outputs_0_lst=[ref], |
| 111 | outputs_1_lst=[warm], |
| 112 | name_0="cold_cache", |
| 113 | name_1="warm_cache", |
| 114 | ) |