A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] Add num_cached_tokens for PoolingRequestOutput (#27378)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Oct 23, 2025 at 14:03 UTC
3729ed00ba4c5a9f351b1241b15a2ca3ca12a5e1
7 files changed
+75
-5
tests/models/language/pooling/test_auto_prefix_cache_support.py
+24
-4
@@ -19,14 +19,25 @@ def test_classify_models(
19
model: str,
20
dtype: str,
21
) -> None:
22
- example_prompts = example_prompts * 2
22
+ # example_prompts is too short for testing prefix_caching
23
+ example_prompts = [s * 10 for s in example_prompts]
24
25
with vllm_runner(
26
model, max_model_len=512, dtype=dtype, enable_prefix_caching=True
27
) as vllm_model:
28
cache_config = vllm_model.llm.llm_engine.cache_config
29
assert cache_config.enable_prefix_caching
29
- vllm_outputs = vllm_model.classify(example_prompts)
30
+
31
+ # First Run
32
+ vllm_model.classify(example_prompts)
33
+
34
+ # assert prefix_caching works
35
+ pooling_outputs = vllm_model.llm.encode(
36
+ example_prompts, pooling_task="classify"
37
+ )
38
+ for output in pooling_outputs:
39
+ assert output.num_cached_tokens > 0
40
+ vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs]
41
42
with hf_runner(
43
model, dtype=dtype, auto_cls=AutoModelForSequenceClassification
@@ -54,7 +65,8 @@ def test_embed_models(
65
model: str,
66
dtype: str,
67
):
57
- example_prompts = [str(s).strip() for s in example_prompts] * 2
68
+ # example_prompts is too short for testing prefix_caching
69
+ example_prompts = [str(s).strip() * 10 for s in example_prompts]
70
71
with vllm_runner(
72
model,
@@ -64,7 +76,15 @@ def test_embed_models(
76
) as vllm_model:
77
cache_config = vllm_model.llm.llm_engine.cache_config
78
assert cache_config.enable_prefix_caching
67
- vllm_outputs = vllm_model.embed(example_prompts)
79
+
80
+ # First Run
81
+ vllm_model.embed(example_prompts)
82
+
83
+ # assert prefix_caching works
84
+ pooling_outputs = vllm_model.llm.encode(example_prompts, pooling_task="embed")
85
+ for output in pooling_outputs:
86
+ assert output.num_cached_tokens > 0
87
+ vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs]
88
89
with hf_runner(
90
model,
tests/models/language/pooling/test_extract_hidden_states.py
new
+33
@@ -0,0 +1,33 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from vllm import TokensPrompt
7
+
8
+
9
+@pytest.mark.parametrize(
10
+ "model",
11
+ ["Qwen/Qwen3-0.6B"],
12
+)
13
+@torch.inference_mode
14
+def test_embed_models(hf_runner, vllm_runner, model: str):
15
+ n_prompt_tokens = [55, 56, 57]
16
+ token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
17
+
18
+ with vllm_runner(
19
+ model,
20
+ max_model_len=128,
21
+ enforce_eager=True,
22
+ runner="pooling",
23
+ enable_chunked_prefill=False,
24
+ enable_prefix_caching=False,
25
+ ) as vllm_model:
26
+ pooling_outputs = vllm_model.llm.encode(
27
+ [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
28
+ pooling_task="token_embed",
29
+ )
30
+
31
+ for n, output in zip(n_prompt_tokens, pooling_outputs):
32
+ assert len(output.prompt_token_ids) == n
33
+ assert output.num_cached_tokens == 0
vllm/entrypoints/llm.py
+3
@@ -1078,6 +1078,9 @@ class LLM:
1078
PoolingRequestOutput[Any](
1079
request_id="",
1080
outputs=processed_outputs,
1081
+ num_cached_tokens=getattr(
1082
+ processed_outputs, "num_cached_tokens", 0
1083
+ ),
1084
prompt_token_ids=[],
1085
finished=True,
1086
)
vllm/entrypoints/openai/serving_embedding.py
+1
@@ -583,6 +583,7 @@ class EmbeddingMixin(OpenAIServing):
583
request_id=aggregator["request_id"],
584
prompt_token_ids=original_token_ids,
585
outputs=pooling_output_data,
586
+ num_cached_tokens=0,
587
finished=True,
588
)
589
vllm/entrypoints/score_utils.py
+1
@@ -66,6 +66,7 @@ def _cosine_similarity(
66
request_id=f"{emb_1.request_id}_{emb_2.request_id}",
67
outputs=pair_score,
68
prompt_token_ids=tokens,
69
+ num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
70
finished=True,
71
)
72
)
vllm/outputs.py
+12
-1
@@ -201,14 +201,21 @@ class PoolingRequestOutput(Generic[_O]):
201
request_id (str): A unique identifier for the pooling request.
202
outputs (PoolingOutput): The pooling results for the given input.
203
prompt_token_ids (list[int]): A list of token IDs used in the prompt.
204
+ num_cached_tokens: The number of tokens with prefix cache hit.
205
finished (bool): A flag indicating whether the pooling is completed.
206
"""
207
208
def __init__(
208
- self, request_id: str, outputs: _O, prompt_token_ids: list[int], finished: bool
209
+ self,
210
+ request_id: str,
211
+ outputs: _O,
212
+ prompt_token_ids: list[int],
213
+ num_cached_tokens: int,
214
+ finished: bool,
215
):
216
self.request_id = request_id
217
self.prompt_token_ids = prompt_token_ids
218
+ self.num_cached_tokens = num_cached_tokens
219
self.finished = finished
220
self.outputs = outputs
221
@@ -217,6 +224,7 @@ class PoolingRequestOutput(Generic[_O]):
224
f"{type(self).__name__}(request_id={self.request_id!r}, "
225
f"outputs={self.outputs!r}, "
226
f"prompt_token_ids={self.prompt_token_ids}, "
227
+ f"num_cached_tokens={self.num_cached_tokens}, "
228
f"finished={self.finished})"
229
)
230
@@ -255,6 +263,7 @@ class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]):
263
request_id=request_output.request_id,
264
outputs=EmbeddingOutput.from_base(request_output.outputs),
265
prompt_token_ids=request_output.prompt_token_ids,
266
+ num_cached_tokens=request_output.num_cached_tokens,
267
finished=request_output.finished,
268
)
269
@@ -294,6 +303,7 @@ class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
303
request_id=request_output.request_id,
304
outputs=ClassificationOutput.from_base(request_output.outputs),
305
prompt_token_ids=request_output.prompt_token_ids,
306
+ num_cached_tokens=request_output.num_cached_tokens,
307
finished=request_output.finished,
308
)
309
@@ -330,5 +340,6 @@ class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]):
340
request_id=request_output.request_id,
341
outputs=ScoringOutput.from_base(request_output.outputs),
342
prompt_token_ids=request_output.prompt_token_ids,
343
+ num_cached_tokens=request_output.num_cached_tokens,
344
finished=request_output.finished,
345
)
vllm/v1/engine/output_processor.py
+1
@@ -230,6 +230,7 @@ class RequestState:
230
return PoolingRequestOutput(
231
request_id=request_id,
232
outputs=first_output,
233
+ num_cached_tokens=self.num_cached_tokens,
234
prompt_token_ids=self.prompt_token_ids,
235
finished=finished,
236
)