[Test][Core] Add hybrid model prefix cache hit-rate coverage (#57127)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Codex <noreply@openai.com>
Jiangyun Zhu committed
Sep 17, 2026 at 10:17 UTC
b1f9ebc0c367ccc2aa1e584647e55351b97a5ec3
1 file changed
+156
tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Replay fixed chat histories and check per-turn cache hits and metrics."""
4
+
5
+import asyncio
6
+import json
7
+from uuid import uuid4
8
+
9
+import httpx
10
+import pytest
11
+from openai.types.chat import ChatCompletionMessageParam
12
+from prometheus_client.parser import text_string_to_metric_families
13
+
14
+from tests.utils import RemoteOpenAIServer
15
+
16
+MODEL = "Qwen/Qwen3.5-0.8B"
17
+NUM_CONVERSATIONS = 8
18
+MAX_TOKENS = 32
19
+# Long first turns leave room for hybrid cache alignment and MTP recomputation.
20
+MIN_PROMPT_TOKENS = 2048
21
+METRICS_TIMEOUT = 30
22
+
23
+
24
+def _user_turns(conversation: int) -> list[str]:
25
+ records = "\n".join(
26
+ f"Sensor {conversation}-{i}: temperature {20 + i % 10} C, "
27
+ f"humidity {40 + i % 20} percent."
28
+ for i in range(160 + conversation)
29
+ )
30
+ return [
31
+ f"Session {conversation}. Summarize these sensor readings briefly:\n{records}",
32
+ "Which readings should an operator investigate first? Explain briefly.",
33
+ "Suggest two checks the operator should perform next.",
34
+ ]
35
+
36
+
37
+async def _prefix_cache_counters(
38
+ client: httpx.AsyncClient,
39
+) -> tuple[float, float]:
40
+ response = await client.get("/metrics")
41
+ response.raise_for_status()
42
+ names = ("vllm:prefix_cache_queries_total", "vllm:prefix_cache_hits_total")
43
+ counters: dict[str, float] = {}
44
+ for family in text_string_to_metric_families(response.text):
45
+ for sample in family.samples:
46
+ if sample.name in names:
47
+ counters[sample.name] = counters.get(sample.name, 0) + sample.value
48
+ assert all(name in counters for name in names), (
49
+ f"missing prefix-cache metrics: {set(names) - counters.keys()}"
50
+ )
51
+ return counters[names[0]], counters[names[1]]
52
+
53
+
54
+async def _wait_for_counters(
55
+ client: httpx.AsyncClient, expected: tuple[float, float]
56
+) -> None:
57
+ """Allow engine stats to reach the frontend after completions are returned."""
58
+ deadline = asyncio.get_running_loop().time() + METRICS_TIMEOUT
59
+ while True:
60
+ actual = await _prefix_cache_counters(client)
61
+ if actual == expected:
62
+ return
63
+ assert all(a <= e for a, e in zip(actual, expected)), (
64
+ f"prefix-cache counters exceeded usage: {actual=}, {expected=}"
65
+ )
66
+ assert asyncio.get_running_loop().time() < deadline, (
67
+ f"prefix-cache counters did not match usage: {actual=}, {expected=}"
68
+ )
69
+ await asyncio.sleep(0.1)
70
+
71
+
72
+@pytest.fixture(scope="module")
73
+def server(request):
74
+ args = ["--enable-prompt-tokens-details"]
75
+ if request.param:
76
+ args += [
77
+ "--speculative-config",
78
+ json.dumps({"method": "mtp", "num_speculative_tokens": 1}),
79
+ ]
80
+ with RemoteOpenAIServer(MODEL, args) as remote_server:
81
+ yield remote_server
82
+
83
+
84
+@pytest.mark.asyncio
85
+@pytest.mark.parametrize(
86
+ "server, min_hit_rate",
87
+ # MTP re-prefills an extra block; both floors reject one more missed block.
88
+ [pytest.param(False, 0.90, id="base"), pytest.param(True, 0.75, id="mtp")],
89
+ indirect=["server"],
90
+)
91
+async def test_prefix_cache_hit_rate(
92
+ server: RemoteOpenAIServer, min_hit_rate: float
93
+) -> None:
94
+ # Isolate sessions so another conversation cannot hide a cache miss.
95
+ salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)]
96
+ histories: list[list[ChatCompletionMessageParam]] = [
97
+ [] for _ in range(NUM_CONVERSATIONS)
98
+ ]
99
+ prompt_lengths: list[list[int]] = [[] for _ in range(NUM_CONVERSATIONS)]
100
+
101
+ async with (
102
+ server.get_async_client() as client,
103
+ httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics,
104
+ ):
105
+
106
+ async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]:
107
+ history = histories[session]
108
+ queries = hits = 0
109
+ for turn, user_text in enumerate(_user_turns(session)):
110
+ if not replay:
111
+ history.append({"role": "user", "content": user_text})
112
+ # Reuse the original replies, even if replay generates different ones.
113
+ response = await client.chat.completions.create(
114
+ model=MODEL,
115
+ messages=history[: 2 * turn + 1],
116
+ temperature=0,
117
+ max_tokens=MAX_TOKENS,
118
+ extra_body={"cache_salt": salts[session]},
119
+ )
120
+ usage = response.usage
121
+ assert usage is not None
122
+ assert usage.prompt_tokens_details is not None
123
+ prompt = usage.prompt_tokens
124
+ cached = usage.prompt_tokens_details.cached_tokens
125
+ context = f"{session=}, {turn=}, {replay=}: {cached=}, {prompt=}"
126
+ assert cached is not None, context
127
+ assert 0 <= cached < prompt, context
128
+ assert prompt >= MIN_PROMPT_TOKENS, context
129
+ if not replay and turn == 0:
130
+ assert cached == 0, context
131
+ else:
132
+ assert cached / prompt >= min_hit_rate, (
133
+ f"{context}, {min_hit_rate=}"
134
+ )
135
+ if replay:
136
+ assert prompt == prompt_lengths[session][turn], context
137
+ else:
138
+ prompt_lengths[session].append(prompt)
139
+ history.append(
140
+ {
141
+ "role": "assistant",
142
+ "content": response.choices[0].message.content or "",
143
+ }
144
+ )
145
+ queries += prompt
146
+ hits += cached
147
+ return queries, hits
148
+
149
+ queries, hits = await _prefix_cache_counters(metrics)
150
+ for replay in (False, True):
151
+ counts = await asyncio.gather(
152
+ *(run_conversation(i, replay=replay) for i in range(NUM_CONVERSATIONS))
153
+ )
154
+ queries += sum(prompt for prompt, _ in counts)
155
+ hits += sum(cached for _, cached in counts)
156
+ await _wait_for_counters(metrics, (queries, hits))